[PM-6172] Run localStorage migrations for web (#7900)

* Create MigrationRunner

- Create MigrationRunner Service for running migrations in StateService
- Create web override so that migrations also run against `localStorage`

* Fix Web StateService

* Fix WebMigrationRunner

* Fix CLI

* Fix ElectronStateService

* Update Comment

* More Common Scenarios

libs/common/src/platform/services/migration-builder.service.spec.ts

@@ -0,0 +1,89 @@
+import { mock } from "jest-mock-extended";
+
+import { FakeStorageService } from "../../../spec/fake-storage.service";
+import { MigrationHelper } from "../../state-migrations/migration-helper";
+
+import { MigrationBuilderService } from "./migration-builder.service";
+
+describe("MigrationBuilderService", () => {
+  // All migrations from 10+ should be capable of having a null account object or null global object
+  const startingStateVersion = 10;
+
+  const noAccounts = {
+    stateVersion: startingStateVersion,
+    authenticatedAccounts: <string[]>[],
+  };
+
+  const nullAndUndefinedAccounts = {
+    stateVersion: startingStateVersion,
+    authenticatedAccounts: ["account1", "account2"],
+    account1: <object>null,
+    account2: <object>undefined,
+  };
+
+  const emptyAccountObject = {
+    stateVersion: startingStateVersion,
+    authenticatedAccounts: ["account1"],
+    account1: {},
+  };
+
+  const nullCommonAccountProperties = {
+    stateVersion: startingStateVersion,
+    authenticatedAccounts: ["account1"],
+    account1: {
+      data: <object>null,
+      keys: <object>null,
+      profile: <object>null,
+      settings: <object>null,
+      tokens: <object>null,
+    },
+  };
+
+  const emptyCommonAccountProperties = {
+    stateVersion: startingStateVersion,
+    authenticatedAccounts: ["account1"],
+    account1: {
+      data: {},
+      keys: {},
+      profile: {},
+      settings: {},
+      tokens: {},
+    },
+  };
+
+  const nullGlobal = {
+    stateVersion: startingStateVersion,
+    global: <object>null,
+  };
+
+  const undefinedGlobal = {
+    stateVersion: startingStateVersion,
+    global: <object>undefined,
+  };
+
+  const emptyGlobalObject = {
+    stateVersion: startingStateVersion,
+    global: {},
+  };
+
+  it.each([
+    noAccounts,
+    nullAndUndefinedAccounts,
+    emptyAccountObject,
+    nullCommonAccountProperties,
+    emptyCommonAccountProperties,
+    nullGlobal,
+    undefinedGlobal,
+    emptyGlobalObject,
+  ])("should not produce migrations that throw when given data: %s", async (startingState) => {
+    const sut = new MigrationBuilderService();
+
+    const helper = new MigrationHelper(
+      startingStateVersion,
+      new FakeStorageService(startingState),
+      mock(),
+    );
+
+    await sut.build().migrate(helper);
+  });
+});

libs/common/src/platform/services/migration-builder.service.ts

@@ -0,0 +1,10 @@
+import { createMigrationBuilder } from "../../state-migrations";
+import { MigrationBuilder } from "../../state-migrations/migration-builder";
+
+export class MigrationBuilderService {
+  private migrationBuilderCache: MigrationBuilder;
+
+  build() {
+    return (this.migrationBuilderCache ??= createMigrationBuilder());
+  }
+}

libs/common/src/platform/services/migration-runner.spec.ts

@@ -0,0 +1,102 @@
+import { mock } from "jest-mock-extended";
+
+import { awaitAsync } from "../../../spec";
+import { CURRENT_VERSION } from "../../state-migrations";
+import { MigrationBuilder } from "../../state-migrations/migration-builder";
+import { LogService } from "../abstractions/log.service";
+import { AbstractStorageService } from "../abstractions/storage.service";
+
+import { MigrationBuilderService } from "./migration-builder.service";
+import { MigrationRunner } from "./migration-runner";
+
+describe("MigrationRunner", () => {
+  const storage = mock<AbstractStorageService>();
+  const logService = mock<LogService>();
+  const migrationBuilderService = mock<MigrationBuilderService>();
+  const mockMigrationBuilder = mock<MigrationBuilder>();
+
+  migrationBuilderService.build.mockReturnValue(mockMigrationBuilder);
+
+  const sut = new MigrationRunner(storage, logService, migrationBuilderService);
+
+  describe("migrate", () => {
+    it("should not run migrations if state is empty", async () => {
+      storage.get.mockReturnValueOnce(null);
+      await sut.run();
+      expect(migrationBuilderService.build).not.toHaveBeenCalled();
+    });
+
+    it("should set to current version if state is empty", async () => {
+      storage.get.mockReturnValueOnce(null);
+      await sut.run();
+      expect(storage.save).toHaveBeenCalledWith("stateVersion", CURRENT_VERSION);
+    });
+
+    it("should run migration if there is a stateVersion", async () => {
+      storage.get.mockResolvedValueOnce(12);
+
+      await sut.run();
+
+      expect(mockMigrationBuilder.migrate).toHaveBeenCalled();
+    });
+  });
+
+  describe("waitForCompletion", () => {
+    it("should wait until stateVersion is current before completing", async () => {
+      let stateVersion: number | null = null;
+
+      storage.get.mockImplementation((key) => {
+        if (key === "stateVersion") {
+          return Promise.resolve(stateVersion);
+        }
+      });
+
+      let promiseCompleted = false;
+
+      const completionPromise = sut.waitForCompletion().then(() => (promiseCompleted = true));
+
+      await awaitAsync(10);
+
+      expect(promiseCompleted).toBe(false);
+
+      stateVersion = CURRENT_VERSION;
+      await completionPromise;
+    });
+
+    // Skipped for CI since this test takes a while to complete, remove `.skip` to test
+    it.skip(
+      "will complete after 8 second step wait if migrations still aren't complete",
+      async () => {
+        storage.get.mockImplementation((key) => {
+          if (key === "stateVersion") {
+            return Promise.resolve(null);
+          }
+        });
+
+        let promiseCompleted = false;
+
+        void sut.waitForCompletion().then(() => (promiseCompleted = true));
+
+        await awaitAsync(2 + 4 + 8 + 16);
+
+        expect(promiseCompleted).toBe(false);
+
+        await awaitAsync(32 + 64 + 128 + 256);
+
+        expect(promiseCompleted).toBe(false);
+
+        await awaitAsync(512 + 1024 + 2048 + 4096);
+
+        expect(promiseCompleted).toBe(false);
+
+        const SKEW = 20;
+
+        await awaitAsync(8192 + SKEW);
+
+        expect(promiseCompleted).toBe(true);
+      },
+      // Have to combine all the steps into the timeout to get this to run
+      2 + 4 + 8 + 16 + 32 + 64 + 128 + 256 + 512 + 1024 + 2048 + 4096 + 8192 + 100,
+    );
+  });
+});

libs/common/src/platform/services/migration-runner.ts

@@ -0,0 +1,37 @@
+import { waitForMigrations } from "../../state-migrations";
+import { CURRENT_VERSION, currentVersion } from "../../state-migrations/migrate";
+import { MigrationHelper } from "../../state-migrations/migration-helper";
+import { LogService } from "../abstractions/log.service";
+import { AbstractStorageService } from "../abstractions/storage.service";
+
+import { MigrationBuilderService } from "./migration-builder.service";
+
+export class MigrationRunner {
+  constructor(
+    protected diskStorage: AbstractStorageService,
+    protected logService: LogService,
+    protected migrationBuilderService: MigrationBuilderService,
+  ) {}
+
+  async run(): Promise<void> {
+    const migrationHelper = new MigrationHelper(
+      await currentVersion(this.diskStorage, this.logService),
+      this.diskStorage,
+      this.logService,
+    );
+
+    if (migrationHelper.currentVersion < 0) {
+      // Cannot determine state, assuming empty so we don't repeatedly apply a migration.
+      await this.diskStorage.save("stateVersion", CURRENT_VERSION);
+      return;
+    }
+
+    const migrationBuilder = this.migrationBuilderService.build();
+
+    await migrationBuilder.migrate(migrationHelper);
+  }
+
+  async waitForCompletion(): Promise<void> {
+    await waitForMigrations(this.diskStorage, this.logService);
+  }
+}

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

@@ -14,8 +14,6 @@ import { BiometricKey } from "../../auth/types/biometric-key";
 import { VaultTimeoutAction } from "../../enums/vault-timeout-action.enum";
 import { EventData } from "../../models/data/event.data";
 import { WindowState } from "../../models/domain/window-state";
-import { migrate } from "../../state-migrations";
-import { waitForMigrations } from "../../state-migrations/migrate";
 import { GeneratorOptions } from "../../tools/generator/generator-options";
 import { GeneratedPasswordHistory, PasswordGeneratorOptions } from "../../tools/generator/password";
 import { UsernameGeneratorOptions } from "../../tools/generator/username";
@@ -57,6 +55,8 @@ import { State } from "../models/domain/state";
 import { StorageOptions } from "../models/domain/storage-options";
 import { SymmetricCryptoKey } from "../models/domain/symmetric-crypto-key";
 
+import { MigrationRunner } from "./migration-runner";
+
 const keys = {
   state: "state",
   stateVersion: "stateVersion",
@@ -108,6 +108,7 @@ export class StateService<
     protected stateFactory: StateFactory<TGlobalState, TAccount>,
     protected accountService: AccountService,
     protected environmentService: EnvironmentService,
+    private migrationRunner: MigrationRunner,
     protected useAccountCache: boolean = true,
   ) {
     // If the account gets changed, verify the new account is unlocked
@@ -136,11 +137,11 @@ export class StateService<
     }
 
     if (runMigrations) {
-      await migrate(this.storageService, this.logService);
+      await this.migrationRunner.run();
     } else {
       // It may have been requested to not run the migrations but we should defensively not
       // continue this method until migrations have a chance to be completed elsewhere.
-      await waitForMigrations(this.storageService, this.logService);
+      await this.migrationRunner.waitForCompletion();
     }
 
     await this.state().then(async (state) => {

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

@@ -1 +1 @@
-export { migrate, CURRENT_VERSION } from "./migrate";
+export { createMigrationBuilder, waitForMigrations, CURRENT_VERSION } from "./migrate";

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

@@ -5,34 +5,7 @@ import { LogService } from "../platform/abstractions/log.service";
 // eslint-disable-next-line import/no-restricted-paths -- Needed to interface with storage locations
 import { AbstractStorageService } from "../platform/abstractions/storage.service";
 
-import { CURRENT_VERSION, currentVersion, migrate } from "./migrate";
-import { MigrationBuilder } from "./migration-builder";
-
-jest.mock("./migration-builder", () => {
-  return {
-    MigrationBuilder: {
-      create: jest.fn().mockReturnThis(),
-    },
-  };
-});
-
-describe("migrate", () => {
-  it("should not run migrations if state is empty", async () => {
-    const storage = mock<AbstractStorageService>();
-    const logService = mock<LogService>();
-    storage.get.mockReturnValueOnce(null);
-    await migrate(storage, logService);
-    expect(MigrationBuilder.create).not.toHaveBeenCalled();
-  });
-
-  it("should set to current version if state is empty", async () => {
-    const storage = mock<AbstractStorageService>();
-    const logService = mock<LogService>();
-    storage.get.mockReturnValueOnce(null);
-    await migrate(storage, logService);
-    expect(storage.save).toHaveBeenCalledWith("stateVersion", CURRENT_VERSION);
-  });
-});
+import { currentVersion } from "./migrate";
 
 describe("currentVersion", () => {
   let storage: MockProxy<AbstractStorageService>;

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

@@ -4,7 +4,6 @@ import { LogService } from "../platform/abstractions/log.service";
 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 { MoveEnvironmentStateToProviders } from "./migrations/12-move-environment-state-to-providers";
@@ -27,21 +26,8 @@ export const MIN_VERSION = 2;
 export const CURRENT_VERSION = 18;
 export type MinVersion = typeof MIN_VERSION;
 
-export async function migrate(
-  storageService: AbstractStorageService,
-  logService: LogService,
-): Promise<void> {
-  const migrationHelper = new MigrationHelper(
-    await currentVersion(storageService, logService),
-    storageService,
-    logService,
-  );
-  if (migrationHelper.currentVersion < 0) {
-    // Cannot determine state, assuming empty so we don't repeatedly apply a migration.
-    await storageService.save("stateVersion", CURRENT_VERSION);
-    return;
-  }
-  await MigrationBuilder.create()
+export function createMigrationBuilder() {
+  return MigrationBuilder.create()
     .with(MinVersionMigrator)
     .with(FixPremiumMigrator, 2, 3)
     .with(RemoveEverBeenUnlockedMigrator, 3, 4)
@@ -58,9 +44,7 @@ export async function migrate(
     .with(FolderMigrator, 14, 15)
     .with(LastSyncMigrator, 15, 16)
     .with(EnablePasskeysMigrator, 16, 17)
-    .with(AutofillSettingsKeyMigrator, 17, CURRENT_VERSION)
-
-    .migrate(migrationHelper);
+    .with(AutofillSettingsKeyMigrator, 17, CURRENT_VERSION);
 }
 
 export async function currentVersion(