[PM-5533] Migrate Asymmetric User Keys to State Providers (#7665)

libs/common/src/platform/abstractions/state.service.ts

@@ -225,8 +225,6 @@ export abstract class StateService<T extends Account = Account> {
    * @deprecated Do not call this, use PolicyService
    */
   setDecryptedPolicies: (value: Policy[], options?: StorageOptions) => Promise<void>;
-  getDecryptedPrivateKey: (options?: StorageOptions) => Promise<Uint8Array>;
-  setDecryptedPrivateKey: (value: Uint8Array, options?: StorageOptions) => Promise<void>;
   /**
    * @deprecated Do not call this directly, use SendService
    */
@@ -346,8 +344,6 @@ export abstract class StateService<T extends Account = Account> {
     value: { [id: string]: PolicyData },
     options?: StorageOptions,
   ) => Promise<void>;
-  getEncryptedPrivateKey: (options?: StorageOptions) => Promise<string>;
-  setEncryptedPrivateKey: (value: string, options?: StorageOptions) => Promise<void>;
   /**
    * @deprecated Do not call this directly, use SendService
    */
@@ -434,8 +430,6 @@ export abstract class StateService<T extends Account = Account> {
   setProtectedPin: (value: string, options?: StorageOptions) => Promise<void>;
   getProviders: (options?: StorageOptions) => Promise<{ [id: string]: ProviderData }>;
   setProviders: (value: { [id: string]: ProviderData }, options?: StorageOptions) => Promise<void>;
-  getPublicKey: (options?: StorageOptions) => Promise<Uint8Array>;
-  setPublicKey: (value: Uint8Array, options?: StorageOptions) => Promise<void>;
   getRefreshToken: (options?: StorageOptions) => Promise<string>;
   setRefreshToken: (value: string, options?: StorageOptions) => Promise<void>;
   getRememberedEmail: (options?: StorageOptions) => Promise<string>;

libs/common/src/platform/models/domain/account.ts

@@ -125,7 +125,6 @@ export class AccountKeys {
   masterKey?: MasterKey;
   masterKeyEncryptedUserKey?: string;
   deviceKey?: ReturnType<SymmetricCryptoKey["toJSON"]>;
-  privateKey?: EncryptionPair<string, Uint8Array> = new EncryptionPair<string, Uint8Array>();
   publicKey?: Uint8Array;
   apiKeyClientSecret?: string;
 
@@ -163,9 +162,6 @@ export class AccountKeys {
         obj?.cryptoSymmetricKey,
         SymmetricCryptoKey.fromJSON,
       ),
-      privateKey: EncryptionPair.fromJSON<string, Uint8Array>(obj?.privateKey, (decObj: string) =>
-        Utils.fromByteStringToArray(decObj),
-      ),
       publicKey: Utils.fromByteStringToArray(obj?.publicKey),
     });
   }

libs/common/src/platform/services/crypto.service.ts

@@ -9,7 +9,16 @@ import { AccountService } from "../../auth/abstractions/account.service";
 import { KdfConfig } from "../../auth/models/domain/kdf-config";
 import { Utils } from "../../platform/misc/utils";
 import { OrganizationId, ProviderId, UserId } from "../../types/guid";
-import { OrgKey, UserKey, MasterKey, ProviderKey, PinKey, CipherKey } from "../../types/key";
+import {
+  OrgKey,
+  UserKey,
+  MasterKey,
+  ProviderKey,
+  PinKey,
+  CipherKey,
+  UserPrivateKey,
+  UserPublicKey,
+} from "../../types/key";
 import { CryptoFunctionService } from "../abstractions/crypto-function.service";
 import { CryptoService as CryptoServiceAbstraction } from "../abstractions/crypto.service";
 import { EncryptService } from "../abstractions/encrypt.service";
@@ -38,7 +47,12 @@ import {
   USER_ORGANIZATION_KEYS,
 } from "./key-state/org-keys.state";
 import { USER_ENCRYPTED_PROVIDER_KEYS, USER_PROVIDER_KEYS } from "./key-state/provider-keys.state";
-import { USER_EVER_HAD_USER_KEY } from "./key-state/user-key.state";
+import {
+  USER_ENCRYPTED_PRIVATE_KEY,
+  USER_EVER_HAD_USER_KEY,
+  USER_PRIVATE_KEY,
+  USER_PUBLIC_KEY,
+} from "./key-state/user-key.state";
 
 export class CryptoService implements CryptoServiceAbstraction {
   private readonly activeUserEverHadUserKey: ActiveUserState<boolean>;
@@ -49,12 +63,16 @@ export class CryptoService implements CryptoServiceAbstraction {
   private readonly activeUserEncryptedProviderKeysState: ActiveUserState<
     Record<ProviderId, EncryptedString>
   >;
-  private readonly activeUserProviderKeysState: DerivedState<Record<OrganizationId, ProviderKey>>;
+  private readonly activeUserProviderKeysState: DerivedState<Record<ProviderId, ProviderKey>>;
+  private readonly activeUserEncryptedPrivateKeyState: ActiveUserState<EncryptedString>;
+  private readonly activeUserPrivateKeyState: DerivedState<UserPrivateKey>;
+  private readonly activeUserPublicKeyState: DerivedState<UserPublicKey>;
 
   readonly activeUserOrgKeys$: Observable<Record<OrganizationId, OrgKey>>;
   readonly activeUserProviderKeys$: Observable<Record<ProviderId, ProviderKey>>;
-
-  readonly everHadUserKey$;
+  readonly activeUserPrivateKey$: Observable<UserPrivateKey>;
+  readonly activeUserPublicKey$: Observable<UserPublicKey>;
+  readonly everHadUserKey$: Observable<boolean>;
 
   constructor(
     protected cryptoFunctionService: CryptoFunctionService,
@@ -65,7 +83,31 @@ export class CryptoService implements CryptoServiceAbstraction {
     protected accountService: AccountService,
     protected stateProvider: StateProvider,
   ) {
+    // User Key
     this.activeUserEverHadUserKey = stateProvider.getActive(USER_EVER_HAD_USER_KEY);
+    this.everHadUserKey$ = this.activeUserEverHadUserKey.state$.pipe(map((x) => x ?? false));
+
+    // User Asymmetric Key Pair
+    this.activeUserEncryptedPrivateKeyState = stateProvider.getActive(USER_ENCRYPTED_PRIVATE_KEY);
+    this.activeUserPrivateKeyState = stateProvider.getDerived(
+      this.activeUserEncryptedPrivateKeyState.combinedState$,
+      USER_PRIVATE_KEY,
+      {
+        encryptService: this.encryptService,
+        cryptoService: this,
+      },
+    );
+    this.activeUserPrivateKey$ = this.activeUserPrivateKeyState.state$; // may be null
+    this.activeUserPublicKeyState = stateProvider.getDerived(
+      this.activeUserPrivateKey$,
+      USER_PUBLIC_KEY,
+      {
+        cryptoFunctionService: this.cryptoFunctionService,
+      },
+    );
+    this.activeUserPublicKey$ = this.activeUserPublicKeyState.state$; // may be null
+
+    // Organization keys
     this.activeUserEncryptedOrgKeysState = stateProvider.getActive(
       USER_ENCRYPTED_ORGANIZATION_KEYS,
     );
@@ -74,6 +116,9 @@ export class CryptoService implements CryptoServiceAbstraction {
       USER_ORGANIZATION_KEYS,
       { cryptoService: this },
     );
+    this.activeUserOrgKeys$ = this.activeUserOrgKeysState.state$; // null handled by `derive` function
+
+    // Provider keys
     this.activeUserEncryptedProviderKeysState = stateProvider.getActive(
       USER_ENCRYPTED_PROVIDER_KEYS,
     );
@@ -82,9 +127,6 @@ export class CryptoService implements CryptoServiceAbstraction {
       USER_PROVIDER_KEYS,
       { encryptService: this.encryptService, cryptoService: this },
     );
-
-    this.everHadUserKey$ = this.activeUserEverHadUserKey.state$.pipe(map((x) => x ?? false));
-    this.activeUserOrgKeys$ = this.activeUserOrgKeysState.state$; // null handled by `derive` function
     this.activeUserProviderKeys$ = this.activeUserProviderKeysState.state$; // null handled by `derive` function
   }
 
@@ -465,19 +507,7 @@ export class CryptoService implements CryptoServiceAbstraction {
   }
 
   async getPublicKey(): Promise<Uint8Array> {
-    const inMemoryPublicKey = await this.stateService.getPublicKey();
-    if (inMemoryPublicKey != null) {
-      return inMemoryPublicKey;
-    }
-
-    const privateKey = await this.getPrivateKey();
-    if (privateKey == null) {
-      return null;
-    }
-
-    const publicKey = await this.cryptoFunctionService.rsaExtractPublicKey(privateKey);
-    await this.stateService.setPublicKey(publicKey);
-    return publicKey;
+    return await firstValueFrom(this.activeUserPublicKey$);
   }
 
   async makeOrgKey<T extends OrgKey | ProviderKey>(): Promise<[EncString, T]> {
@@ -487,32 +517,16 @@ export class CryptoService implements CryptoServiceAbstraction {
     return [encShareKey, new SymmetricCryptoKey(shareKey) as T];
   }
 
-  async setPrivateKey(encPrivateKey: string): Promise<void> {
+  async setPrivateKey(encPrivateKey: EncryptedString): Promise<void> {
     if (encPrivateKey == null) {
       return;
     }
 
-    await this.stateService.setDecryptedPrivateKey(null);
-    await this.stateService.setEncryptedPrivateKey(encPrivateKey);
+    await this.activeUserEncryptedPrivateKeyState.update(() => encPrivateKey);
   }
 
   async getPrivateKey(): Promise<Uint8Array> {
-    const decryptedPrivateKey = await this.stateService.getDecryptedPrivateKey();
-    if (decryptedPrivateKey != null) {
-      return decryptedPrivateKey;
-    }
-
-    const encPrivateKey = await this.stateService.getEncryptedPrivateKey();
-    if (encPrivateKey == null) {
-      return null;
-    }
-
-    const privateKey = await this.encryptService.decryptToBytes(
-      new EncString(encPrivateKey),
-      await this.getUserKeyWithLegacySupport(),
-    );
-    await this.stateService.setDecryptedPrivateKey(privateKey);
-    return privateKey;
+    return await firstValueFrom(this.activeUserPrivateKey$);
   }
 
   async getFingerprint(fingerprintMaterial: string, publicKey?: Uint8Array): Promise<string[]> {
@@ -543,14 +557,23 @@ export class CryptoService implements CryptoServiceAbstraction {
   }
 
   async clearKeyPair(memoryOnly?: boolean, userId?: UserId): Promise<void[]> {
-    const keysToClear: Promise<void>[] = [
-      this.stateService.setDecryptedPrivateKey(null, { userId: userId }),
-      this.stateService.setPublicKey(null, { userId: userId }),
-    ];
-    if (!memoryOnly) {
-      keysToClear.push(this.stateService.setEncryptedPrivateKey(null, { userId: userId }));
+    const activeUserId = (await firstValueFrom(this.accountService.activeAccount$))?.id;
+    const userIdIsActive = userId == null || userId === activeUserId;
+    if (memoryOnly && userIdIsActive) {
+      // key pair is only cached for active users
+      await this.activeUserPrivateKeyState.forceValue(null);
+      await this.activeUserPublicKeyState.forceValue(null);
+      return;
+    } else {
+      if (userId == null && activeUserId == null) {
+        // nothing to do
+        return;
+      }
+      // below updates decrypted private key and public keys if this is the active user as well since those are derived from the encrypted private key
+      await this.stateProvider
+        .getUser(userId ?? activeUserId, USER_ENCRYPTED_PRIVATE_KEY)
+        .update(() => null);
     }
-    return Promise.all(keysToClear);
   }
 
   async makePinKey(pin: string, salt: string, kdf: KdfType, kdfConfig: KdfConfig): Promise<PinKey> {
@@ -735,16 +758,23 @@ export class CryptoService implements CryptoServiceAbstraction {
     }
 
     try {
-      const encPrivateKey = await this.stateService.getEncryptedPrivateKey();
+      const [userId, encPrivateKey] = await firstValueFrom(
+        this.activeUserEncryptedPrivateKeyState.combinedState$,
+      );
       if (encPrivateKey == null) {
         return false;
       }
 
-      const privateKey = await this.encryptService.decryptToBytes(
-        new EncString(encPrivateKey),
-        key,
-      );
-      await this.cryptoFunctionService.rsaExtractPublicKey(privateKey);
+      // Can decrypt private key
+      const privateKey = await USER_PRIVATE_KEY.derive([userId, encPrivateKey], {
+        encryptService: this.encryptService,
+        cryptoService: this,
+      });
+
+      // Can successfully derive public key
+      await USER_PUBLIC_KEY.derive(privateKey, {
+        cryptoFunctionService: this.cryptoFunctionService,
+      });
     } catch (e) {
       return false;
     }
@@ -765,7 +795,7 @@ export class CryptoService implements CryptoServiceAbstraction {
     const userKey = new SymmetricCryptoKey(rawKey) as UserKey;
     const [publicKey, privateKey] = await this.makeKeyPair(userKey);
     await this.setUserKey(userKey);
-    await this.stateService.setEncryptedPrivateKey(privateKey.encryptedString);
+    await this.activeUserEncryptedPrivateKeyState.update(() => privateKey.encryptedString);
 
     return {
       userKey,

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

@@ -0,0 +1,130 @@
+import { mock } from "jest-mock-extended";
+
+import { makeStaticByteArray } from "../../../../spec";
+import { UserId } from "../../../types/guid";
+import { UserKey, UserPrivateKey, UserPublicKey } from "../../../types/key";
+import { CryptoFunctionService } from "../../abstractions/crypto-function.service";
+import { EncryptService } from "../../abstractions/encrypt.service";
+import { EncryptionType } from "../../enums";
+import { Utils } from "../../misc/utils";
+import { EncString } from "../../models/domain/enc-string";
+import { CryptoService } from "../crypto.service";
+
+import {
+  USER_ENCRYPTED_PRIVATE_KEY,
+  USER_EVER_HAD_USER_KEY,
+  USER_PRIVATE_KEY,
+  USER_PUBLIC_KEY,
+} from "./user-key.state";
+
+function makeEncString(data?: string) {
+  data ??= Utils.newGuid();
+  return new EncString(EncryptionType.AesCbc256_HmacSha256_B64, data, "test", "test");
+}
+
+describe("Ever had user key", () => {
+  const sut = USER_EVER_HAD_USER_KEY;
+
+  it("should deserialize ever had user key", () => {
+    const everHadUserKey = true;
+
+    const result = sut.deserializer(JSON.parse(JSON.stringify(everHadUserKey)));
+
+    expect(result).toEqual(everHadUserKey);
+  });
+});
+
+describe("Encrypted private key", () => {
+  const sut = USER_ENCRYPTED_PRIVATE_KEY;
+
+  it("should deserialize encrypted private key", () => {
+    const encryptedPrivateKey = makeEncString().encryptedString;
+
+    const result = sut.deserializer(JSON.parse(JSON.stringify(encryptedPrivateKey)));
+
+    expect(result).toEqual(encryptedPrivateKey);
+  });
+});
+
+describe("User public key", () => {
+  const sut = USER_PUBLIC_KEY;
+  const userPrivateKey = makeStaticByteArray(64, 1) as UserPrivateKey;
+  const userPublicKey = makeStaticByteArray(64, 2) as UserPublicKey;
+
+  it("should deserialize user public key", () => {
+    const userPublicKey = makeStaticByteArray(64, 1);
+
+    const result = sut.deserialize(JSON.parse(JSON.stringify(userPublicKey)));
+
+    expect(result).toEqual(userPublicKey);
+  });
+
+  it("should derive user public key", async () => {
+    const cryptoFunctionService = mock<CryptoFunctionService>();
+    cryptoFunctionService.rsaExtractPublicKey.mockResolvedValue(userPublicKey);
+
+    const result = await sut.derive(userPrivateKey, { cryptoFunctionService });
+
+    expect(result).toEqual(userPublicKey);
+  });
+});
+
+describe("Derived decrypted private key", () => {
+  const sut = USER_PRIVATE_KEY;
+  const userId = "userId" as UserId;
+  const userKey = mock<UserKey>();
+  const encryptedPrivateKey = makeEncString().encryptedString;
+  const decryptedPrivateKey = makeStaticByteArray(64, 1);
+
+  afterEach(() => {
+    jest.resetAllMocks();
+  });
+
+  it("should deserialize decrypted private key", () => {
+    const decryptedPrivateKey = makeStaticByteArray(64, 1);
+
+    const result = sut.deserialize(JSON.parse(JSON.stringify(decryptedPrivateKey)));
+
+    expect(result).toEqual(decryptedPrivateKey);
+  });
+
+  it("should derive decrypted private key", async () => {
+    const cryptoService = mock<CryptoService>();
+    cryptoService.getUserKey.mockResolvedValue(userKey);
+    const encryptService = mock<EncryptService>();
+    encryptService.decryptToBytes.mockResolvedValue(decryptedPrivateKey);
+
+    const result = await sut.derive([userId, encryptedPrivateKey], {
+      encryptService,
+      cryptoService,
+    });
+
+    expect(result).toEqual(decryptedPrivateKey);
+  });
+
+  it("should handle null input values", async () => {
+    const cryptoService = mock<CryptoService>();
+    cryptoService.getUserKey.mockResolvedValue(userKey);
+    const encryptService = mock<EncryptService>();
+
+    const result = await sut.derive([userId, null], {
+      encryptService,
+      cryptoService,
+    });
+
+    expect(result).toEqual(null);
+  });
+
+  it("should handle null user key", async () => {
+    const cryptoService = mock<CryptoService>();
+    cryptoService.getUserKey.mockResolvedValue(null);
+    const encryptService = mock<EncryptService>();
+
+    const result = await sut.derive([userId, encryptedPrivateKey], {
+      encryptService,
+      cryptoService,
+    });
+
+    expect(result).toEqual(null);
+  });
+});

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

@@ -1,5 +1,59 @@
-import { KeyDefinition, CRYPTO_DISK } from "../../state";
+import { UserPrivateKey, UserPublicKey } from "../../../types/key";
+import { CryptoFunctionService } from "../../abstractions/crypto-function.service";
+import { EncryptService } from "../../abstractions/encrypt.service";
+import { EncString, EncryptedString } from "../../models/domain/enc-string";
+import { KeyDefinition, CRYPTO_DISK, DeriveDefinition } from "../../state";
+import { CryptoService } from "../crypto.service";
 
 export const USER_EVER_HAD_USER_KEY = new KeyDefinition<boolean>(CRYPTO_DISK, "everHadUserKey", {
   deserializer: (obj) => obj,
 });
+
+export const USER_ENCRYPTED_PRIVATE_KEY = new KeyDefinition<EncryptedString>(
+  CRYPTO_DISK,
+  "privateKey",
+  {
+    deserializer: (obj) => obj,
+  },
+);
+
+export const USER_PRIVATE_KEY = DeriveDefinition.fromWithUserId<
+  EncryptedString,
+  UserPrivateKey,
+  // TODO: update cryptoService to user key directly
+  { encryptService: EncryptService; cryptoService: CryptoService }
+>(USER_ENCRYPTED_PRIVATE_KEY, {
+  deserializer: (obj) => new Uint8Array(Object.values(obj)) as UserPrivateKey,
+  derive: async ([userId, encPrivateKeyString], { encryptService, cryptoService }) => {
+    if (encPrivateKeyString == null) {
+      return null;
+    }
+
+    const userKey = await cryptoService.getUserKey(userId);
+    if (userKey == null) {
+      return null;
+    }
+
+    const encPrivateKey = new EncString(encPrivateKeyString);
+    const privateKey = (await encryptService.decryptToBytes(
+      encPrivateKey,
+      userKey,
+    )) as UserPrivateKey;
+    return privateKey;
+  },
+});
+
+export const USER_PUBLIC_KEY = DeriveDefinition.from<
+  UserPrivateKey,
+  UserPublicKey,
+  { cryptoFunctionService: CryptoFunctionService }
+>([USER_PRIVATE_KEY, "publicKey"], {
+  deserializer: (obj) => new Uint8Array(Object.values(obj)) as UserPublicKey,
+  derive: async (privateKey, { cryptoFunctionService }) => {
+    if (privateKey == null) {
+      return null;
+    }
+
+    return (await cryptoFunctionService.rsaExtractPublicKey(privateKey)) as UserPublicKey;
+  },
+});

libs/common/src/platform/services/state.service.ts

@@ -1048,23 +1048,6 @@ export class StateService<
     );
   }
 
-  async getDecryptedPrivateKey(options?: StorageOptions): Promise<Uint8Array> {
-    return (
-      await this.getAccount(this.reconcileOptions(options, await this.defaultInMemoryOptions()))
-    )?.keys?.privateKey.decrypted;
-  }
-
-  async setDecryptedPrivateKey(value: Uint8Array, options?: StorageOptions): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-    account.keys.privateKey.decrypted = value;
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-  }
-
   @withPrototypeForArrayMembers(SendView)
   async getDecryptedSends(options?: StorageOptions): Promise<SendView[]> {
     return (
@@ -1752,24 +1735,6 @@ export class StateService<
     );
   }
 
-  async getEncryptedPrivateKey(options?: StorageOptions): Promise<string> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultOnDiskOptions()),
-    );
-    return account?.keys?.privateKey?.encrypted;
-  }
-
-  async setEncryptedPrivateKey(value: string, options?: StorageOptions): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultOnDiskOptions()),
-    );
-    account.keys.privateKey.encrypted = value;
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultOnDiskOptions()),
-    );
-  }
-
   @withPrototypeForObjectValues(SendData)
   async getEncryptedSends(options?: StorageOptions): Promise<{ [id: string]: SendData }> {
     return (
@@ -2274,24 +2239,6 @@ export class StateService<
     );
   }
 
