[PM-5533] Migrate Org Keys to state providers (#7521)

* Move org keys to state providers

* Create state for org keys and derive decrypted for use

* Make state readonly

* Remove org keys from state service

* Migrate user keys state

* Review feedback

* Correct test name

* Refix key types

* `npm run prettier` 🤖

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

@@ -4,6 +4,7 @@ import { ProfileOrganizationResponse } from "../../admin-console/models/response
 import { ProfileProviderOrganizationResponse } from "../../admin-console/models/response/profile-provider-organization.response";
 import { ProfileProviderResponse } from "../../admin-console/models/response/profile-provider.response";
 import { KdfConfig } from "../../auth/models/domain/kdf-config";
+import { OrganizationId } from "../../types/guid";
 import { UserKey, MasterKey, OrgKey, ProviderKey, PinKey, CipherKey } from "../../types/key";
 import { KeySuffixOptions, KdfType, HashPurpose } from "../enums";
 import { EncArrayBuffer } from "../models/domain/enc-array-buffer";
@@ -199,16 +200,19 @@ export abstract class CryptoService {
     orgs: ProfileOrganizationResponse[],
     providerOrgs: ProfileProviderOrganizationResponse[],
   ) => Promise<void>;
+  activeUserOrgKeys$: Observable<Record<OrganizationId, OrgKey>>;
   /**
    * Returns the organization's symmetric key
+   * @deprecated Use the observable activeUserOrgKeys$ and `map` to the desired orgKey instead
    * @param orgId The desired organization
    * @returns The organization's symmetric key
    */
   getOrgKey: (orgId: string) => Promise<OrgKey>;
   /**
-   * @returns A map of the organization Ids to their symmetric keys
+   * @deprecated Use the observable activeUserOrgKeys$ instead
+   * @returns A record of the organization Ids to their symmetric keys
    */
-  getOrgKeys: () => Promise<Map<string, SymmetricCryptoKey>>;
+  getOrgKeys: () => Promise<Record<string, SymmetricCryptoKey>>;
   /**
    * Uses the org key to derive a new symmetric key for encrypting data
    * @param orgKey The organization's symmetric key

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

@@ -1,6 +1,5 @@
 import { Observable } from "rxjs";
 
-import { EncryptedOrganizationKeyData } from "../../admin-console/models/data/encrypted-organization-key.data";
 import { OrganizationData } from "../../admin-console/models/data/organization.data";
 import { PolicyData } from "../../admin-console/models/data/policy.data";
 import { ProviderData } from "../../admin-console/models/data/provider.data";
@@ -193,13 +192,6 @@ export abstract class StateService<T extends Account = Account> {
   setDecryptedCiphers: (value: CipherView[], options?: StorageOptions) => Promise<void>;
   getDecryptedCollections: (options?: StorageOptions) => Promise<CollectionView[]>;
   setDecryptedCollections: (value: CollectionView[], options?: StorageOptions) => Promise<void>;
-  getDecryptedOrganizationKeys: (
-    options?: StorageOptions,
-  ) => Promise<Map<string, SymmetricCryptoKey>>;
-  setDecryptedOrganizationKeys: (
-    value: Map<string, SymmetricCryptoKey>,
-    options?: StorageOptions,
-  ) => Promise<void>;
   getDecryptedPasswordGenerationHistory: (
     options?: StorageOptions,
   ) => Promise<GeneratedPasswordHistory[]>;
@@ -344,13 +336,6 @@ export abstract class StateService<T extends Account = Account> {
     value: { [id: string]: FolderData },
     options?: StorageOptions,
   ) => Promise<void>;
-  getEncryptedOrganizationKeys: (
-    options?: StorageOptions,
-  ) => Promise<{ [orgId: string]: EncryptedOrganizationKeyData }>;
-  setEncryptedOrganizationKeys: (
-    value: { [orgId: string]: EncryptedOrganizationKeyData },
-    options?: StorageOptions,
-  ) => Promise<void>;
   getEncryptedPasswordGenerationHistory: (
     options?: StorageOptions,
   ) => Promise<GeneratedPasswordHistory[]>;

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

@@ -1,6 +1,5 @@
 import { Jsonify } from "type-fest";
 
-import { EncryptedOrganizationKeyData } from "../../../admin-console/models/data/encrypted-organization-key.data";
 import { OrganizationData } from "../../../admin-console/models/data/organization.data";
 import { PolicyData } from "../../../admin-console/models/data/policy.data";
 import { ProviderData } from "../../../admin-console/models/data/provider.data";
@@ -127,13 +126,6 @@ export class AccountKeys {
   masterKey?: MasterKey;
   masterKeyEncryptedUserKey?: string;
   deviceKey?: ReturnType<SymmetricCryptoKey["toJSON"]>;
-  organizationKeys?: EncryptionPair<
-    { [orgId: string]: EncryptedOrganizationKeyData },
-    Record<string, SymmetricCryptoKey>
-  > = new EncryptionPair<
-    { [orgId: string]: EncryptedOrganizationKeyData },
-    Record<string, SymmetricCryptoKey>
-  >();
   providerKeys?: EncryptionPair<any, Record<string, SymmetricCryptoKey>> = new EncryptionPair<
     any,
     Record<string, SymmetricCryptoKey>
@@ -176,7 +168,6 @@ export class AccountKeys {
         obj?.cryptoSymmetricKey,
         SymmetricCryptoKey.fromJSON,
       ),
-      organizationKeys: AccountKeys.initRecordEncryptionPairsFromJSON(obj?.organizationKeys),
       providerKeys: AccountKeys.initRecordEncryptionPairsFromJSON(obj?.providerKeys),
       privateKey: EncryptionPair.fromJSON<string, Uint8Array>(obj?.privateKey, (decObj: string) =>
         Utils.fromByteStringToArray(decObj),

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

@@ -15,7 +15,9 @@ import { StateService } from "../abstractions/state.service";
 import { Utils } from "../misc/utils";
 import { EncString } from "../models/domain/enc-string";
 import { SymmetricCryptoKey } from "../models/domain/symmetric-crypto-key";
-import { CryptoService, USER_EVER_HAD_USER_KEY } from "../services/crypto.service";
+import { CryptoService } from "../services/crypto.service";
+
+import { USER_EVER_HAD_USER_KEY } from "./key-state/user-key.state";
 
 describe("cryptoService", () => {
   let cryptoService: CryptoService;

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

@@ -1,15 +1,14 @@
 import * as bigInt from "big-integer";
-import { firstValueFrom, map } from "rxjs";
+import { Observable, firstValueFrom, map } from "rxjs";
 
 import { EncryptedOrganizationKeyData } from "../../admin-console/models/data/encrypted-organization-key.data";
-import { BaseEncryptedOrganizationKey } from "../../admin-console/models/domain/encrypted-organization-key";
 import { ProfileOrganizationResponse } from "../../admin-console/models/response/profile-organization.response";
 import { ProfileProviderOrganizationResponse } from "../../admin-console/models/response/profile-provider-organization.response";
 import { ProfileProviderResponse } from "../../admin-console/models/response/profile-provider.response";
 import { AccountService } from "../../auth/abstractions/account.service";
 import { KdfConfig } from "../../auth/models/domain/kdf-config";
 import { Utils } from "../../platform/misc/utils";
-import { UserId } from "../../types/guid";
+import { OrganizationId, UserId } from "../../types/guid";
 import { UserKey, MasterKey, OrgKey, ProviderKey, PinKey, CipherKey } from "../../types/key";
 import { CryptoFunctionService } from "../abstractions/crypto-function.service";
 import { CryptoService as CryptoServiceAbstraction } from "../abstractions/crypto.service";
@@ -32,14 +31,22 @@ import { EFFLongWordList } from "../misc/wordlist";
 import { EncArrayBuffer } from "../models/domain/enc-array-buffer";
 import { EncString } from "../models/domain/enc-string";
 import { SymmetricCryptoKey } from "../models/domain/symmetric-crypto-key";
-import { ActiveUserState, CRYPTO_DISK, KeyDefinition, StateProvider } from "../state";
+import { ActiveUserState, DerivedState, StateProvider } from "../state";
 
-export const USER_EVER_HAD_USER_KEY = new KeyDefinition<boolean>(CRYPTO_DISK, "everHadUserKey", {
-  deserializer: (obj) => obj,
-});
+import {
+  USER_ENCRYPTED_ORGANIZATION_KEYS,
+  USER_ORGANIZATION_KEYS,
+} from "./key-state/org-keys.state";
+import { USER_EVER_HAD_USER_KEY } from "./key-state/user-key.state";
 
 export class CryptoService implements CryptoServiceAbstraction {
-  private activeUserEverHadUserKey: ActiveUserState<boolean>;
+  private readonly activeUserEverHadUserKey: ActiveUserState<boolean>;
+  private readonly activeUserEncryptedOrgKeysState: ActiveUserState<
+    Record<OrganizationId, EncryptedOrganizationKeyData>
+  >;
+  private readonly activeUserOrgKeysState: DerivedState<Record<OrganizationId, OrgKey>>;
+
+  readonly activeUserOrgKeys$: Observable<Record<OrganizationId, OrgKey>>;
 
   readonly everHadUserKey$;
 
@@ -53,8 +60,17 @@ export class CryptoService implements CryptoServiceAbstraction {
     protected stateProvider: StateProvider,
   ) {
     this.activeUserEverHadUserKey = stateProvider.getActive(USER_EVER_HAD_USER_KEY);
+    this.activeUserEncryptedOrgKeysState = stateProvider.getActive(
+      USER_ENCRYPTED_ORGANIZATION_KEYS,
+    );
+    this.activeUserOrgKeysState = stateProvider.getDerived(
+      this.activeUserEncryptedOrgKeysState.state$,
+      USER_ORGANIZATION_KEYS,
+      { cryptoService: this },
+    );
 
     this.everHadUserKey$ = this.activeUserEverHadUserKey.state$.pipe(map((x) => x ?? false));
+    this.activeUserOrgKeys$ = this.activeUserOrgKeysState.state$; // null handled by `derive` function
   }
 
   async setUserKey(key: UserKey, userId?: UserId): Promise<void> {
@@ -320,72 +336,35 @@ export class CryptoService implements CryptoServiceAbstraction {
     orgs: ProfileOrganizationResponse[] = [],
     providerOrgs: ProfileProviderOrganizationResponse[] = [],
   ): Promise<void> {
-    const encOrgKeyData: { [orgId: string]: EncryptedOrganizationKeyData } = {};
+    this.activeUserEncryptedOrgKeysState.update((_) => {
+      const encOrgKeyData: { [orgId: string]: EncryptedOrganizationKeyData } = {};
 
-    orgs.forEach((org) => {
-      encOrgKeyData[org.id] = {
-        type: "organization",
-        key: org.key,
-      };
+      orgs.forEach((org) => {
+        encOrgKeyData[org.id] = {
+          type: "organization",
+          key: org.key,
+        };
+      });
+
+      providerOrgs.forEach((org) => {
+        encOrgKeyData[org.id] = {
+          type: "provider",
+          providerId: org.providerId,
+          key: org.key,
+        };
+      });
+
+      return encOrgKeyData;
     });
-
-    providerOrgs.forEach((org) => {
-      encOrgKeyData[org.id] = {
-        type: "provider",
-        providerId: org.providerId,
-        key: org.key,
-      };
-    });
-
-    await this.stateService.setDecryptedOrganizationKeys(null);
-    return await this.stateService.setEncryptedOrganizationKeys(encOrgKeyData);
   }
 
-  async getOrgKey(orgId: string): Promise<OrgKey> {
-    if (orgId == null) {
-      return null;
-    }
-
-    const orgKeys = await this.getOrgKeys();
-    if (orgKeys == null || !orgKeys.has(orgId)) {
-      return null;
-    }
-
-    return orgKeys.get(orgId);
+  async getOrgKey(orgId: OrganizationId): Promise<OrgKey> {
+    return (await firstValueFrom(this.activeUserOrgKeys$))[orgId];
   }
 
   @sequentialize(() => "getOrgKeys")
-  async getOrgKeys(): Promise<Map<string, OrgKey>> {
-    const result: Map<string, OrgKey> = new Map<string, OrgKey>();
-    const decryptedOrganizationKeys = await this.stateService.getDecryptedOrganizationKeys();
-    if (decryptedOrganizationKeys != null && decryptedOrganizationKeys.size > 0) {
-      return decryptedOrganizationKeys as Map<string, OrgKey>;
-    }
-
-    const encOrgKeyData = await this.stateService.getEncryptedOrganizationKeys();
-    if (encOrgKeyData == null) {
-      return result;
-    }
-
-    let setKey = false;
-
-    for (const orgId of Object.keys(encOrgKeyData)) {
-      if (result.has(orgId)) {
-        continue;
-      }
-
-      const encOrgKey = BaseEncryptedOrganizationKey.fromData(encOrgKeyData[orgId]);
-      const decOrgKey = (await encOrgKey.decrypt(this)) as OrgKey;
-      result.set(orgId, decOrgKey);
-
-      setKey = true;
-    }
-
-    if (setKey) {
-      await this.stateService.setDecryptedOrganizationKeys(result);
-    }
-
-    return result;
+  async getOrgKeys(): Promise<Record<string, OrgKey>> {
+    return await firstValueFrom(this.activeUserOrgKeys$);
   }
 
   async makeDataEncKey<T extends OrgKey | UserKey>(
@@ -400,9 +379,19 @@ export class CryptoService implements CryptoServiceAbstraction {
   }
 
   async clearOrgKeys(memoryOnly?: boolean, userId?: UserId): Promise<void> {
-    await this.stateService.setDecryptedOrganizationKeys(null, { userId: userId });
-    if (!memoryOnly) {
-      await this.stateService.setEncryptedOrganizationKeys(null, { userId: userId });
+    const activeUserId = (await firstValueFrom(this.accountService.activeAccount$))?.id;
+    const userIdIsActive = userId == null || userId === activeUserId;
+    if (memoryOnly && userIdIsActive) {
+      // org keys are only cached for active users
+      await this.activeUserOrgKeysState.forceValue({});
+    } else {
+      if (userId == null && activeUserId == null) {
+        // nothing to do
+        return;
+      }
+      await this.stateProvider
+        .getUser(userId ?? activeUserId, USER_ENCRYPTED_ORGANIZATION_KEYS)
+        .update(() => null);
     }
   }
 

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

@@ -0,0 +1,114 @@
+import { mock } from "jest-mock-extended";
+
+import { makeStaticByteArray } from "../../../../spec";
+import { ProviderEncryptedOrganizationKey } from "../../../admin-console/models/domain/encrypted-organization-key";
+import { OrgKey } from "../../../types/key";
+import { CryptoService } from "../../abstractions/crypto.service";
+import { EncryptionType } from "../../enums";
+import { Utils } from "../../misc/utils";
+import { EncString } from "../../models/domain/enc-string";
+import { SymmetricCryptoKey } from "../../models/domain/symmetric-crypto-key";
+
+import { USER_ENCRYPTED_ORGANIZATION_KEYS, USER_ORGANIZATION_KEYS } from "./org-keys.state";
+
+function makeEncString(data?: string) {
+  data ??= Utils.newGuid();
+  return new EncString(EncryptionType.AesCbc256_HmacSha256_B64, data, "test", "test");
+}
+ProviderEncryptedOrganizationKey;
+
+describe("encrypted org keys", () => {
+  const sut = USER_ENCRYPTED_ORGANIZATION_KEYS;
+
+  it("should deserialize encrypted org keys", () => {
+    const encryptedOrgKeys = {
+      "org-id-1": {
+        type: "organization",
+        key: makeEncString().encryptedString,
+      },
+      "org-id-2": {
+        type: "provider",
+        key: makeEncString().encryptedString,
+        providerId: "provider-id-2",
+      },
+    };
+
+    const result = sut.deserializer(JSON.parse(JSON.stringify(encryptedOrgKeys)));
+
+    expect(result).toEqual(encryptedOrgKeys);
+  });
+});
+
+describe("derived decrypted org keys", () => {
+  const cryptoService = mock<CryptoService>();
+  const sut = USER_ORGANIZATION_KEYS;
+
+  afterEach(() => {
+    jest.resetAllMocks();
+  });
+
+  it("should deserialize org keys", async () => {
+    const decryptedOrgKeys = {
+      "org-id-1": new SymmetricCryptoKey(makeStaticByteArray(64, 1)) as OrgKey,
+      "org-id-2": new SymmetricCryptoKey(makeStaticByteArray(64, 2)) as OrgKey,
+    };
+
+    const result = sut.deserialize(JSON.parse(JSON.stringify(decryptedOrgKeys)));
+
+    expect(result).toEqual(decryptedOrgKeys);
+  });
+
+  it("should derive org keys", async () => {
+    const encryptedOrgKeys = {
+      "org-id-1": {
+        type: "organization",
+        key: makeEncString().encryptedString,
+      },
+      "org-id-2": {
+        type: "organization",
+        key: makeEncString().encryptedString,
+      },
+    };
+
+    const decryptedOrgKeys = {
+      "org-id-1": new SymmetricCryptoKey(makeStaticByteArray(64, 1)) as OrgKey,
+      "org-id-2": new SymmetricCryptoKey(makeStaticByteArray(64, 2)) as OrgKey,
+    };
+
+    // TODO: How to not have to mock these decryptions. They are internal concerns of EncryptedOrganizationKey
+    cryptoService.rsaDecrypt.mockResolvedValueOnce(decryptedOrgKeys["org-id-1"].key);
+    cryptoService.rsaDecrypt.mockResolvedValueOnce(decryptedOrgKeys["org-id-2"].key);
+
+    const result = await sut.derive(encryptedOrgKeys, { cryptoService });
+
+    expect(result).toEqual(decryptedOrgKeys);
+  });
+
+  it("should derive org keys from providers", async () => {
+    const encryptedOrgKeys = {
+      "org-id-1": {
+        type: "provider",
+        key: makeEncString().encryptedString,
+        providerId: "provider-id-1",
+      },
+      "org-id-2": {
+        type: "provider",
+        key: makeEncString().encryptedString,
+        providerId: "provider-id-2",
+      },
+    };
+
+    const decryptedOrgKeys = {
+      "org-id-1": new SymmetricCryptoKey(makeStaticByteArray(64, 1)) as OrgKey,
+      "org-id-2": new SymmetricCryptoKey(makeStaticByteArray(64, 2)) as OrgKey,
+    };
+
+    // TODO: How to not have to mock these decryptions. They are internal concerns of ProviderEncryptedOrganizationKey
+    cryptoService.decryptToBytes.mockResolvedValueOnce(decryptedOrgKeys["org-id-1"].key);
+    cryptoService.decryptToBytes.mockResolvedValueOnce(decryptedOrgKeys["org-id-2"].key);
+
+    const result = await sut.derive(encryptedOrgKeys, { cryptoService });
+
+    expect(result).toEqual(decryptedOrgKeys);
+  });
+});

libs/common/src/platform/services/key-state/org-keys.state.ts

@@ -0,0 +1,42 @@
+import { EncryptedOrganizationKeyData } from "../../../admin-console/models/data/encrypted-organization-key.data";
+import { BaseEncryptedOrganizationKey } from "../../../admin-console/models/domain/encrypted-organization-key";
+import { OrganizationId } from "../../../types/guid";
+import { OrgKey } from "../../../types/key";
+import { CryptoService } from "../../abstractions/crypto.service";
+import { SymmetricCryptoKey } from "../../models/domain/symmetric-crypto-key";
+import { KeyDefinition, CRYPTO_DISK, DeriveDefinition } from "../../state";
+
+export const USER_ENCRYPTED_ORGANIZATION_KEYS = KeyDefinition.record<
+  EncryptedOrganizationKeyData,
+  OrganizationId
+>(CRYPTO_DISK, "organizationKeys", {
+  deserializer: (obj) => obj,
+});
+
+export const USER_ORGANIZATION_KEYS = DeriveDefinition.from<
+  Record<OrganizationId, EncryptedOrganizationKeyData>,
+  Record<OrganizationId, OrgKey>,
+  { cryptoService: CryptoService }
+>(USER_ENCRYPTED_ORGANIZATION_KEYS, {
+  deserializer: (obj) => {
+    const result: Record<OrganizationId, OrgKey> = {};
+    for (const orgId of Object.keys(obj ?? {}) as OrganizationId[]) {
+      result[orgId] = SymmetricCryptoKey.fromJSON(obj[orgId]) as OrgKey;
+    }
+    return result;
+  },
+  derive: async (from, { cryptoService }) => {
+    const result: Record<OrganizationId, OrgKey> = {};
+    for (const orgId of Object.keys(from ?? {}) as OrganizationId[]) {
+      if (result[orgId] != null) {
+        continue;
+      }
+      const encrypted = BaseEncryptedOrganizationKey.fromData(from[orgId]);
+      const decrypted = await encrypted.decrypt(cryptoService);
+
+      result[orgId] = decrypted;
+    }
+
+    return result;
+  },
+});

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

@@ -0,0 +1,5 @@
+import { KeyDefinition, CRYPTO_DISK } from "../../state";
+
+export const USER_EVER_HAD_USER_KEY = new KeyDefinition<boolean>(CRYPTO_DISK, "everHadUserKey", {
+  deserializer: (obj) => obj,
+});

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

@@ -2,7 +2,6 @@ import { BehaviorSubject, concatMap } from "rxjs";
 import { Jsonify, JsonValue } from "type-fest";
 
 import { AutofillOverlayVisibility } from "../../../../../apps/browser/src/autofill/utils/autofill-overlay.enum";
-import { EncryptedOrganizationKeyData } from "../../admin-console/models/data/encrypted-organization-key.data";
 import { OrganizationData } from "../../admin-console/models/data/organization.data";
 import { PolicyData } from "../../admin-console/models/data/policy.data";
 import { ProviderData } from "../../admin-console/models/data/provider.data";
@@ -991,29 +990,6 @@ export class StateService<
     );
   }
 
-  async getDecryptedOrganizationKeys(
-    options?: StorageOptions,
-  ): Promise<Map<string, SymmetricCryptoKey>> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-    return Utils.recordToMap(account?.keys?.organizationKeys?.decrypted);
-  }
-
-  async setDecryptedOrganizationKeys(
-    value: Map<string, SymmetricCryptoKey>,
-    options?: StorageOptions,
-  ): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-    account.keys.organizationKeys.decrypted = Utils.mapToRecord(value);
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-  }
-
   @withPrototypeForArrayMembers(GeneratedPasswordHistory)
   async getDecryptedPasswordGenerationHistory(
     options?: StorageOptions,
@@ -1856,28 +1832,6 @@ export class StateService<
     );
   }
 
-  async getEncryptedOrganizationKeys(
-    options?: StorageOptions,
-  ): Promise<{ [orgId: string]: EncryptedOrganizationKeyData }> {
-    return (
-      await this.getAccount(this.reconcileOptions(options, await this.defaultOnDiskOptions()))
-    )?.keys?.organizationKeys.encrypted;
-  }
-
-  async setEncryptedOrganizationKeys(
-    value: { [orgId: string]: EncryptedOrganizationKeyData },
-    options?: StorageOptions,
-  ): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultOnDiskOptions()),
-    );
-    account.keys.organizationKeys.encrypted = value;
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultOnDiskOptions()),
-    );
-  }
-
   @withPrototypeForArrayMembers(GeneratedPasswordHistory)
   async getEncryptedPasswordGenerationHistory(
     options?: StorageOptions,