From 3edf098aafb01f739f147f66a9c36e68e24b29fb Mon Sep 17 00:00:00 2001
From: Jason Ng <jng@bitwarden.com>
Date: Wed, 14 Feb 2024 17:03:03 -0500
Subject: [PATCH] PM-5274 Migrate Collection Service State (#7732)

* update collection service to use new state provider framework, remove stateservice from collection service, update collections state provider with migrate file and unit test
---
 .../browser/src/background/main.background.ts |   2 +-
 .../collection-service.factory.ts             |   8 +-
 apps/cli/src/bw.ts                            |   5 +-
 apps/web/src/app/core/state/state.service.ts  |  16 --
 .../vault-header/vault-header.component.html  |   2 +-
 .../vault-header/vault-header.component.ts    |   2 +-
 .../vault-header/vault-header.component.html  |   2 +-
 .../src/services/jslib-services.module.ts     |   2 +-
 .../platform/abstractions/state.service.ts    |   9 -
 .../src/platform/models/domain/account.ts     |   8 +-
 .../src/platform/services/state.service.ts    |  43 ----
 .../src/platform/state/state-definitions.ts   |   3 +
 .../vault-timeout/vault-timeout.service.ts    |   2 +-
 libs/common/src/state-migrations/migrate.ts   |   6 +-
 ...ollections-state-to-state-provider.spec.ts | 196 ++++++++++++++++++
 ...ove-collections-state-to-state-provider.ts |  63 ++++++
 libs/common/src/types/guid.ts                 |   1 +
 .../vault/abstractions/collection.service.ts  |   9 +-
 .../src/vault/models/data/collection.data.ts  |  11 +-
 .../vault/models/domain/collection.spec.ts    |   5 +-
 .../models/response/collection.response.ts    |   5 +-
 .../src/vault/models/view/collection.view.ts  |   6 +
 .../src/vault/services/collection.service.ts  | 190 ++++++++++-------
 23 files changed, 426 insertions(+), 170 deletions(-)
 create mode 100644 libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.spec.ts
 create mode 100644 libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.ts

diff --git a/apps/browser/src/background/main.background.ts b/apps/browser/src/background/main.background.ts
index 13732ed0f3a..9420e7d4157 100644
--- a/apps/browser/src/background/main.background.ts
+++ b/apps/browser/src/background/main.background.ts
@@ -454,7 +454,7 @@ export default class MainBackground {
     this.collectionService = new CollectionService(
       this.cryptoService,
       this.i18nService,
-      this.stateService,
+      this.stateProvider,
     );
     this.syncNotifierService = new SyncNotifierService();
     this.organizationService = new BrowserOrganizationService(
diff --git a/apps/browser/src/vault/background/service_factories/collection-service.factory.ts b/apps/browser/src/vault/background/service_factories/collection-service.factory.ts
index fd576ca12c3..69d5cd6f22f 100644
--- a/apps/browser/src/vault/background/service_factories/collection-service.factory.ts
+++ b/apps/browser/src/vault/background/service_factories/collection-service.factory.ts
@@ -14,10 +14,8 @@ import {
   i18nServiceFactory,
   I18nServiceInitOptions,
 } from "../../../platform/background/service-factories/i18n-service.factory";
-import {
-  stateServiceFactory as stateServiceFactory,
-  StateServiceInitOptions,
-} from "../../../platform/background/service-factories/state-service.factory";
+import { stateProviderFactory } from "../../../platform/background/service-factories/state-provider.factory";
+import { StateServiceInitOptions } from "../../../platform/background/service-factories/state-service.factory";
 
 type CollectionServiceFactoryOptions = FactoryOptions;
 
@@ -38,7 +36,7 @@ export function collectionServiceFactory(
       new CollectionService(
         await cryptoServiceFactory(cache, opts),
         await i18nServiceFactory(cache, opts),
-        await stateServiceFactory(cache, opts),
+        await stateProviderFactory(cache, opts),
       ),
   );
 }
diff --git a/apps/cli/src/bw.ts b/apps/cli/src/bw.ts
index 05ad43128e5..d1c46e9e0b2 100644
--- a/apps/cli/src/bw.ts
+++ b/apps/cli/src/bw.ts
@@ -87,6 +87,7 @@ import {
 } from "@bitwarden/common/tools/password-strength";
 import { SendApiService } from "@bitwarden/common/tools/send/services/send-api.service";
 import { SendService } from "@bitwarden/common/tools/send/services/send.service";
+import { UserId } from "@bitwarden/common/types/guid";
 import { InternalFolderService } from "@bitwarden/common/vault/abstractions/folder/folder.service.abstraction";
 import { CipherService } from "@bitwarden/common/vault/services/cipher.service";
 import { CollectionService } from "@bitwarden/common/vault/services/collection.service";
@@ -358,7 +359,7 @@ export class Main {
     this.collectionService = new CollectionService(
       this.cryptoService,
       this.i18nService,
-      this.stateService,
+      this.stateProvider,
     );
 
     this.providerService = new ProviderService(this.stateService);
@@ -616,7 +617,7 @@ export class Main {
       this.settingsService.clear(userId),
       this.cipherService.clear(userId),
       this.folderService.clear(userId),
-      this.collectionService.clear(userId),
+      this.collectionService.clear(userId as UserId),
       this.policyService.clear(userId),
       this.passwordGenerationService.clear(),
     ]);
diff --git a/apps/web/src/app/core/state/state.service.ts b/apps/web/src/app/core/state/state.service.ts
index a40a24b2155..b46c2b590a0 100644
--- a/apps/web/src/app/core/state/state.service.ts
+++ b/apps/web/src/app/core/state/state.service.ts
@@ -19,7 +19,6 @@ import { MigrationRunner } from "@bitwarden/common/platform/services/migration-r
 import { StateService as BaseStateService } from "@bitwarden/common/platform/services/state.service";
 import { SendData } from "@bitwarden/common/tools/send/models/data/send.data";
 import { CipherData } from "@bitwarden/common/vault/models/data/cipher.data";
-import { CollectionData } from "@bitwarden/common/vault/models/data/collection.data";
 
 import { Account } from "./account";
 import { GlobalState } from "./global-state";
@@ -69,21 +68,6 @@ export class StateService extends BaseStateService<GlobalState, Account> {
     return await super.setEncryptedCiphers(value, options);
   }
 
-  async getEncryptedCollections(
-    options?: StorageOptions,
-  ): Promise<{ [id: string]: CollectionData }> {
-    options = this.reconcileOptions(options, await this.defaultInMemoryOptions());
-    return await super.getEncryptedCollections(options);
-  }
-
-  async setEncryptedCollections(
-    value: { [id: string]: CollectionData },
-    options?: StorageOptions,
-  ): Promise<void> {
-    options = this.reconcileOptions(options, await this.defaultInMemoryOptions());
-    return await super.setEncryptedCollections(value, options);
-  }
-
   async getEncryptedSends(options?: StorageOptions): Promise<{ [id: string]: SendData }> {
     options = this.reconcileOptions(options, await this.defaultInMemoryOptions());
     return await super.getEncryptedSends(options);
diff --git a/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.html b/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.html
index 851e252e39a..ca8c71cf5df 100644
--- a/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.html
+++ b/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.html
@@ -28,7 +28,7 @@
         aria-hidden="true"
       ></i>
       <span>{{ title }}</span>
-      <ng-container *ngIf="collection !== undefined && (canEditCollection || canDeleteCollection)">
+      <ng-container *ngIf="collection != null && (canEditCollection || canDeleteCollection)">
         <button
           bitIconButton="bwi-angle-down"
           [bitMenuTriggerFor]="editCollectionMenu"
diff --git a/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.ts b/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.ts
index 6374178f24b..0416971e0d0 100644
--- a/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.ts
+++ b/apps/web/src/app/vault/individual-vault/vault-header/vault-header.component.ts
@@ -125,7 +125,7 @@ export class VaultHeaderComponent {
 
   get canEditCollection(): boolean {
     // Only edit collections if not editing "Unassigned"
-    if (this.collection === undefined) {
+    if (this.collection == null) {
       return false;
     }
 
diff --git a/apps/web/src/app/vault/org-vault/vault-header/vault-header.component.html b/apps/web/src/app/vault/org-vault/vault-header/vault-header.component.html
index 1d65ae53d68..01c8e2f1beb 100644
--- a/apps/web/src/app/vault/org-vault/vault-header/vault-header.component.html
+++ b/apps/web/src/app/vault/org-vault/vault-header/vault-header.component.html
@@ -33,7 +33,7 @@
         aria-hidden="true"
       ></i>
       <span>{{ title }}</span>
-      <ng-container *ngIf="collection !== undefined && (canEditCollection || canDeleteCollection)">
+      <ng-container *ngIf="collection != null && (canEditCollection || canDeleteCollection)">
         <button
           bitIconButton="bwi-angle-down"
           [bitMenuTriggerFor]="editCollectionMenu"
diff --git a/libs/angular/src/services/jslib-services.module.ts b/libs/angular/src/services/jslib-services.module.ts
index a3e6d1222ee..b833a3b929c 100644
--- a/libs/angular/src/services/jslib-services.module.ts
+++ b/libs/angular/src/services/jslib-services.module.ts
@@ -418,7 +418,7 @@ import { ModalService } from "./modal.service";
     {
       provide: CollectionServiceAbstraction,
       useClass: CollectionService,
-      deps: [CryptoServiceAbstraction, I18nServiceAbstraction, StateServiceAbstraction],
+      deps: [CryptoServiceAbstraction, I18nServiceAbstraction, StateProvider],
     },
     {
       provide: EnvironmentServiceAbstraction,
diff --git a/libs/common/src/platform/abstractions/state.service.ts b/libs/common/src/platform/abstractions/state.service.ts
index b59c69a92ff..9e28536b65d 100644
--- a/libs/common/src/platform/abstractions/state.service.ts
+++ b/libs/common/src/platform/abstractions/state.service.ts
@@ -19,10 +19,8 @@ import { UserId } from "../../types/guid";
 import { DeviceKey, MasterKey, UserKey } from "../../types/key";
 import { UriMatchType } from "../../vault/enums";
 import { CipherData } from "../../vault/models/data/cipher.data";
-import { CollectionData } from "../../vault/models/data/collection.data";
 import { LocalData } from "../../vault/models/data/local.data";
 import { CipherView } from "../../vault/models/view/cipher.view";
-import { CollectionView } from "../../vault/models/view/collection.view";
 import { AddEditCipherInfo } from "../../vault/types/add-edit-cipher-info";
 import { KdfType, ThemeType } from "../enums";
 import { ServerConfigData } from "../models/data/server-config.data";
@@ -200,8 +198,6 @@ export abstract class StateService<T extends Account = Account> {
   setBiometricPromptCancelled: (value: boolean, options?: StorageOptions) => Promise<void>;
   getDecryptedCiphers: (options?: StorageOptions) => Promise<CipherView[]>;
   setDecryptedCiphers: (value: CipherView[], options?: StorageOptions) => Promise<void>;
-  getDecryptedCollections: (options?: StorageOptions) => Promise<CollectionView[]>;
-  setDecryptedCollections: (value: CollectionView[], options?: StorageOptions) => Promise<void>;
   getDecryptedPasswordGenerationHistory: (
     options?: StorageOptions,
   ) => Promise<GeneratedPasswordHistory[]>;
@@ -313,11 +309,6 @@ export abstract class StateService<T extends Account = Account> {
     value: { [id: string]: CipherData },
     options?: StorageOptions,
   ) => Promise<void>;
-  getEncryptedCollections: (options?: StorageOptions) => Promise<{ [id: string]: CollectionData }>;
-  setEncryptedCollections: (
-    value: { [id: string]: CollectionData },
-    options?: StorageOptions,
-  ) => Promise<void>;
   getEncryptedPasswordGenerationHistory: (
     options?: StorageOptions,
   ) => Promise<GeneratedPasswordHistory[]>;
diff --git a/libs/common/src/platform/models/domain/account.ts b/libs/common/src/platform/models/domain/account.ts
index fc2350941d7..233276fa078 100644
--- a/libs/common/src/platform/models/domain/account.ts
+++ b/libs/common/src/platform/models/domain/account.ts
@@ -22,10 +22,8 @@ import { DeepJsonify } from "../../../types/deep-jsonify";
 import { MasterKey, UserKey } from "../../../types/key";
 import { UriMatchType } from "../../../vault/enums";
 import { CipherData } from "../../../vault/models/data/cipher.data";
-import { CollectionData } from "../../../vault/models/data/collection.data";
 import { FolderData } from "../../../vault/models/data/folder.data";
 import { CipherView } from "../../../vault/models/view/cipher.view";
-import { CollectionView } from "../../../vault/models/view/collection.view";
 import { AddEditCipherInfo } from "../../../vault/types/add-edit-cipher-info";
 import { KdfType } from "../../enums";
 import { Utils } from "../../misc/utils";
@@ -73,7 +71,7 @@ export class EncryptionPair<TEncrypted, TDecrypted> {
 }
 
 export class DataEncryptionPair<TEncrypted, TDecrypted> {
-  encrypted?: { [id: string]: TEncrypted };
+  encrypted?: Record<string, TEncrypted>;
   decrypted?: TDecrypted[];
 }
 
@@ -92,10 +90,6 @@ export class AccountData {
   folders? = new TemporaryDataEncryption<FolderData>();
   localData?: any;
   sends?: DataEncryptionPair<SendData, SendView> = new DataEncryptionPair<SendData, SendView>();
-  collections?: DataEncryptionPair<CollectionData, CollectionView> = new DataEncryptionPair<
-    CollectionData,
-    CollectionView
-  >();
   policies?: DataEncryptionPair<PolicyData, Policy> = new DataEncryptionPair<PolicyData, Policy>();
   passwordGenerationHistory?: EncryptionPair<
     GeneratedPasswordHistory[],
diff --git a/libs/common/src/platform/services/state.service.ts b/libs/common/src/platform/services/state.service.ts
index 0f42b6660a6..be29627ef70 100644
--- a/libs/common/src/platform/services/state.service.ts
+++ b/libs/common/src/platform/services/state.service.ts
@@ -23,10 +23,8 @@ import { UserId } from "../../types/guid";
 import { DeviceKey, MasterKey, UserKey } from "../../types/key";
 import { UriMatchType } from "../../vault/enums";
 import { CipherData } from "../../vault/models/data/cipher.data";
-import { CollectionData } from "../../vault/models/data/collection.data";
 import { LocalData } from "../../vault/models/data/local.data";
 import { CipherView } from "../../vault/models/view/cipher.view";
-import { CollectionView } from "../../vault/models/view/collection.view";
 import { AddEditCipherInfo } from "../../vault/types/add-edit-cipher-info";
 import { EnvironmentService } from "../abstractions/environment.service";
 import { LogService } from "../abstractions/log.service";
@@ -939,24 +937,6 @@ export class StateService<
     );
   }
 
-  @withPrototypeForArrayMembers(CollectionView)
-  async getDecryptedCollections(options?: StorageOptions): Promise<CollectionView[]> {
-    return (
-      await this.getAccount(this.reconcileOptions(options, await this.defaultInMemoryOptions()))
-    )?.data?.collections?.decrypted;
-  }
-
-  async setDecryptedCollections(value: CollectionView[], options?: StorageOptions): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-    account.data.collections.decrypted = value;
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-  }
-
   /**
    * @deprecated Use UserKey instead
    */
@@ -1628,29 +1608,6 @@ export class StateService<
     );
   }
 
-  @withPrototypeForObjectValues(CollectionData)
-  async getEncryptedCollections(
-    options?: StorageOptions,
-  ): Promise<{ [id: string]: CollectionData }> {
-    return (
-      await this.getAccount(this.reconcileOptions(options, await this.defaultOnDiskMemoryOptions()))
-    )?.data?.collections?.encrypted;
-  }
-
-  async setEncryptedCollections(
-    value: { [id: string]: CollectionData },
-    options?: StorageOptions,
-  ): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultOnDiskMemoryOptions()),
-    );
-    account.data.collections.encrypted = value;
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultOnDiskMemoryOptions()),
-    );
-  }
-
   /**
    * @deprecated Use UserKey instead
    */