-  async getPublicKey(options?: StorageOptions): Promise<Uint8Array> {
-    const keys = (
-      await this.getAccount(this.reconcileOptions(options, await this.defaultInMemoryOptions()))
-    )?.keys;
-    return keys?.publicKey;
-  }
-
-  async setPublicKey(value: Uint8Array, options?: StorageOptions): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-    account.keys.publicKey = value;
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-  }
-
   async getRefreshToken(options?: StorageOptions): Promise<string> {
     options = await this.getTimeoutBasedStorageOptions(options);
     return (await this.getAccount(options))?.tokens?.refreshToken;

libs/common/src/platform/state/derive-definition.spec.ts

@@ -0,0 +1,42 @@
+import { DeriveDefinition } from "./derive-definition";
+import { KeyDefinition } from "./key-definition";
+import { StateDefinition } from "./state-definition";
+
+const derive: () => any = () => null;
+const deserializer: any = (obj: any) => obj;
+
+const STATE_DEFINITION = new StateDefinition("test", "disk");
+const TEST_KEY = new KeyDefinition(STATE_DEFINITION, "test", {
+  deserializer,
+});
+const TEST_DERIVE = new DeriveDefinition(STATE_DEFINITION, "test", {
+  derive,
+  deserializer,
+});
+
+describe("DeriveDefinition", () => {
+  describe("from", () => {
+    it("should create a new DeriveDefinition from a KeyDefinition", () => {
+      const result = DeriveDefinition.from(TEST_KEY, {
+        derive,
+        deserializer,
+      });
+
+      expect(result).toEqual(TEST_DERIVE);
+    });
+
+    it("should create a new DeriveDefinition from a DeriveDefinition", () => {
+      const result = DeriveDefinition.from([TEST_DERIVE, "newDerive"], {
+        derive,
+        deserializer,
+      });
+
+      expect(result).toEqual(
+        new DeriveDefinition(STATE_DEFINITION, "newDerive", {
+          derive,
+          deserializer,
+        }),
+      );
+    });
+  });
+});

libs/common/src/platform/state/derive-definition.ts

@@ -1,5 +1,6 @@
 import { Jsonify } from "type-fest";
 
+import { UserId } from "../../types/guid";
 import { DerivedStateDependencies, StorageKey } from "../../types/state";
 
 import { KeyDefinition } from "./key-definition";
@@ -95,18 +96,60 @@ export class DeriveDefinition<TFrom, TTo, TDeps extends DerivedStateDependencies
   ) {}
 
   /**
-   * Factory that produces a {@link DeriveDefinition} from a {@link KeyDefinition} and a set of options. The returned
-   * definition will have the same key as the given key definition, but will not collide with it in storage, even if
-   * they both reside in memory.
-   * @param keyDefinition
+   * Factory that produces a {@link DeriveDefinition} from a {@link KeyDefinition} or {@link DeriveDefinition} and new name.
+   *
+   * If a `KeyDefinition` is passed in, the returned definition will have the same key as the given key definition, but
+   * will not collide with it in storage, even if they both reside in memory.
+   *
+   * If a `DeriveDefinition` is passed in, the returned definition will instead use the name given in the second position
+   * of the tuple. It is up to you to ensure this is unique within the domain of derived state.
+   *
+   * @param options A set of options to customize the behavior of {@link DeriveDefinition}.
+   * @param options.derive A function to use to convert values from TFrom to TTo. This is called on each emit of the parent state observable
+   * and the resulting value will be emitted from the derived state observable.
+   * @param options.cleanupDelayMs The number of milliseconds to wait before cleaning up the state after the last subscriber has unsubscribed.
+   * Defaults to 1000ms.
+   * @param options.dependencyShape An object defining the dependencies of the derive function. The keys of the object are the names of the dependencies
+   * and the values are the types of the dependencies.
+   * for example:
+   * ```
+   * {
+   *   myService: MyService,
+   *   myOtherService: MyOtherService,
+   * }
+   * ```
+   *
+   * @param options.deserializer A function to use to safely convert your type from json to your expected type.
+   *   Your data may be serialized/deserialized at any time and this needs callback needs to be able to faithfully re-initialize
+   *   from the JSON object representation of your type.
+   * @param definition
    * @param options
    * @returns
    */
   static from<TFrom, TTo, TDeps extends DerivedStateDependencies = never>(
-    keyDefinition: KeyDefinition<TFrom>,
+    definition:
+      | KeyDefinition<TFrom>
+      | [DeriveDefinition<unknown, TFrom, DerivedStateDependencies>, string],
     options: DeriveDefinitionOptions<TFrom, TTo, TDeps>,
   ) {
-    return new DeriveDefinition(keyDefinition.stateDefinition, keyDefinition.key, options);
+    if (isKeyDefinition(definition)) {
+      return new DeriveDefinition(definition.stateDefinition, definition.key, options);
+    } else {
+      return new DeriveDefinition(definition[0].stateDefinition, definition[1], options);
+    }
+  }
+
+  static fromWithUserId<TKeyDef, TTo, TDeps extends DerivedStateDependencies = never>(
+    definition:
+      | KeyDefinition<TKeyDef>
+      | [DeriveDefinition<unknown, TKeyDef, DerivedStateDependencies>, string],
+    options: DeriveDefinitionOptions<[UserId, TKeyDef], TTo, TDeps>,
+  ) {
+    if (isKeyDefinition(definition)) {
+      return new DeriveDefinition(definition.stateDefinition, definition.key, options);
+    } else {
+      return new DeriveDefinition(definition[0].stateDefinition, definition[1], options);
+    }
   }
 
   get derive() {
@@ -137,3 +180,11 @@ export class DeriveDefinition<TFrom, TTo, TDeps extends DerivedStateDependencies
     return `derived_${this.stateDefinition.name}_${this.uniqueDerivationName}` as StorageKey;
   }
 }
