[PM-31763] Migrate userkey state to be client-managed SDK-available state (#18906)

* Migrate userkey state to be client-managed SDK-available state

* Update sdk to 556

* Fix types

* Fix build

* Linting

libs/common/src/key-management/user-key-mapper.ts

@@ -0,0 +1,21 @@
+import { SdkRecordMapper } from "@bitwarden/common/platform/services/sdk/client-managed-state";
+import { UserKeyDefinition } from "@bitwarden/common/platform/state";
+import { UserKeyState } from "@bitwarden/sdk-internal";
+
+import { SymmetricCryptoKey } from "../platform/models/domain/symmetric-crypto-key";
+import { USER_KEY } from "../platform/services/key-state/user-key.state";
+import { UserKey } from "../types/key";
+
+export class UserKeyRecordMapper implements SdkRecordMapper<UserKey, UserKeyState> {
+  userKeyDefinition(): UserKeyDefinition<Record<string, UserKey>> {
+    return USER_KEY;
+  }
+
+  toSdk(value: UserKey): UserKeyState {
+    return { decrypted_user_key: value.toBase64() } as UserKeyState;
+  }
+
+  fromSdk(value: UserKeyState): UserKey {
+    return SymmetricCryptoKey.fromString(value.decrypted_user_key) as UserKey;
+  }
+}

libs/common/src/platform/services/key-state/user-key.state.ts

@@ -11,7 +11,7 @@ export const USER_EVER_HAD_USER_KEY = new UserKeyDefinition<boolean>(
   },
 );
 