diff --git a/libs/common/src/platform/state/state-definitions.ts b/libs/common/src/platform/state/state-definitions.ts
index bc51af6ae25..38c7ba9b8ff 100644
--- a/libs/common/src/platform/state/state-definitions.ts
+++ b/libs/common/src/platform/state/state-definitions.ts
@@ -50,6 +50,9 @@ export const VAULT_SETTINGS_DISK = new StateDefinition("vaultSettings", "disk",
   web: "disk-local",
 });
 
+export const COLLECTION_DATA = new StateDefinition("collection", "disk", {
+  web: "memory",
+});
 export const AUTOFILL_SETTINGS_DISK = new StateDefinition("autofillSettings", "disk");
 export const AUTOFILL_SETTINGS_DISK_LOCAL = new StateDefinition("autofillSettingsLocal", "disk", {
   web: "disk-local",
diff --git a/libs/common/src/services/vault-timeout/vault-timeout.service.ts b/libs/common/src/services/vault-timeout/vault-timeout.service.ts
index 584f9713d9f..6163005367f 100644
--- a/libs/common/src/services/vault-timeout/vault-timeout.service.ts
+++ b/libs/common/src/services/vault-timeout/vault-timeout.service.ts
@@ -84,6 +84,7 @@ export class VaultTimeoutService implements VaultTimeoutServiceAbstraction {
     if (userId == null || userId === (await this.stateService.getUserId())) {
       this.searchService.clearIndex();
       await this.folderService.clearCache();
+      await this.collectionService.clearActiveUserCache();
     }
 
     await this.stateService.setEverBeenUnlocked(true, { userId: userId });
@@ -96,7 +97,6 @@ export class VaultTimeoutService implements VaultTimeoutServiceAbstraction {
     await this.cryptoService.clearKeyPair(true, userId);
 
     await this.cipherService.clearCache(userId);
-    await this.collectionService.clearCache(userId);
 
     this.messagingService.send("locked", { userId: userId });
 
diff --git a/libs/common/src/state-migrations/migrate.ts b/libs/common/src/state-migrations/migrate.ts
index 649fbc14a73..959ae4c0217 100644
--- a/libs/common/src/state-migrations/migrate.ts
+++ b/libs/common/src/state-migrations/migrate.ts
@@ -15,6 +15,7 @@ import { EnablePasskeysMigrator } from "./migrations/17-move-enable-passkeys-to-
 import { AutofillSettingsKeyMigrator } from "./migrations/18-move-autofill-settings-to-state-providers";
 import { RequirePasswordOnStartMigrator } from "./migrations/19-migrate-require-password-on-start";
 import { PrivateKeyMigrator } from "./migrations/20-move-private-key-to-state-providers";
+import { CollectionMigrator } from "./migrations/21-move-collections-state-to-state-provider";
 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";
@@ -25,7 +26,7 @@ import { MoveBrowserSettingsToGlobal } from "./migrations/9-move-browser-setting
 import { MinVersionMigrator } from "./migrations/min-version";
 
 export const MIN_VERSION = 2;
-export const CURRENT_VERSION = 20;
+export const CURRENT_VERSION = 21;
 export type MinVersion = typeof MIN_VERSION;
 
 export function createMigrationBuilder() {
@@ -48,7 +49,8 @@ export function createMigrationBuilder() {
     .with(EnablePasskeysMigrator, 16, 17)
     .with(AutofillSettingsKeyMigrator, 17, 18)
     .with(RequirePasswordOnStartMigrator, 18, 19)
-    .with(PrivateKeyMigrator, 19, CURRENT_VERSION);
+    .with(PrivateKeyMigrator, 19, 20)
+    .with(CollectionMigrator, 20, CURRENT_VERSION);
 }
 
 export async function currentVersion(
diff --git a/libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.spec.ts b/libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.spec.ts
new file mode 100644
index 00000000000..33699117635
--- /dev/null
+++ b/libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.spec.ts
@@ -0,0 +1,196 @@
+import { MockProxy, any } from "jest-mock-extended";
+
+import { MigrationHelper } from "../migration-helper";
+import { mockMigrationHelper } from "../migration-helper.spec";
+
+import { CollectionMigrator } from "./21-move-collections-state-to-state-provider";
+
+function exampleJSON() {
+  return {
+    global: {
+      otherStuff: "otherStuff1",
+    },
+    authenticatedAccounts: ["user-1", "user-2"],
+    "user-1": {
+      data: {
+        collections: {
+          encrypted: {
+            "877fef70-be32-439e-8678-b0d80125653d": {
+              id: "877fef70-be32-439e-8678-b0d80125653d",
+              organizationId: "fe1ff6ef-d2d4-49f3-9c07-b0c7013998f9",
+              name: "2.MD9OMDsvYiU1CTSUxjHorw==|uFc4cZhnmQmK2LFCWbyeZg==|syk2d9JESeplxInLvP36BK5RhqS1c/i+ZQp5NR7EUA4=",
+              externalId: "",
+              readOnly: false,
+              manage: true,
+              hidePasswords: false,
+            },
+            "0d3fee82-3f81-434c-aed0-b0c200ee6c7a": {
+              id: "0d3fee82-3f81-434c-aed0-b0c200ee6c7a",
+              organizationId: "5f277723-6391-4b5c-add9-b0c200ee6967",
+              name: "2.GxnXkIbBCGFr57F6lT7+Ow==|3ctMg95FKquG3l+qfv8BgvaCbYzMmuhnukCEHXhUukE=|cJRZWq05xjPBayUgx6P6gsbtNVLi8exQwo8F1SfqQQ4=",
+              externalId: "",
+              readOnly: false,
+              manage: false,
+              hidePasswords: false,
+            },
+          },
+        },
+        otherStuff: "otherStuff2",
+      },
+      otherStuff: "otherStuff3",
+    },
+    "user-2": {
+      data: {
+        otherStuff: "otherStuff4",
+      },
+      otherStuff: "otherStuff5",
+    },
+  };
+}
+
+function rollbackJSON() {
+  return {
+    "user_user-1_collection_collections": {
+      "877fef70-be32-439e-8678-b0d80125653d": {
+        id: "877fef70-be32-439e-8678-b0d80125653d",
+        organizationId: "fe1ff6ef-d2d4-49f3-9c07-b0c7013998f9",
+        name: "2.MD9OMDsvYiU1CTSUxjHorw==|uFc4cZhnmQmK2LFCWbyeZg==|syk2d9JESeplxInLvP36BK5RhqS1c/i+ZQp5NR7EUA4=",
+        externalId: "",
+        readOnly: false,
+        manage: true,
+        hidePasswords: false,
+      },
+      "0d3fee82-3f81-434c-aed0-b0c200ee6c7a": {
+        id: "0d3fee82-3f81-434c-aed0-b0c200ee6c7a",
+        organizationId: "5f277723-6391-4b5c-add9-b0c200ee6967",
+        name: "2.GxnXkIbBCGFr57F6lT7+Ow==|3ctMg95FKquG3l+qfv8BgvaCbYzMmuhnukCEHXhUukE=|cJRZWq05xjPBayUgx6P6gsbtNVLi8exQwo8F1SfqQQ4=",
+        externalId: "",
+        readOnly: false,
+        manage: false,
+        hidePasswords: false,
+      },
+    },
+    "user_user-2_collection_data": null as any,
+    global: {
+      otherStuff: "otherStuff1",
+    },
+    authenticatedAccounts: ["user-1", "user-2"],
+    "user-1": {
+      data: {
+        otherStuff: "otherStuff2",
+      },
+      otherStuff: "otherStuff3",
+    },
+    "user-2": {
+      data: {
+        otherStuff: "otherStuff4",
+      },
+      otherStuff: "otherStuff5",
+    },
+  };
+}
+
+describe("CollectionMigrator", () => {
+  let helper: MockProxy<MigrationHelper>;
+  let sut: CollectionMigrator;
+  const keyDefinitionLike = {
+    key: "collections",
+    stateDefinition: {
+      name: "collection",
+    },
+  };
+
+  describe("migrate", () => {
+    beforeEach(() => {
+      helper = mockMigrationHelper(exampleJSON(), 20);
+      sut = new CollectionMigrator(20, 21);
+    });
+
+    it("should remove collections from all accounts", async () => {
+      await sut.migrate(helper);
+      expect(helper.set).toHaveBeenCalledWith("user-1", {
+        data: {
+          otherStuff: "otherStuff2",
+        },
+        otherStuff: "otherStuff3",
+      });
+    });
+
+    it("should set collections value for each account", async () => {
+      await sut.migrate(helper);
+
+      expect(helper.setToUser).toHaveBeenCalledWith("user-1", keyDefinitionLike, {
+        "877fef70-be32-439e-8678-b0d80125653d": {
+          id: "877fef70-be32-439e-8678-b0d80125653d",
+          organizationId: "fe1ff6ef-d2d4-49f3-9c07-b0c7013998f9",
+          name: "2.MD9OMDsvYiU1CTSUxjHorw==|uFc4cZhnmQmK2LFCWbyeZg==|syk2d9JESeplxInLvP36BK5RhqS1c/i+ZQp5NR7EUA4=",
+          externalId: "",
+          readOnly: false,
+          manage: true,
+          hidePasswords: false,
+        },
+        "0d3fee82-3f81-434c-aed0-b0c200ee6c7a": {
+          id: "0d3fee82-3f81-434c-aed0-b0c200ee6c7a",
+          organizationId: "5f277723-6391-4b5c-add9-b0c200ee6967",
+          name: "2.GxnXkIbBCGFr57F6lT7+Ow==|3ctMg95FKquG3l+qfv8BgvaCbYzMmuhnukCEHXhUukE=|cJRZWq05xjPBayUgx6P6gsbtNVLi8exQwo8F1SfqQQ4=",
+          externalId: "",
+          readOnly: false,
+          manage: false,
+          hidePasswords: false,
+        },
+      });
+    });
+  });
+
+  describe("rollback", () => {
+    beforeEach(() => {
+      helper = mockMigrationHelper(rollbackJSON(), 21);
+      sut = new CollectionMigrator(20, 21);
+    });
+
+    it.each(["user-1", "user-2"])("should null out new values", async (userId) => {
+      await sut.rollback(helper);
+      expect(helper.setToUser).toHaveBeenCalledWith(userId, keyDefinitionLike, null);
+    });
+
+    it("should add collection values back to accounts", async () => {
+      await sut.rollback(helper);
+
+      expect(helper.set).toHaveBeenCalled();
+      expect(helper.set).toHaveBeenCalledWith("user-1", {
+        data: {
+          collections: {
+            encrypted: {
+              "877fef70-be32-439e-8678-b0d80125653d": {
+                id: "877fef70-be32-439e-8678-b0d80125653d",
+                organizationId: "fe1ff6ef-d2d4-49f3-9c07-b0c7013998f9",
+                name: "2.MD9OMDsvYiU1CTSUxjHorw==|uFc4cZhnmQmK2LFCWbyeZg==|syk2d9JESeplxInLvP36BK5RhqS1c/i+ZQp5NR7EUA4=",
+                externalId: "",
+                readOnly: false,
+                manage: true,
+                hidePasswords: false,
+              },
+              "0d3fee82-3f81-434c-aed0-b0c200ee6c7a": {
+                id: "0d3fee82-3f81-434c-aed0-b0c200ee6c7a",
+                organizationId: "5f277723-6391-4b5c-add9-b0c200ee6967",
+                name: "2.GxnXkIbBCGFr57F6lT7+Ow==|3ctMg95FKquG3l+qfv8BgvaCbYzMmuhnukCEHXhUukE=|cJRZWq05xjPBayUgx6P6gsbtNVLi8exQwo8F1SfqQQ4=",
+                externalId: "",
+                readOnly: false,
+                manage: false,
+                hidePasswords: false,
+              },
+            },
+          },
+          otherStuff: "otherStuff2",
+        },
+        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());
+    });
+  });
+});
diff --git a/libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.ts b/libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.ts
new file mode 100644
index 00000000000..0192161f6a1
--- /dev/null
+++ b/libs/common/src/state-migrations/migrations/21-move-collections-state-to-state-provider.ts
@@ -0,0 +1,63 @@
+import { KeyDefinitionLike, MigrationHelper } from "../migration-helper";
+import { Migrator } from "../migrator";
+
+type CollectionDataType = {
+  id: string;
+  organizationId: string;
+  name: string;
+  externalId: string;
+  readOnly: boolean;
+  manage: boolean;
+  hidePasswords: boolean;
+};
+
+type ExpectedAccountType = {
+  data?: {
+    collections?: {
+      encrypted?: Record<string, CollectionDataType>;
+    };
+  };
+};
+
+const USER_ENCRYPTED_COLLECTIONS: KeyDefinitionLike = {
+  key: "collections",
+  stateDefinition: {
+    name: "collection",
+  },
+};
+
+export class CollectionMigrator extends Migrator<20, 21> {
+  async migrate(helper: MigrationHelper): Promise<void> {
+    const accounts = await helper.getAccounts<ExpectedAccountType>();
+
+    async function migrateAccount(userId: string, account: ExpectedAccountType): Promise<void> {
+      const value = account?.data?.collections?.encrypted;
+      if (value != null) {
+        await helper.setToUser(userId, USER_ENCRYPTED_COLLECTIONS, value);
+        delete account.data.collections;
+        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(userId, USER_ENCRYPTED_COLLECTIONS);
+      if (account) {
+        account.data = Object.assign(account.data ?? {}, {
+          collections: {
+            encrypted: value,
+          },
+        });
+
+        await helper.set(userId, account);
+      }
+      await helper.setToUser(userId, USER_ENCRYPTED_COLLECTIONS, null);
+    }
+    await Promise.all([...accounts.map(({ userId, account }) => rollbackAccount(userId, account))]);
+  }
+}
diff --git a/libs/common/src/types/guid.ts b/libs/common/src/types/guid.ts
index 7f208999081..50ef5d0a8c5 100644
--- a/libs/common/src/types/guid.ts
+++ b/libs/common/src/types/guid.ts
@@ -4,4 +4,5 @@ export type Guid = Opaque<string, "Guid">;
 
 export type UserId = Opaque<string, "UserId">;
 export type OrganizationId = Opaque<string, "OrganizationId">;
+export type CollectionId = Opaque<string, "CollectionId">;
 export type ProviderId = Opaque<string, "ProviderId">;
diff --git a/libs/common/src/vault/abstractions/collection.service.ts b/libs/common/src/vault/abstractions/collection.service.ts
index 316794a5562..ab606153292 100644
--- a/libs/common/src/vault/abstractions/collection.service.ts
+++ b/libs/common/src/vault/abstractions/collection.service.ts
@@ -1,11 +1,18 @@
+import { Observable } from "rxjs";
+
+import { CollectionId } from "../../types/guid";
 import { CollectionData } from "../models/data/collection.data";
 import { Collection } from "../models/domain/collection";
 import { TreeNode } from "../models/domain/tree-node";
 import { CollectionView } from "../models/view/collection.view";
 
 export abstract class CollectionService {
-  clearCache: (userId?: string) => Promise<void>;
+  clearActiveUserCache: () => Promise<void>;
   encrypt: (model: CollectionView) => Promise<Collection>;
+  decryptedCollectionViews$: (ids: CollectionId[]) => Observable<CollectionView[]>;
+  /**
+   * @deprecated This method will soon be made private, use `decryptedCollectionViews$` instead.
+   */
   decryptMany: (collections: Collection[]) => Promise<CollectionView[]>;
   get: (id: string) => Promise<Collection>;
   getAll: () => Promise<Collection[]>;
diff --git a/libs/common/src/vault/models/data/collection.data.ts b/libs/common/src/vault/models/data/collection.data.ts
index 512bb4f8266..6f575a9dd81 100644
--- a/libs/common/src/vault/models/data/collection.data.ts
+++ b/libs/common/src/vault/models/data/collection.data.ts
@@ -1,8 +1,11 @@
+import { Jsonify } from "type-fest";
+
+import { CollectionId, OrganizationId } from "../../../types/guid";
 import { CollectionDetailsResponse } from "../response/collection.response";
 
 export class CollectionData {
-  id: string;
-  organizationId: string;
+  id: CollectionId;
+  organizationId: OrganizationId;
   name: string;
   externalId: string;
   readOnly: boolean;
@@ -18,4 +21,8 @@ export class CollectionData {
     this.manage = response.manage;
     this.hidePasswords = response.hidePasswords;
   }
+
+  static fromJSON(obj: Jsonify<CollectionData>) {
+    return Object.assign(new CollectionData(new CollectionDetailsResponse({})), obj);
+  }
 }
diff --git a/libs/common/src/vault/models/domain/collection.spec.ts b/libs/common/src/vault/models/domain/collection.spec.ts
index 99fa39384b4..848633e501d 100644
--- a/libs/common/src/vault/models/domain/collection.spec.ts
+++ b/libs/common/src/vault/models/domain/collection.spec.ts
@@ -1,4 +1,5 @@
 import { mockEnc } from "../../../../spec";
+import { CollectionId, OrganizationId } from "../../../types/guid";
 import { CollectionData } from "../data/collection.data";
 
 import { Collection } from "./collection";
@@ -8,8 +9,8 @@ describe("Collection", () => {
 
   beforeEach(() => {
     data = {
-      id: "id",
-      organizationId: "orgId",
+      id: "id" as CollectionId,
+      organizationId: "orgId" as OrganizationId,
       name: "encName",
       externalId: "extId",
       readOnly: true,
diff --git a/libs/common/src/vault/models/response/collection.response.ts b/libs/common/src/vault/models/response/collection.response.ts
index 6a97b78238a..f16fe547e07 100644
--- a/libs/common/src/vault/models/response/collection.response.ts
+++ b/libs/common/src/vault/models/response/collection.response.ts
@@ -1,9 +1,10 @@
 import { SelectionReadOnlyResponse } from "../../../admin-console/models/response/selection-read-only.response";
 import { BaseResponse } from "../../../models/response/base.response";
+import { CollectionId, OrganizationId } from "../../../types/guid";
 
 export class CollectionResponse extends BaseResponse {
-  id: string;
-  organizationId: string;
+  id: CollectionId;
+  organizationId: OrganizationId;
   name: string;
   externalId: string;
 
diff --git a/libs/common/src/vault/models/view/collection.view.ts b/libs/common/src/vault/models/view/collection.view.ts
index 3ff8b7ad705..1177d23220b 100644
--- a/libs/common/src/vault/models/view/collection.view.ts
+++ b/libs/common/src/vault/models/view/collection.view.ts
@@ -1,3 +1,5 @@
+import { Jsonify } from "type-fest";
+
 import { Organization } from "../../../admin-console/models/domain/organization";
 import { View } from "../../../models/view/view";
 import { Collection } from "../domain/collection";
@@ -56,4 +58,8 @@ export class CollectionView implements View, ITreeNodeObject {
       ? org?.canDeleteAnyCollection || (!org?.limitCollectionCreationDeletion && this.manage)
       : org?.canDeleteAnyCollection || org?.canDeleteAssignedCollections;
   }
+
+  static fromJSON(obj: Jsonify<CollectionView>) {
+    return Object.assign(new CollectionView(new Collection()), obj);
+  }
 }
diff --git a/libs/common/src/vault/services/collection.service.ts b/libs/common/src/vault/services/collection.service.ts
index ed051f05aaa..897fb4a772d 100644
--- a/libs/common/src/vault/services/collection.service.ts
+++ b/libs/common/src/vault/services/collection.service.ts
@@ -1,7 +1,18 @@
+import { firstValueFrom, map, Observable } from "rxjs";
+import { Jsonify } from "type-fest";
+
 import { CryptoService } from "../../platform/abstractions/crypto.service";
 import { I18nService } from "../../platform/abstractions/i18n.service";
-import { StateService } from "../../platform/abstractions/state.service";
 import { Utils } from "../../platform/misc/utils";
+import {
+  ActiveUserState,
+  KeyDefinition,
+  StateProvider,
+  COLLECTION_DATA,
+  DeriveDefinition,
+  DerivedState,
+} from "../../platform/state";
+import { CollectionId, UserId } from "../../types/guid";
 import { CollectionService as CollectionServiceAbstraction } from "../../vault/abstractions/collection.service";
 import { CollectionData } from "../models/data/collection.data";
 import { Collection } from "../models/domain/collection";
@@ -9,17 +20,71 @@ import { TreeNode } from "../models/domain/tree-node";
 import { CollectionView } from "../models/view/collection.view";
 import { ServiceUtils } from "../service-utils";
 
+const ENCRYPTED_COLLECTION_DATA_KEY = KeyDefinition.record<CollectionData, CollectionId>(
+  COLLECTION_DATA,
+  "collections",
+  {
+    deserializer: (jsonData: Jsonify<CollectionData>) => CollectionData.fromJSON(jsonData),
+  },
+);
+
+const DECRYPTED_COLLECTION_DATA_KEY = DeriveDefinition.from<
+  Record<CollectionId, CollectionData>,
+  CollectionView[],
+  { collectionService: CollectionService }
+>(ENCRYPTED_COLLECTION_DATA_KEY, {
+  deserializer: (obj) => obj.map((collection) => CollectionView.fromJSON(collection)),
+  derive: async (collections: Record<CollectionId, CollectionData>, { collectionService }) => {
+    const data: Collection[] = [];
+    for (const id in collections ?? {}) {
+      const collectionId = id as CollectionId;
+      data.push(new Collection(collections[collectionId]));
+    }
+    return await collectionService.decryptMany(data);
+  },
+});
+
 const NestingDelimiter = "/";
 
 export class CollectionService implements CollectionServiceAbstraction {
+  private encryptedCollectionDataState: ActiveUserState<Record<CollectionId, CollectionData>>;
+  encryptedCollections$: Observable<Collection[]>;
+  private decryptedCollectionDataState: DerivedState<CollectionView[]>;
+  decryptedCollections$: Observable<CollectionView[]>;
+
+  decryptedCollectionViews$(ids: CollectionId[]): Observable<CollectionView[]> {
+    return this.decryptedCollections$.pipe(
+      map((collections) => collections.filter((c) => ids.includes(c.id as CollectionId))),
+    );
+  }
+
   constructor(
     private cryptoService: CryptoService,
     private i18nService: I18nService,
-    private stateService: StateService,
-  ) {}
+    protected stateProvider: StateProvider,
+  ) {
+    this.encryptedCollectionDataState = this.stateProvider.getActive(ENCRYPTED_COLLECTION_DATA_KEY);
+    this.encryptedCollections$ = this.encryptedCollectionDataState.state$.pipe(
+      map((collections) => {
+        const response: Collection[] = [];
+        for (const id in collections ?? {}) {
+          response.push(new Collection(collections[id as CollectionId]));
+        }
+        return response;
+      }),
+    );
 
-  async clearCache(userId?: string): Promise<void> {
-    await this.stateService.setDecryptedCollections(null, { userId: userId });
+    this.decryptedCollectionDataState = this.stateProvider.getDerived(
+      this.encryptedCollectionDataState.state$,
+      DECRYPTED_COLLECTION_DATA_KEY,
+      { collectionService: this },
+    );
+
+    this.decryptedCollections$ = this.decryptedCollectionDataState.state$;
+  }
+
+  async clearActiveUserCache(): Promise<void> {
+    await this.decryptedCollectionDataState.forceValue(null);
   }
 
   async encrypt(model: CollectionView): Promise<Collection> {
@@ -52,43 +117,19 @@ export class CollectionService implements CollectionServiceAbstraction {
   }
 
   async get(id: string): Promise<Collection> {
-    const collections = await this.stateService.getEncryptedCollections();
-    // eslint-disable-next-line
-    if (collections == null || !collections.hasOwnProperty(id)) {
-      return null;
-    }
-
-    return new Collection(collections[id]);
+    return (
+      (await firstValueFrom(
+        this.encryptedCollections$.pipe(map((cs) => cs.find((c) => c.id === id))),
+      )) ?? null
+    );
   }
 
   async getAll(): Promise<Collection[]> {
-    const collections = await this.stateService.getEncryptedCollections();
-    const response: Collection[] = [];
-    for (const id in collections) {
-      // eslint-disable-next-line
-      if (collections.hasOwnProperty(id)) {
-        response.push(new Collection(collections[id]));
-      }
-    }
-    return response;
+    return await firstValueFrom(this.encryptedCollections$);
   }
 
   async getAllDecrypted(): Promise<CollectionView[]> {
-    let decryptedCollections = await this.stateService.getDecryptedCollections();
-    if (decryptedCollections != null) {
-      return decryptedCollections;
-    }
-
-    const hasKey = await this.cryptoService.hasUserKey();
-    if (!hasKey) {
-      throw new Error("No key.");
-    }
-
-    const collections = await this.getAll();
-    decryptedCollections = await this.decryptMany(collections);
-
-    await this.stateService.setDecryptedCollections(decryptedCollections);
-    return decryptedCollections;
+    return await firstValueFrom(this.decryptedCollections$);
   }
 
   async getAllNested(collections: CollectionView[] = null): Promise<TreeNode<CollectionView>[]> {
@@ -115,48 +156,51 @@ export class CollectionService implements CollectionServiceAbstraction {
     return ServiceUtils.getTreeNodeObjectFromList(collections, id) as TreeNode<CollectionView>;
   }
 
-  async upsert(collection: CollectionData | CollectionData[]): Promise<any> {
-    let collections = await this.stateService.getEncryptedCollections();
-    if (collections == null) {
-      collections = {};
-    }
-
-    if (collection instanceof CollectionData) {
-      const c = collection as CollectionData;
-      collections[c.id] = c;
-    } else {
-      (collection as CollectionData[]).forEach((c) => {
-        collections[c.id] = c;
-      });
-    }
-
-    await this.replace(collections);
-  }
-
-  async replace(collections: { [id: string]: CollectionData }): Promise<any> {
-    await this.clearCache();
-    await this.stateService.setEncryptedCollections(collections);
-  }
-
-  async clear(userId?: string): Promise<any> {
-    await this.clearCache(userId);
-    await this.stateService.setEncryptedCollections(null, { userId: userId });
-  }
-
-  async delete(id: string | string[]): Promise<any> {
-    const collections = await this.stateService.getEncryptedCollections();
-    if (collections == null) {
+  async upsert(toUpdate: CollectionData | CollectionData[]): Promise<void> {
+    if (toUpdate == null) {
       return;
     }
+    await this.encryptedCollectionDataState.update((collections) => {
+      if (collections == null) {
+        collections = {};
+      }
+      if (Array.isArray(toUpdate)) {
+        toUpdate.forEach((c) => {
+          collections[c.id] = c;
+        });
+      } else {
+        collections[toUpdate.id] = toUpdate;
+      }
+      return collections;
+    });
+  }
 
-    if (typeof id === "string") {
-      delete collections[id];
+  async replace(collections: Record<CollectionId, CollectionData>): Promise<void> {
+    await this.encryptedCollectionDataState.update(() => collections);
+  }
+
+  async clear(userId?: UserId): Promise<any> {
+    if (userId == null) {
+      await this.encryptedCollectionDataState.update(() => null);
+      await this.decryptedCollectionDataState.forceValue(null);
     } else {
-      (id as string[]).forEach((i) => {
-        delete collections[i];
-      });
+      await this.stateProvider.getUser(userId, ENCRYPTED_COLLECTION_DATA_KEY).update(() => null);
     }
+  }
 
-    await this.replace(collections);
+  async delete(id: CollectionId | CollectionId[]): Promise<any> {
+    await this.encryptedCollectionDataState.update((collections) => {
+      if (collections == null) {
+        collections = {};
+      }
+      if (typeof id === "string") {
+        delete collections[id];
+      } else {
+        (id as CollectionId[]).forEach((i) => {
+          delete collections[i];
+        });
+      }
+      return collections;
+    });
   }
 }