[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/admin-console/models/domain/encrypted-organization-key.ts

@@ -1,6 +1,7 @@
 import { CryptoService } from "../../../platform/abstractions/crypto.service";
 import { EncString } from "../../../platform/models/domain/enc-string";
 import { SymmetricCryptoKey } from "../../../platform/models/domain/symmetric-crypto-key";
+import { OrgKey } from "../../../types/key";
 import { EncryptedOrganizationKeyData } from "../data/encrypted-organization-key.data";
 
 export abstract class BaseEncryptedOrganizationKey {
@@ -25,7 +26,7 @@ export class EncryptedOrganizationKey implements BaseEncryptedOrganizationKey {
 
   async decrypt(cryptoService: CryptoService) {
     const decValue = await cryptoService.rsaDecrypt(this.key);
-    return new SymmetricCryptoKey(decValue);
+    return new SymmetricCryptoKey(decValue) as OrgKey;
   }
 
   toData(): EncryptedOrganizationKeyData {
@@ -45,7 +46,7 @@ export class ProviderEncryptedOrganizationKey implements BaseEncryptedOrganizati
   async decrypt(cryptoService: CryptoService) {
     const providerKey = await cryptoService.getProviderKey(this.providerId);
     const decValue = await cryptoService.decryptToBytes(new EncString(this.key), providerKey);
-    return new SymmetricCryptoKey(decValue);
+    return new SymmetricCryptoKey(decValue) as OrgKey;
   }
 
   toData(): EncryptedOrganizationKeyData {

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,

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

@@ -6,6 +6,7 @@ import { AbstractStorageService } from "../platform/abstractions/storage.service
 import { MigrationBuilder } from "./migration-builder";
 import { MigrationHelper } from "./migration-helper";
 import { EverHadUserKeyMigrator } from "./migrations/10-move-ever-had-user-key-to-state-providers";
+import { OrganizationKeyMigrator } from "./migrations/11-move-org-keys-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";
@@ -16,7 +17,7 @@ import { MoveBrowserSettingsToGlobal } from "./migrations/9-move-browser-setting
 import { MinVersionMigrator } from "./migrations/min-version";
 
 export const MIN_VERSION = 2;
-export const CURRENT_VERSION = 10;
+export const CURRENT_VERSION = 11;
 export type MinVersion = typeof MIN_VERSION;
 
 export async function migrate(
@@ -42,7 +43,9 @@ export async function migrate(
     .with(MoveBiometricAutoPromptToAccount, 6, 7)
     .with(MoveStateVersionMigrator, 7, 8)
     .with(MoveBrowserSettingsToGlobal, 8, 9)
-    .with(EverHadUserKeyMigrator, 9, CURRENT_VERSION)
+    .with(EverHadUserKeyMigrator, 9, 10)
+    .with(OrganizationKeyMigrator, 10, CURRENT_VERSION)
+
     .migrate(migrationHelper);
 }
 

libs/common/src/state-migrations/migrations/11-move-org-keys-to-state-providers.spec.ts

@@ -0,0 +1,163 @@
+import { MockProxy, any } from "jest-mock-extended";
+
+import { MigrationHelper } from "../migration-helper";
+import { mockMigrationHelper } from "../migration-helper.spec";
+
+import { OrganizationKeyMigrator } from "./11-move-org-keys-to-state-providers";
+
+function exampleJSON() {
+  return {
+    global: {
+      otherStuff: "otherStuff1",
+    },
+    authenticatedAccounts: ["user-1", "user-2", "user-3"],
+    "user-1": {
+      keys: {
+        organizationKeys: {
+          encrypted: {
+            "org-id-1": {
+              type: "organization",
+              key: "org-key-1",
+            },
+            "org-id-2": {
+              type: "provider",
+              key: "org-key-2",
+              providerId: "provider-id-2",
+            },
+          },
+        },
+        otherStuff: "overStuff2",
+      },
+      otherStuff: "otherStuff3",
+    },
+    "user-2": {
+      keys: {
+        otherStuff: "otherStuff4",
+      },
+      otherStuff: "otherStuff5",
+    },
+  };
+}
+
+function rollbackJSON() {
+  return {
+    "user_user-1_crypto_organizationKeys": {
+      "org-id-1": {
+        type: "organization",
+        key: "org-key-1",
+      },
+      "org-id-2": {
+        type: "provider",
+        key: "org-key-2",
+        providerId: "provider-id-2",
+      },
+    },
+    "user_user-2_crypto_organizationKeys": 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("OrganizationKeysMigrator", () => {
+  let helper: MockProxy<MigrationHelper>;
+  let sut: OrganizationKeyMigrator;
+  const keyDefinitionLike = {
+    key: "organizationKeys",
+    stateDefinition: {
+      name: "crypto",
+    },
+  };
+
+  describe("migrate", () => {
+    beforeEach(() => {
+      helper = mockMigrationHelper(exampleJSON(), 10);
+      sut = new OrganizationKeyMigrator(10, 11);
+    });
+
+    it("should remove organizationKeys 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 organizationKeys value for each account", async () => {
+      await sut.migrate(helper);
+
+      expect(helper.setToUser).toHaveBeenCalledTimes(1);
+      expect(helper.setToUser).toHaveBeenCalledWith("user-1", keyDefinitionLike, {
+        "org-id-1": {
+          type: "organization",
+          key: "org-key-1",
+        },
+        "org-id-2": {
+          type: "provider",
+          key: "org-key-2",
+          providerId: "provider-id-2",
+        },
+      });
+    });
+  });
+
+  describe("rollback", () => {
+    beforeEach(() => {
+      helper = mockMigrationHelper(rollbackJSON(), 10);
+      sut = new OrganizationKeyMigrator(10, 11);
+    });
+
+    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: {
+          organizationKeys: {
+            encrypted: {
+              "org-id-1": {
+                type: "organization",
+                key: "org-key-1",
+              },
+              "org-id-2": {
+                type: "provider",
+                key: "org-key-2",
+                providerId: "provider-id-2",
+              },
+            },
+          },
+          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/11-move-org-keys-to-state-providers.ts

@@ -0,0 +1,59 @@
+import { KeyDefinitionLike, MigrationHelper } from "../migration-helper";
+import { Migrator } from "../migrator";
+
+type OrgKeyDataType = {
+  type: "organization" | "provider";
+  key: string;
+  providerId?: string;
+};
+
+type ExpectedAccountType = {
+  keys?: {
+    organizationKeys?: {
+      encrypted?: Record<string, OrgKeyDataType>;
+    };
+  };
+};
+
+const USER_ENCRYPTED_ORGANIZATION_KEYS: KeyDefinitionLike = {
+  key: "organizationKeys",
+  stateDefinition: {
+    name: "crypto",
+  },
+};
+
+export class OrganizationKeyMigrator extends Migrator<10, 11> {
+  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?.organizationKeys?.encrypted;
+      if (value != null) {
+        await helper.setToUser(userId, USER_ENCRYPTED_ORGANIZATION_KEYS, value);
+        delete account.keys.organizationKeys;
+        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, OrgKeyDataType>>(
+        userId,
+        USER_ENCRYPTED_ORGANIZATION_KEYS,
+      );
+      if (account && value) {
+        account.keys = Object.assign(account.keys ?? {}, {
+          organizationKeys: {
+            encrypted: value,
+          },
+        });
+        await helper.set(userId, account);
+      }
+      await helper.setToUser(userId, USER_ENCRYPTED_ORGANIZATION_KEYS, null);
+    }
+
+    await Promise.all([...accounts.map(({ userId, account }) => rollbackAccount(userId, account))]);
+  }
+}

libs/common/src/types/guid.ts

@@ -3,3 +3,4 @@ import { Opaque } from "type-fest";
 export type Guid = Opaque<string, "Guid">;
 
 export type UserId = Opaque<string, "UserId">;
+export type OrganizationId = Opaque<string, "OrganizationId">;

libs/common/src/vault/services/cipher.service.ts

@@ -290,7 +290,7 @@ export class CipherService implements CipherServiceAbstraction {
     const ciphers = await this.getAll();
     const orgKeys = await this.cryptoService.getOrgKeys();
     const userKey = await this.cryptoService.getUserKeyWithLegacySupport();
-    if (orgKeys?.size === 0 && userKey == null) {
+    if (Object.keys(orgKeys).length === 0 && userKey == null) {
       // return early if there are no keys to decrypt with
       return;
     }
@@ -308,7 +308,7 @@ export class CipherService implements CipherServiceAbstraction {
     const decCiphers = (
       await Promise.all(
         Object.entries(grouped).map(([orgId, groupedCiphers]) =>
-          this.encryptService.decryptItems(groupedCiphers, orgKeys.get(orgId) ?? userKey),
+          this.encryptService.decryptItems(groupedCiphers, orgKeys[orgId] ?? userKey),
         ),
       )
     )