-export const USER_KEY = new UserKeyDefinition<UserKey>(CRYPTO_MEMORY, "userKey", {
+export const USER_KEY = UserKeyDefinition.record<UserKey>(CRYPTO_MEMORY, "userKey", {
   deserializer: (obj) => SymmetricCryptoKey.fromJSON(obj) as UserKey,
   clearOn: ["logout", "lock"],
 });

libs/common/src/platform/services/sdk/client-managed-state.ts

@@ -4,6 +4,7 @@ import { UserId } from "@bitwarden/common/types/guid";
 import { CipherRecordMapper } from "@bitwarden/common/vault/models/domain/cipher-sdk-mapper";
 import { StateClient, Repository } from "@bitwarden/sdk-internal";
 
+import { UserKeyRecordMapper } from "../../../key-management/user-key-mapper";
 import { StateProvider, UserKeyDefinition } from "../../state";
 
 export async function initializeState(
@@ -11,9 +12,11 @@ export async function initializeState(
   stateClient: StateClient,
   stateProvider: StateProvider,
 ): Promise<void> {
-  await stateClient.register_cipher_repository(
-    new RepositoryRecord(userId, stateProvider, new CipherRecordMapper()),
-  );
+  stateClient.register_client_managed_repositories({
+    cipher: new RepositoryRecord(userId, stateProvider, new CipherRecordMapper()),
+    folder: null,
+    user_key_state: new RepositoryRecord(userId, stateProvider, new UserKeyRecordMapper()),
+  });
 }
 
 export interface SdkRecordMapper<ClientType, SdkType> {

libs/key-management/src/abstractions/key.service.ts

@@ -57,7 +57,7 @@ export abstract class KeyService {
    * @note this observable represents only user keys stored in memory. A null value does not indicate that we cannot load a user key from storage.
    * @param userId The desired user
    */
-  abstract getInMemoryUserKeyFor$(userId: UserId): Observable<UserKey>;
+  abstract getInMemoryUserKeyFor$(userId: UserId): Observable<UserKey | null>;
   /**
    * Sets the provided user key and stores
    * any other necessary versions (such as auto, biometrics,
@@ -91,7 +91,7 @@ export abstract class KeyService {
    *
    * @deprecated Use {@link userKey$} with a required {@link UserId} instead.
    */
-  abstract getUserKey(userId?: string): Promise<UserKey>;
+  abstract getUserKey(userId?: string): Promise<UserKey | null>;
 
   /**
    * Retrieves the user key from storage

libs/key-management/src/key.service.spec.ts

@@ -91,6 +91,12 @@ describe("keyService", () => {
     );
   });
 
+  const setUserKeyState = (userId: UserId, userKey: UserKey | null) => {
+    stateProvider.singleUser
+      .getFake(userId, USER_KEY)
+      .nextState(userKey == null ? null : ({ "": userKey } as Record<string, UserKey>));
+  };
+
   afterEach(() => {
     jest.resetAllMocks();
   });
@@ -110,7 +116,7 @@ describe("keyService", () => {
     );
 
     it("throws error if user key not found", async () => {
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(null);
+      setUserKeyState(mockUserId, null);
 
       await expect(keyService.refreshAdditionalKeys(mockUserId)).rejects.toThrow(
         "No user key found for: " + mockUserId,
@@ -119,7 +125,7 @@ describe("keyService", () => {
 
     it("refreshes additional keys when user key is available", async () => {
       const mockUserKey = new SymmetricCryptoKey(new Uint8Array(64)) as UserKey;
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(mockUserKey);
+      setUserKeyState(mockUserId, mockUserKey);
       const setUserKeySpy = jest.spyOn(keyService, "setUserKey");
 
       await keyService.refreshAdditionalKeys(mockUserId);
@@ -143,7 +149,7 @@ describe("keyService", () => {
     });
 
     it("returns the User Key if available", async () => {
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(mockUserKey);
+      setUserKeyState(mockUserId, mockUserKey);
 
       const userKey = await keyService.getUserKey(mockUserId);
 
@@ -173,7 +179,7 @@ describe("keyService", () => {
     );
 
     it.each([true, false])("returns %s if the user key is set", async (hasKey) => {
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(hasKey ? mockUserKey : null);
+      setUserKeyState(mockUserId, hasKey ? mockUserKey : null);
       expect(await keyService.hasUserKey(mockUserId)).toBe(hasKey);
     });
   });
@@ -365,7 +371,7 @@ describe("keyService", () => {
       mockUserKey = makeSymmetricCryptoKey<UserKey>(64);
       mockEncryptedPrivateKey = makeEncString("encryptedPrivateKey").encryptedString!;
       mockUserPrivateKey = makeStaticByteArray(10, 1);
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(mockUserKey);
+      setUserKeyState(mockUserId, mockUserKey);
       accountCryptographicStateService.accountCryptographicState$.mockReturnValue(
         of({ V1: { private_key: mockEncryptedPrivateKey } }),
       );
@@ -392,7 +398,7 @@ describe("keyService", () => {
     });
 
     it("returns null if user key is not set", async () => {
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(null);
+      setUserKeyState(mockUserId, null);
 
       const result = await firstValueFrom(keyService.userPrivateKey$(mockUserId));
 
@@ -423,14 +429,14 @@ describe("keyService", () => {
       expect(result).toEqual(mockUserPrivateKey);
 
       // Change user key to null
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(null);
+      setUserKeyState(mockUserId, null);
 
       result = await firstValueFrom(keyService.userPrivateKey$(mockUserId));
 
       expect(result).toBeNull();
 
       // Restore user key, remove encrypted private key
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(mockUserKey);
+      setUserKeyState(mockUserId, mockUserKey);
       accountStateSubject.next(null);
 
       result = await firstValueFrom(keyService.userPrivateKey$(mockUserId));
@@ -480,8 +486,7 @@ describe("keyService", () => {
 
     function updateKeys(keys: Partial<UpdateKeysParams> = {}) {
       if ("userKey" in keys) {
-        const userKeyState = stateProvider.singleUser.getFake(mockUserId, USER_KEY);
-        userKeyState.nextState(keys.userKey!);
+        setUserKeyState(mockUserId, keys.userKey!);
       }
 
       if ("encryptedPrivateKey" in keys) {
@@ -915,7 +920,9 @@ describe("keyService", () => {
       masterPasswordService.masterKeySubject.next(fakeMasterKey);
       userKeyState.nextState(null);
       const fakeUserKey = makeUserKey ? makeSymmetricCryptoKey<UserKey>(64) : null;
-      userKeyState.nextState(fakeUserKey);
+      userKeyState.nextState(
+        fakeUserKey == null ? null : ({ "": fakeUserKey } as Record<string, UserKey>),
+      );
       return [fakeUserKey, fakeMasterKey];
     }
 
@@ -1054,7 +1061,7 @@ describe("keyService", () => {
 
     it("throws when user already has a user key", async () => {
       const existingUserKey = makeSymmetricCryptoKey<UserKey>(64);
-      stateProvider.singleUser.getFake(mockUserId, USER_KEY).nextState(existingUserKey);
+      setUserKeyState(mockUserId, existingUserKey);
 
       await expect(keyService.initAccount(mockUserId)).rejects.toThrow(
         "Cannot initialize account, keys already exist.",

libs/key-management/src/key.service.ts

@@ -68,6 +68,8 @@ import {
 } from "./abstractions/key.service";
 import { KdfConfig } from "./models/kdf-config";
 
+const USER_KEY_STATE_KEY: string = "";
+
 export class DefaultKeyService implements KeyServiceAbstraction {
   /**
    * Retrieves a stream of the active users organization keys,
@@ -108,7 +110,7 @@ export class DefaultKeyService implements KeyServiceAbstraction {
     }
 
     // Set userId to ensure we have one for the account status update
-    await this.stateProvider.setUserState(USER_KEY, key, userId);
+    await this.stateProvider.setUserState(USER_KEY, this.userKeyToStateObject(key), userId);
     await this.stateProvider.setUserState(USER_EVER_HAD_USER_KEY, true, userId);
 
     await this.storeAdditionalKeys(key, userId);
@@ -140,16 +142,18 @@ export class DefaultKeyService implements KeyServiceAbstraction {
       .state$.pipe(map((x) => x ?? false));
   }
 
-  getInMemoryUserKeyFor$(userId: UserId): Observable<UserKey> {
-    return this.stateProvider.getUserState$(USER_KEY, userId);
+  getInMemoryUserKeyFor$(userId: UserId): Observable<UserKey | null> {
+    return this.stateProvider
+      .getUserState$(USER_KEY, userId)
+      .pipe(map((userKey) => this.stateObjectToUserKey(userKey)));
   }
 
   /**
    * @deprecated Use {@link userKey$} with a required {@link UserId} instead.
    */
-  async getUserKey(userId?: UserId): Promise<UserKey> {
+  async getUserKey(userId?: UserId): Promise<UserKey | null> {
     const userKey = await firstValueFrom(this.stateProvider.getUserState$(USER_KEY, userId));
-    return userKey;
+    return this.stateObjectToUserKey(userKey);
   }
 
   async getUserKeyFromStorage(
@@ -640,7 +644,9 @@ export class DefaultKeyService implements KeyServiceAbstraction {
   }
 
   userKey$(userId: UserId): Observable<UserKey | null> {
-    return this.stateProvider.getUser(userId, USER_KEY).state$;
+    return this.stateProvider
+      .getUser(userId, USER_KEY)
+      .state$.pipe(map((key) => (key != null ? (key[""] as UserKey) : null)));
   }
 
   userPublicKey$(userId: UserId) {
@@ -919,4 +925,18 @@ export class DefaultKeyService implements KeyServiceAbstraction {
       }),
     );
   }
+
+  private userKeyToStateObject(userKey: UserKey | null): Record<string, UserKey> | null {
+    if (userKey == null) {
+      return null;
+    }
+    return { [USER_KEY_STATE_KEY]: userKey };
+  }
+
+  private stateObjectToUserKey(stateObject: Record<string, UserKey> | null): UserKey | null {
+    if (stateObject == null) {
+      return null;
+    }
+    return stateObject[USER_KEY_STATE_KEY] ?? null;
+  }
 }

libs/vault/src/components/add-edit-folder-dialog/add-edit-folder-dialog.component.ts

@@ -127,7 +127,7 @@ export class AddEditFolderDialogComponent implements AfterViewInit, OnInit {
 
     try {
       const activeUserId = await firstValueFrom(this.activeUserId$);
-      const userKey = await this.keyService.getUserKey(activeUserId!);
+      const userKey = (await this.keyService.getUserKey(activeUserId!))!;
       const folder = await this.folderService.encrypt(this.folder, userKey);
       await this.folderApiService.save(folder, activeUserId!);