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

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;