[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/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))]);
+  }
+}