Set userkey to state from SDK

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 { USER_KEY } from "../platform/services/key-state/user-key.state";
+import { UserKey } from "../types/key";
+import { SymmetricCryptoKey } from "../platform/models/domain/symmetric-crypto-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

@@ -3,6 +3,7 @@ import { SignedPublicKey, WrappedSigningKey } from "../../../key-management/type
 import { UserKey } from "../../../types/key";
 import { SymmetricCryptoKey } from "../../models/domain/symmetric-crypto-key";
 import { CRYPTO_DISK, CRYPTO_MEMORY, UserKeyDefinition } from "../../state";
+import { record } from "@bitwarden/serialization";
 
 export const USER_EVER_HAD_USER_KEY = new UserKeyDefinition<boolean>(
   CRYPTO_DISK,
@@ -22,7 +23,8 @@ export const USER_ENCRYPTED_PRIVATE_KEY = new UserKeyDefinition<EncryptedString>
   },
 );
 
-export const USER_KEY = new UserKeyDefinition<UserKey>(CRYPTO_MEMORY, "userKey", {
+// This is a map with a single entry to conform with the repository pattern in the SDK.
+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

@@ -5,14 +5,19 @@ import { CipherRecordMapper } from "@bitwarden/common/vault/models/domain/cipher
 import { StateClient, Repository } from "@bitwarden/sdk-internal";
 
 import { StateProvider, UserKeyDefinition } from "../../state";
+import { UserKeyRecordMapper } from "@bitwarden/common/key-management/user-key-mapper";
 
 export async function initializeState(
   userId: UserId,
   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()),
+    }
   );
 }
 

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

@@ -108,7 +108,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);
@@ -165,7 +165,7 @@ export class DefaultKeyService implements KeyServiceAbstraction {
   }
 
   getInMemoryUserKeyFor$(userId: UserId): Observable<UserKey> {
-    return this.stateProvider.getUserState$(USER_KEY, userId);
+    return this.stateProvider.getUserState$(USER_KEY, userId).pipe(map(userKey => this.stateObjectToUserKey(userKey)));
   }
 
   /**
@@ -173,7 +173,7 @@ export class DefaultKeyService implements KeyServiceAbstraction {
    */
   async getUserKey(userId?: UserId): Promise<UserKey> {
     const userKey = await firstValueFrom(this.stateProvider.getUserState$(USER_KEY, userId));
-    return userKey;
+    return this.stateObjectToUserKey(userKey);
   }
 
   async getUserKeyFromStorage(
@@ -751,7 +751,8 @@ 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) {
@@ -1016,4 +1017,18 @@ export class DefaultKeyService implements KeyServiceAbstraction {
   userSignedPublicKey$(userId: UserId): Observable<SignedPublicKey | null> {
     return this.stateProvider.getUserState$(USER_SIGNED_PUBLIC_KEY, userId);
   }
+
+  private userKeyToStateObject(userKey: UserKey | null): Record<string, UserKey> | null {
+    if (userKey == null) {
+      return null;
+    }
+    return { "": userKey };
+  }
+
+  private stateObjectToUserKey(stateObject: Record<string, UserKey> | null): UserKey | null {
+    if (stateObject == null) {
+      return null;
+    }
+    return stateObject[""] ?? null;
+  }
 }