+
+function isKeyDefinition(
+  definition:
+    | KeyDefinition<unknown>
+    | [DeriveDefinition<unknown, unknown, DerivedStateDependencies>, string],
+): definition is KeyDefinition<unknown> {
+  return Object.prototype.hasOwnProperty.call(definition, "key");
+}

libs/common/src/platform/state/state.provider.ts

@@ -49,7 +49,7 @@ export abstract class StateProvider {
   getGlobal: <T>(keyDefinition: KeyDefinition<T>) => GlobalState<T>;
   getDerived: <TFrom, TTo, TDeps extends DerivedStateDependencies>(
     parentState$: Observable<TFrom>,
-    deriveDefinition: DeriveDefinition<unknown, TTo, TDeps>,
+    deriveDefinition: DeriveDefinition<TFrom, TTo, TDeps>,
     dependencies: TDeps,
   ) => DerivedState<TTo>;
 }

libs/common/src/state-migrations/migrate.ts

@@ -14,6 +14,7 @@ import { LastSyncMigrator } from "./migrations/16-move-last-sync-to-state-provid
 import { EnablePasskeysMigrator } from "./migrations/17-move-enable-passkeys-to-state-providers";
 import { AutofillSettingsKeyMigrator } from "./migrations/18-move-autofill-settings-to-state-providers";
 import { RequirePasswordOnStartMigrator } from "./migrations/19-migrate-require-password-on-start";
