From f6e8cffd9476bcf246138452f70b529f89c3ee0f Mon Sep 17 00:00:00 2001
From: Thomas Rittson <trittson@bitwarden.com>
Date: Thu, 12 May 2022 12:15:20 +1000
Subject: [PATCH] Finish migration and add tests

---
 ...vice.ts => stateMigration.service.spec.ts} | 53 ++++++++++++++++++-
 common/src/enums/stateVersion.ts              |  2 +-
 common/src/services/stateMigration.service.ts | 14 +++--
 3 files changed, 64 insertions(+), 5 deletions(-)
 rename common/spec/services/{stateMigration.service.ts => stateMigration.service.spec.ts} (62%)

diff --git a/common/spec/services/stateMigration.service.ts b/common/spec/services/stateMigration.service.spec.ts
similarity index 62%
rename from common/spec/services/stateMigration.service.ts
rename to common/spec/services/stateMigration.service.spec.ts
index d325585a..3d0fc94b 100644
--- a/common/spec/services/stateMigration.service.ts
+++ b/common/spec/services/stateMigration.service.spec.ts
@@ -28,7 +28,7 @@ describe("State Migration Service", () => {
     );
   });
 
-  describe("StateVersion 3 to 4 migration", async () => {
+  describe("StateVersion 3 to 4 migration", () => {
     beforeEach(() => {
       const globalVersion3: Partial<GlobalState> = {
         stateVersion: StateVersion.Three,
@@ -81,4 +81,55 @@ describe("State Migration Service", () => {
       );
     });
   });
+
+  describe("StateVersion 4 to 5 migration", () => {
+    beforeEach(() => {
+      const globalVersion4: Partial<GlobalState> = {
+        stateVersion: StateVersion.Four,
+      };
+
+      storageService.get("global", Arg.any()).resolves(globalVersion4);
+      storageService.get("authenticatedAccounts", Arg.any()).resolves([userId]);
+    });
+
+    it("migrates organization keys to new format", async () => {
+      const accountVersion4 = new Account({
+        keys: {
+          organizationKeys: {
+            encrypted: {
+              orgOneId: "orgOneEncKey",
+              orgTwoId: "orgTwoEncKey",
+              orgThreeId: "orgThreeEncKey",
+            },
+          },
+        },
+      } as any);
+
+      storageService.get(userId, Arg.any()).resolves(accountVersion4);
+
+      await stateMigrationService.migrate();
+
+      storageService.received(1).save(
+        userId,
+        Arg.is((account: Account) => {
+          return (
+            account.keys.organizationKeys.encrypted["orgOneId"].key === "orgOneEncKey" &&
+            account.keys.organizationKeys.encrypted["orgTwoId"].key === "orgTwoEncKey" &&
+            account.keys.organizationKeys.encrypted["orgThreeId"].key === "orgThreeEncKey"
+          );
+        }),
+        Arg.any()
+      );
+    });
+
+    it("updates StateVersion number", async () => {
+      await stateMigrationService.migrate();
+
+      storageService.received(1).save(
+        "global",
+        Arg.is((globals: GlobalState) => globals.stateVersion === StateVersion.Five),
+        Arg.any()
+      );
+    });
+  });
 });
diff --git a/common/src/enums/stateVersion.ts b/common/src/enums/stateVersion.ts
index 285d973e..cc36f53f 100644
--- a/common/src/enums/stateVersion.ts
+++ b/common/src/enums/stateVersion.ts
@@ -3,6 +3,6 @@ export enum StateVersion {
   Two = 2, // Move to a typed State object
   Three = 3, // Fix migration of users' premium status
   Four = 4, // Fix 'Never Lock' option by removing stale data
-  Five = 5,
+  Five = 5, // Migrate to new storage of encrypted organization keys
   Latest = Five,
 }
diff --git a/common/src/services/stateMigration.service.ts b/common/src/services/stateMigration.service.ts
index ec834c6c..93b5ce24 100644
--- a/common/src/services/stateMigration.service.ts
+++ b/common/src/services/stateMigration.service.ts
@@ -1,3 +1,5 @@
+import { EncryptedOrganizationKeyData } from "jslib-common/models/data/encryptedOrganizationKeyData";
+
 import { StorageService } from "../abstractions/storage.service";
 import { HtmlStorageLocation } from "../enums/htmlStorageLocation";
 import { KdfType } from "../enums/kdfType";
@@ -493,12 +495,18 @@ export class StateMigrationService<
 
   protected async migrateStateFrom4To5(): Promise<void> {
     const authenticatedUserIds = await this.get<string[]>(keys.authenticatedAccounts);
-    for (const userId in authenticatedUserIds) {
+    for (const userId of authenticatedUserIds) {
       const account = await this.get<TAccount>(userId);
       const encryptedOrgKeys = account.keys.organizationKeys?.encrypted;
-      if (encryptedOrgKeys != null) {
-        // TODO, iterate through KVPs and update value
+      if (encryptedOrgKeys == null) {
+        continue;
       }
+
+      for (const [orgId, encKey] of Object.entries(encryptedOrgKeys)) {
+        encryptedOrgKeys[orgId] = new EncryptedOrganizationKeyData(encKey as unknown as string);
+      }
+
+      this.set(userId, account);
     }
 
     await this.setCurrentStateVersion(StateVersion.Five);