[PM-37190] Clean up ephemeral states that go via state bridge

libs/common/src/key-management/pin/pin-state.service.implementation.ts

@@ -15,8 +15,6 @@ import {
   USER_KEY_ENCRYPTED_PIN,
 } from "./pin.state";
 
-const EPHEMERAL_PIN_ENVELOPE_KEY = "";
-
 export class PinStateService implements PinStateServiceAbstraction {
   constructor(private stateProvider: StateProvider) {}
 
@@ -82,7 +80,7 @@ export class PinStateService implements PinStateServiceAbstraction {
     if (pinLockType === "EPHEMERAL") {
       await this.stateProvider.setUserState(
         PIN_PROTECTED_USER_KEY_ENVELOPE_EPHEMERAL,
-        { [EPHEMERAL_PIN_ENVELOPE_KEY]: { pin_envelope: pinProtectedUserKeyEnvelope } },
+        pinProtectedUserKeyEnvelope,
         userId,
       );
     } else if (pinLockType === "PERSISTENT") {
@@ -121,7 +119,7 @@ export class PinStateService implements PinStateServiceAbstraction {
     if (pinLockType === "EPHEMERAL") {
       return this.stateProvider
         .getUserState$(PIN_PROTECTED_USER_KEY_ENVELOPE_EPHEMERAL, userId)
-        .pipe(map((record) => record?.[EPHEMERAL_PIN_ENVELOPE_KEY]?.pin_envelope ?? null));
+        .pipe(map((envelope) => envelope ?? null));
     } else if (pinLockType === "PERSISTENT") {
       return this.stateProvider.getUserState$(PIN_PROTECTED_USER_KEY_ENVELOPE_PERSISTENT, userId);
     } else {

libs/common/src/key-management/pin/pin.state.ts

@@ -1,5 +1,5 @@
 import { PIN_DISK, PIN_MEMORY, UserKeyDefinition } from "@bitwarden/common/platform/state";
-import { PasswordProtectedKeyEnvelope, EphemeralPinEnvelopeState } from "@bitwarden/sdk-internal";
+import { PasswordProtectedKeyEnvelope } from "@bitwarden/sdk-internal";
 
 import { EncryptedString } from "../crypto/models/enc-string";
 
@@ -23,7 +23,7 @@ export const PIN_PROTECTED_USER_KEY_ENVELOPE_PERSISTENT =
  * The ephemeral (stored in memory) version of the UserKey, stored in a `PasswordProtectedKeyEnvelope`.
  */
 export const PIN_PROTECTED_USER_KEY_ENVELOPE_EPHEMERAL =
-  UserKeyDefinition.record<EphemeralPinEnvelopeState>(
+  new UserKeyDefinition<PasswordProtectedKeyEnvelope>(
     PIN_MEMORY,
     "pinProtectedUserKeyEnvelopeEphemeral",
     {

libs/common/src/key-management/state-bridge.ts

@@ -122,15 +122,18 @@ export class JsWasmStateBridge implements WasmStateBridge {
   }
 
   async set_user_key(userKey: SymmetricKey): Promise<void> {
-    await writeAtomic(this.stateProvider, this.userId, USER_KEY, {
-      "": SymmetricCryptoKey.fromSdk(userKey) as UserKey,
-    });
+    await writeAtomic(
+      this.stateProvider,
+      this.userId,
+      USER_KEY,
+      SymmetricCryptoKey.fromSdk(userKey) as UserKey,
+    );
   }
 
   async get_user_key(): Promise<SymmetricKey | null> {
     const key = await readAtomic(this.stateProvider, this.userId, USER_KEY);
     if (key != null) {
-      return key[""].toSdk();
+      return key.toSdk();
     } else {
       return null;
     }
@@ -141,22 +144,20 @@ export class JsWasmStateBridge implements WasmStateBridge {
   }
 
   async set_ephemeral_pin_envelope(pinEnvelope: PasswordProtectedKeyEnvelope): Promise<void> {
-    await writeAtomic(this.stateProvider, this.userId, PIN_PROTECTED_USER_KEY_ENVELOPE_EPHEMERAL, {
-      "": { pin_envelope: pinEnvelope },
-    });
+    await writeAtomic(
+      this.stateProvider,
+      this.userId,
+      PIN_PROTECTED_USER_KEY_ENVELOPE_EPHEMERAL,
+      pinEnvelope,
+    );
   }
 
   async get_ephemeral_pin_envelope(): Promise<PasswordProtectedKeyEnvelope | null> {
-    const result = await readAtomic(
+    return await readAtomic(
       this.stateProvider,
       this.userId,
       PIN_PROTECTED_USER_KEY_ENVELOPE_EPHEMERAL,
     );
-    if (result != null) {
-      return result[""]?.pin_envelope ?? null;
-    } else {
-      return null;
-    }
   }
 
   async clear_ephemeral_pin_envelope(): Promise<void> {

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

@@ -22,7 +22,7 @@ export const USER_EVER_HAD_USER_KEY = new UserKeyDefinition<boolean>(
   },
 );
 
-export const USER_KEY = UserKeyDefinition.record<UserKey>(CRYPTO_MEMORY, "userKey", {
+export const USER_KEY = new UserKeyDefinition<UserKey>(CRYPTO_MEMORY, "userKey", {
   deserializer: (obj) => SymmetricCryptoKey.fromJSON(obj) as UserKey,
   clearOn: ["logout", "lock"],
   // Prevents the state from caching and rxjs observable becoming hot observable.

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

@@ -4,9 +4,7 @@ import { UserId } from "@bitwarden/common/types/guid";
 import { CipherRecordMapper } from "@bitwarden/common/vault/models/domain/cipher-sdk-mapper";
 import { Repository, StateClient } from "@bitwarden/sdk-internal";
 
-import { EphemeralPinEnvelopeMapper } from "../../../key-management/ephemeral-pin-envelope-mapper";
 import { LocalUserDataKeyRecordMapper } from "../../../key-management/local-user-data-key-mapper";
-import { UserKeyRecordMapper } from "../../../key-management/user-key-mapper";
 import { StateProvider, UserKeyDefinition } from "../../state";
 
 export async function initializeClientManagedState(
@@ -17,17 +15,13 @@ export async function initializeClientManagedState(
   stateClient.register_client_managed_repositories({
     cipher: new RepositoryRecord(userId, stateProvider, new CipherRecordMapper(), true),
     folder: null,
-    user_key_state: new RepositoryRecord(userId, stateProvider, new UserKeyRecordMapper()),
+    user_key_state: null,
     local_user_data_key_state: new RepositoryRecord(
       userId,
       stateProvider,
       new LocalUserDataKeyRecordMapper(),
     ),
-    ephemeral_pin_envelope_state: new RepositoryRecord(
-      userId,
-      stateProvider,
-      new EphemeralPinEnvelopeMapper(),
-    ),
+    ephemeral_pin_envelope_state: null,
     organization_shared_key: null,
   });
 }

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

@@ -92,9 +92,7 @@ describe("keyService", () => {
   });
 
   const setUserKeyState = (userId: UserId, userKey: UserKey | null) => {
-    stateProvider.singleUser
-      .getFake(userId, USER_KEY)
-      .nextState(userKey == null ? null : ({ "": userKey } as Record<string, UserKey>));
+    stateProvider.singleUser.getFake(userId, USER_KEY).nextState(userKey);
   };
 
   afterEach(() => {
@@ -899,9 +897,7 @@ describe("keyService", () => {
       masterPasswordService.masterKeySubject.next(fakeMasterKey);
       userKeyState.nextState(null);
       const fakeUserKey = makeUserKey ? makeSymmetricCryptoKey<UserKey>(64) : null;
-      userKeyState.nextState(
-        fakeUserKey == null ? null : ({ "": fakeUserKey } as Record<string, UserKey>),
-      );
+      userKeyState.nextState(fakeUserKey);
       return [fakeUserKey, fakeMasterKey];
     }
 

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

@@ -68,8 +68,6 @@ 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,
@@ -110,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, this.userKeyToStateObject(key), userId);
+    await this.stateProvider.setUserState(USER_KEY, key, userId);
     await this.stateProvider.setUserState(USER_EVER_HAD_USER_KEY, true, userId);
 
     await this.storeAdditionalKeys(key, userId);
@@ -145,15 +143,14 @@ export class DefaultKeyService implements KeyServiceAbstraction {
   getInMemoryUserKeyFor$(userId: UserId): Observable<UserKey | null> {
     return this.stateProvider
       .getUserState$(USER_KEY, userId)
-      .pipe(map((userKey) => this.stateObjectToUserKey(userKey)));
+      .pipe(map((userKey) => userKey ?? null));
   }
 
   /**
    * @deprecated Use {@link userKey$} with a required {@link UserId} instead.
    */
   async getUserKey(userId?: UserId): Promise<UserKey | null> {
-    const userKey = await firstValueFrom(this.stateProvider.getUserState$(USER_KEY, userId));
-    return this.stateObjectToUserKey(userKey);
+    return (await firstValueFrom(this.stateProvider.getUserState$(USER_KEY, userId))) ?? null;
   }
 
   async getUserKeyFromStorage(
@@ -640,9 +637,7 @@ export class DefaultKeyService implements KeyServiceAbstraction {
   }
 
   userKey$(userId: UserId): Observable<UserKey | null> {
-    return this.stateProvider
-      .getUser(userId, USER_KEY)
-      .state$.pipe(map((key) => (key != null ? (key[""] as UserKey) : null)));
+    return this.stateProvider.getUser(userId, USER_KEY).state$.pipe(map((key) => key ?? null));
   }
 
   userPublicKey$(userId: UserId) {
@@ -921,18 +916,4 @@ 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;
-  }
 }