+import { PrivateKeyMigrator } from "./migrations/20-move-private-key-to-state-providers";
 import { FixPremiumMigrator } from "./migrations/3-fix-premium";
 import { RemoveEverBeenUnlockedMigrator } from "./migrations/4-remove-ever-been-unlocked";
 import { AddKeyTypeToOrgKeysMigrator } from "./migrations/5-add-key-type-to-org-keys";
@@ -24,7 +25,7 @@ import { MoveBrowserSettingsToGlobal } from "./migrations/9-move-browser-setting
 import { MinVersionMigrator } from "./migrations/min-version";
 
 export const MIN_VERSION = 2;
-export const CURRENT_VERSION = 19;
+export const CURRENT_VERSION = 20;
 export type MinVersion = typeof MIN_VERSION;
 
 export function createMigrationBuilder() {
@@ -46,7 +47,8 @@ export function createMigrationBuilder() {
     .with(LastSyncMigrator, 15, 16)
     .with(EnablePasskeysMigrator, 16, 17)
     .with(AutofillSettingsKeyMigrator, 17, 18)
-    .with(RequirePasswordOnStartMigrator, 18, CURRENT_VERSION);
+    .with(RequirePasswordOnStartMigrator, 18, 19)
+    .with(PrivateKeyMigrator, 19, CURRENT_VERSION);
 }
 
 export async function currentVersion(

libs/common/src/state-migrations/migrations/20-move-private-key-to-state-providers.spec.ts

@@ -0,0 +1,127 @@
+import { MockProxy, any } from "jest-mock-extended";
+
+import { MigrationHelper } from "../migration-helper";
+import { mockMigrationHelper } from "../migration-helper.spec";
+
+import { PrivateKeyMigrator } from "./20-move-private-key-to-state-providers";
+
+function exampleJSON() {
+  return {
+    global: {
+      otherStuff: "otherStuff1",
+    },
+    authenticatedAccounts: ["user-1", "user-2", "user-3"],
+    "user-1": {
+      keys: {
+        privateKey: {
+          encrypted: "user-1-encrypted-private-key",
+        },
+        otherStuff: "overStuff2",
+      },
+      otherStuff: "otherStuff3",
+    },
+    "user-2": {
+      keys: {
+        otherStuff: "otherStuff4",
+      },
+      otherStuff: "otherStuff5",
+    },
+  };
+}
+
+function rollbackJSON() {
+  return {
+    "user_user-1_crypto_privateKey": "encrypted-private-key",
+    "user_user-2_crypto_privateKey": null as any,
+    global: {
+      otherStuff: "otherStuff1",
+    },
+    authenticatedAccounts: ["user-1", "user-2", "user-3"],
+    "user-1": {
+      keys: {
+        otherStuff: "overStuff2",
+      },
+      otherStuff: "otherStuff3",
+    },
+    "user-2": {
+      keys: {
+        otherStuff: "otherStuff4",
+      },
+      otherStuff: "otherStuff5",
+    },
+  };
+}
+
+describe("privateKeyMigrator", () => {
+  let helper: MockProxy<MigrationHelper>;
+  let sut: PrivateKeyMigrator;
+  const keyDefinitionLike = {
+    key: "privateKey",
+    stateDefinition: {
+      name: "crypto",
+    },
+  };
+
+  describe("migrate", () => {
+    beforeEach(() => {
+      helper = mockMigrationHelper(exampleJSON(), 19);
+      sut = new PrivateKeyMigrator(19, 20);
+    });
+
+    it("should remove privateKey from all accounts", async () => {
+      await sut.migrate(helper);
+      expect(helper.set).toHaveBeenCalledTimes(1);
+      expect(helper.set).toHaveBeenCalledWith("user-1", {
+        keys: {
+          otherStuff: "overStuff2",
+        },
+        otherStuff: "otherStuff3",
+      });
+    });
+
+    it("should set privateKey value for each account", async () => {
+      await sut.migrate(helper);
+
+      expect(helper.setToUser).toHaveBeenCalledTimes(1);
+      expect(helper.setToUser).toHaveBeenCalledWith(
+        "user-1",
+        keyDefinitionLike,
+        "user-1-encrypted-private-key",
+      );
+    });
+  });
+
+  describe("rollback", () => {
+    beforeEach(() => {
+      helper = mockMigrationHelper(rollbackJSON(), 20);
+      sut = new PrivateKeyMigrator(19, 20);
+    });
+
+    it.each(["user-1", "user-2", "user-3"])("should null out new values %s", async (userId) => {
+      await sut.rollback(helper);
+
+      expect(helper.setToUser).toHaveBeenCalledWith(userId, keyDefinitionLike, null);
+    });
+
+    it("should add explicit value back to accounts", async () => {
+      await sut.rollback(helper);
+
+      expect(helper.set).toHaveBeenCalledTimes(1);
+      expect(helper.set).toHaveBeenCalledWith("user-1", {
+        keys: {
+          privateKey: {
+            encrypted: "encrypted-private-key",
+          },
+          otherStuff: "overStuff2",
+        },
+        otherStuff: "otherStuff3",
+      });
+    });
+
+    it("should not try to restore values to missing accounts", async () => {
+      await sut.rollback(helper);
+
+      expect(helper.set).not.toHaveBeenCalledWith("user-3", any());
+    });
+  });
+});

libs/common/src/state-migrations/migrations/20-move-private-key-to-state-providers.ts

@@ -0,0 +1,53 @@
+import { KeyDefinitionLike, MigrationHelper } from "../migration-helper";
+import { Migrator } from "../migrator";
+
+type ExpectedAccountType = {
+  keys?: {
+    privateKey?: {
+      encrypted?: string; // EncryptedString
+    };
+  };
+};
+
+const USER_ENCRYPTED_PRIVATE_KEY: KeyDefinitionLike = {
+  key: "privateKey",
+  stateDefinition: {
+    name: "crypto",
+  },
+};
+
+export class PrivateKeyMigrator extends Migrator<19, 20> {
+  async migrate(helper: MigrationHelper): Promise<void> {
+    const accounts = await helper.getAccounts<ExpectedAccountType>();
+    async function migrateAccount(userId: string, account: ExpectedAccountType): Promise<void> {
+      const value = account?.keys?.privateKey?.encrypted;
+      if (value != null) {
+        await helper.setToUser(userId, USER_ENCRYPTED_PRIVATE_KEY, value);
+        delete account.keys.privateKey;
+        await helper.set(userId, account);
+      }
+    }
+
+    await Promise.all([...accounts.map(({ userId, account }) => migrateAccount(userId, account))]);
+  }
+  async rollback(helper: MigrationHelper): Promise<void> {
+    const accounts = await helper.getAccounts<ExpectedAccountType>();
+    async function rollbackAccount(userId: string, account: ExpectedAccountType): Promise<void> {
+      const value = await helper.getFromUser<Record<string, string>>(
+        userId,
+        USER_ENCRYPTED_PRIVATE_KEY,
+      );
+      if (account && value) {
+        account.keys = Object.assign(account.keys ?? {}, {
+          privateKey: {
+            encrypted: value,
+          },
+        });
+        await helper.set(userId, account);
+      }
+      await helper.setToUser(userId, USER_ENCRYPTED_PRIVATE_KEY, null);
+    }
+
+    await Promise.all([...accounts.map(({ userId, account }) => rollbackAccount(userId, account))]);
+  }
+}

libs/common/src/types/key.ts

@@ -14,3 +14,4 @@ export type CipherKey = Opaque<SymmetricCryptoKey, "CipherKey">;
 
 // asymmetric keys
 export type UserPrivateKey = Opaque<Uint8Array, "UserPrivateKey">;
+export type UserPublicKey = Opaque<Uint8Array, "UserPublicKey">;