From 6c88677d70e7e2fec1e02c92b05d45c73c0fca48 Mon Sep 17 00:00:00 2001
From: Matt Gibson <mgibson@bitwarden.com>
Date: Thu, 5 Oct 2023 11:41:21 -0400
Subject: [PATCH] Use KeyDerivations in DerivedState

---
 .../src/platform/state/derived-state-definition.ts   |  4 ++--
 libs/common/src/platform/state/key-definition.ts     | 12 ++++++------
 libs/common/src/platform/state/key-definitions.ts    |  8 +++++++-
 libs/common/src/platform/state/state-definitions.ts  |  1 +
 .../src/vault/services/folder/folder.service.ts      |  4 ++--
 5 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/libs/common/src/platform/state/derived-state-definition.ts b/libs/common/src/platform/state/derived-state-definition.ts
index bed8201192..0c85b03878 100644
--- a/libs/common/src/platform/state/derived-state-definition.ts
+++ b/libs/common/src/platform/state/derived-state-definition.ts
@@ -1,7 +1,7 @@
 import { EncryptService } from "../abstractions/encrypt.service";
 import { UserKey } from "../models/domain/symmetric-crypto-key";
 
-import { StorageLocation } from "./state-definition";
+import { KeyDefinition } from "./key-definition";
 
 // TODO: Move type
 export class DeriveContext {
@@ -10,7 +10,7 @@ export class DeriveContext {
 
 export class DerivedStateDefinition<TFrom, TTo> {
   constructor(
-    readonly location: StorageLocation,
+    readonly KeyDefinition: KeyDefinition<TTo>,
     readonly converter: (data: TFrom, context: DeriveContext) => Promise<TTo>
   ) {}
 }
diff --git a/libs/common/src/platform/state/key-definition.ts b/libs/common/src/platform/state/key-definition.ts
index 64ccaf095a..310cf3c644 100644
--- a/libs/common/src/platform/state/key-definition.ts
+++ b/libs/common/src/platform/state/key-definition.ts
@@ -1,7 +1,7 @@
 import { Jsonify } from "type-fest";
 
 import { DeriveContext, DerivedStateDefinition } from "./derived-state-definition";
-import { StateDefinition, StorageLocation } from "./state-definition";
+import { StateDefinition } from "./state-definition";
 
 /**
  *
@@ -47,14 +47,14 @@ export class KeyDefinition<T> {
 
   /**
    * Helper for defining a derived definition that will often be used alongside a given key
-   * @param storageLocation
-   * @param decrypt
+   * @param keyDefinition The key definition detailing storage details for the derived state
+   * @param deriveCallback The callback used to convert from the parent state to the derived state
    * @returns
    */
   createDerivedDefinition<TTo>(
-    storageLocation: StorageLocation,
-    decrypt: (data: T, context: DeriveContext) => Promise<TTo>
+    keyDefinition: KeyDefinition<TTo>,
+    deriveCallback: (data: T, context: DeriveContext) => Promise<TTo>
   ) {
-    return new DerivedStateDefinition(storageLocation, decrypt);
+    return new DerivedStateDefinition(keyDefinition, deriveCallback);
   }
 }
diff --git a/libs/common/src/platform/state/key-definitions.ts b/libs/common/src/platform/state/key-definitions.ts
index be97432943..d522f1718f 100644
--- a/libs/common/src/platform/state/key-definitions.ts
+++ b/libs/common/src/platform/state/key-definitions.ts
@@ -1,7 +1,8 @@
 import { FolderData } from "../../vault/models/data/folder.data";
+import { FolderView } from "../../vault/models/view/folder.view";
 
 import { KeyDefinition } from "./key-definition";
-import { FOLDER_SERVICE_DISK } from "./state-definitions";
+import { FOLDER_SERVICE_DISK, FOLDER_SERVICE_MEMORY } from "./state-definitions";
 
 // FolderService Keys
 export const FOLDERS = KeyDefinition.record<FolderData>(
@@ -9,3 +10,8 @@ export const FOLDERS = KeyDefinition.record<FolderData>(
   "folders",
   FolderData.fromJSON
 );
+export const FOLDER_VIEWS = KeyDefinition.array<FolderView>(
+  FOLDER_SERVICE_MEMORY,
+  "folderViews",
+  FolderView.fromJSON
+);
diff --git a/libs/common/src/platform/state/state-definitions.ts b/libs/common/src/platform/state/state-definitions.ts
index 3fda37149d..fbb5f4454f 100644
--- a/libs/common/src/platform/state/state-definitions.ts
+++ b/libs/common/src/platform/state/state-definitions.ts
@@ -1,3 +1,4 @@
 import { StateDefinition } from "./state-definition";
 
 export const FOLDER_SERVICE_DISK = new StateDefinition("FolderService", "disk");
+export const FOLDER_SERVICE_MEMORY = new StateDefinition("FolderService", "memory");
diff --git a/libs/common/src/vault/services/folder/folder.service.ts b/libs/common/src/vault/services/folder/folder.service.ts
index dae1f4b488..853dbdfbf1 100644
--- a/libs/common/src/vault/services/folder/folder.service.ts
+++ b/libs/common/src/vault/services/folder/folder.service.ts
@@ -8,7 +8,7 @@ import { UserState } from "../../../platform/interfaces/user-state";
 import { Utils } from "../../../platform/misc/utils";
 import { SymmetricCryptoKey } from "../../../platform/models/domain/symmetric-crypto-key";
 import { DerivedUserState } from "../../../platform/state";
-import { FOLDERS } from "../../../platform/state/key-definitions";
+import { FOLDERS, FOLDER_VIEWS } from "../../../platform/state/key-definitions";
 import { UserId } from "../../../types/guid";
 import { CipherService } from "../../../vault/abstractions/cipher.service";
 import { InternalFolderService as InternalFolderServiceAbstraction } from "../../../vault/abstractions/folder/folder.service.abstraction";
@@ -34,7 +34,7 @@ export class FolderService implements InternalFolderServiceAbstraction {
     (window as any).services ||= {};
     (window as any).services.folderService = this;
     const derivedFoldersDefinition = FOLDERS.createDerivedDefinition(
-      "memory",
+      FOLDER_VIEWS,
       async (foldersMap) => {
         const folders = this.flattenMap(foldersMap);
         const decryptedFolders = await this.decryptFolders(folders);