LocalBackedSessionStorage Updates (#8542)

apps/browser/src/background/main.background.ts

@@ -207,6 +207,7 @@ import { BrowserStateService as StateServiceAbstraction } from "../platform/serv
 import { BrowserCryptoService } from "../platform/services/browser-crypto.service";
 import { BrowserEnvironmentService } from "../platform/services/browser-environment.service";
 import BrowserLocalStorageService from "../platform/services/browser-local-storage.service";
+import BrowserMemoryStorageService from "../platform/services/browser-memory-storage.service";
 import BrowserMessagingPrivateModeBackgroundService from "../platform/services/browser-messaging-private-mode-background.service";
 import BrowserMessagingService from "../platform/services/browser-messaging.service";
 import { BrowserStateService } from "../platform/services/browser-state.service";
@@ -230,7 +231,7 @@ import RuntimeBackground from "./runtime.background";
 
 export default class MainBackground {
   messagingService: MessagingServiceAbstraction;
-  storageService: AbstractStorageService;
+  storageService: AbstractStorageService & ObservableStorageService;
   secureStorageService: AbstractStorageService;
   memoryStorageService: AbstractMemoryStorageService;
   memoryStorageForStateProviders: AbstractMemoryStorageService & ObservableStorageService;
@@ -365,22 +366,28 @@ export default class MainBackground {
     this.cryptoFunctionService = new WebCryptoFunctionService(self);
     this.keyGenerationService = new KeyGenerationService(this.cryptoFunctionService);
     this.storageService = new BrowserLocalStorageService();
+
+    const mv3MemoryStorageCreator = (partitionName: string) => {
+      // TODO: Consider using multithreaded encrypt service in popup only context
+      return new LocalBackedSessionStorageService(
+        new EncryptServiceImplementation(this.cryptoFunctionService, this.logService, false),
+        this.keyGenerationService,
+        new BrowserLocalStorageService(),
+        new BrowserMemoryStorageService(),
+        partitionName,
+      );
+    };
+
     this.secureStorageService = this.storageService; // secure storage is not supported in browsers, so we use local storage and warn users when it is used
     this.memoryStorageService = BrowserApi.isManifestVersion(3)
-      ? new LocalBackedSessionStorageService(
+      ? mv3MemoryStorageCreator("stateService")
-          new EncryptServiceImplementation(this.cryptoFunctionService, this.logService, false),
-          this.keyGenerationService,
-        )
       : new MemoryStorageService();
     this.memoryStorageForStateProviders = BrowserApi.isManifestVersion(3)
-      ? new LocalBackedSessionStorageService(
+      ? mv3MemoryStorageCreator("stateProviders")
-          new EncryptServiceImplementation(this.cryptoFunctionService, this.logService, false),
-          this.keyGenerationService,
-        )
       : new BackgroundMemoryStorageService();
 
     const storageServiceProvider = new StorageServiceProvider(
-      this.storageService as BrowserLocalStorageService,
+      this.storageService,
       this.memoryStorageForStateProviders,
     );
 

apps/browser/src/platform/background/service-factories/storage-service.factory.ts

@@ -7,6 +7,7 @@ import { MemoryStorageService } from "@bitwarden/common/platform/services/memory
 
 import { BrowserApi } from "../../browser/browser-api";
 import BrowserLocalStorageService from "../../services/browser-local-storage.service";
+import BrowserMemoryStorageService from "../../services/browser-memory-storage.service";
 import { LocalBackedSessionStorageService } from "../../services/local-backed-session-storage.service";
 import { BackgroundMemoryStorageService } from "../../storage/background-memory-storage.service";
 
@@ -17,13 +18,14 @@ import {
   keyGenerationServiceFactory,
 } from "./key-generation-service.factory";
 
-type StorageServiceFactoryOptions = FactoryOptions;
+export type DiskStorageServiceInitOptions = FactoryOptions;
-
+export type SecureStorageServiceInitOptions = FactoryOptions;
-export type DiskStorageServiceInitOptions = StorageServiceFactoryOptions;
+export type SessionStorageServiceInitOptions = FactoryOptions;
-export type SecureStorageServiceInitOptions = StorageServiceFactoryOptions;
+export type MemoryStorageServiceInitOptions = FactoryOptions &
-export type MemoryStorageServiceInitOptions = StorageServiceFactoryOptions &
   EncryptServiceInitOptions &
-  KeyGenerationServiceInitOptions;
+  KeyGenerationServiceInitOptions &
+  DiskStorageServiceInitOptions &
+  SessionStorageServiceInitOptions;
 
 export function diskStorageServiceFactory(
   cache: { diskStorageService?: AbstractStorageService } & CachedServices,
@@ -47,6 +49,13 @@ export function secureStorageServiceFactory(
   return factory(cache, "secureStorageService", opts, () => new BrowserLocalStorageService());
 }
 
+export function sessionStorageServiceFactory(
+  cache: { sessionStorageService?: AbstractStorageService } & CachedServices,
+  opts: SessionStorageServiceInitOptions,
+): Promise<AbstractStorageService> {
+  return factory(cache, "sessionStorageService", opts, () => new BrowserMemoryStorageService());
+}
+
 export function memoryStorageServiceFactory(
   cache: { memoryStorageService?: AbstractMemoryStorageService } & CachedServices,
   opts: MemoryStorageServiceInitOptions,
@@ -56,6 +65,9 @@ export function memoryStorageServiceFactory(
       return new LocalBackedSessionStorageService(
         await encryptServiceFactory(cache, opts),
         await keyGenerationServiceFactory(cache, opts),
+        await diskStorageServiceFactory(cache, opts),
+        await sessionStorageServiceFactory(cache, opts),
+        "serviceFactories",
       );
     }
     return new MemoryStorageService();

apps/browser/src/platform/services/local-backed-session-storage.service.spec.ts

@@ -2,45 +2,70 @@ import { mock, MockProxy } from "jest-mock-extended";
 
 import { EncryptService } from "@bitwarden/common/platform/abstractions/encrypt.service";
 import { KeyGenerationService } from "@bitwarden/common/platform/abstractions/key-generation.service";
+import {
+  AbstractMemoryStorageService,
+  AbstractStorageService,
+  StorageUpdate,
+} from "@bitwarden/common/platform/abstractions/storage.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
 import { EncString } from "@bitwarden/common/platform/models/domain/enc-string";
 import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
 
-import BrowserLocalStorageService from "./browser-local-storage.service";
-import BrowserMemoryStorageService from "./browser-memory-storage.service";
 import { LocalBackedSessionStorageService } from "./local-backed-session-storage.service";
 
-describe("Browser Session Storage Service", () => {
+describe("LocalBackedSessionStorage", () => {
   let encryptService: MockProxy<EncryptService>;
   let keyGenerationService: MockProxy<KeyGenerationService>;
+  let localStorageService: MockProxy<AbstractStorageService>;
+  let sessionStorageService: MockProxy<AbstractMemoryStorageService>;
 
   let cache: Map<string, any>;
   const testObj = { a: 1, b: 2 };
 
-  let localStorage: BrowserLocalStorageService;
-  let sessionStorage: BrowserMemoryStorageService;
-
   const key = new SymmetricCryptoKey(Utils.fromUtf8ToArray("00000000000000000000000000000000"));
   let getSessionKeySpy: jest.SpyInstance;
+  let sendUpdateSpy: jest.SpyInstance<void, [storageUpdate: StorageUpdate]>;
   const mockEnc = (input: string) => Promise.resolve(new EncString("ENCRYPTED" + input));
 
   let sut: LocalBackedSessionStorageService;
 
+  const mockExistingSessionKey = (key: SymmetricCryptoKey) => {
+    sessionStorageService.get.mockImplementation((storageKey) => {
+      if (storageKey === "localEncryptionKey_test") {
+        return Promise.resolve(key?.toJSON());
+      }
+
+      return Promise.reject("No implementation for " + storageKey);
+    });
+  };
+
   beforeEach(() => {
     encryptService = mock<EncryptService>();
     keyGenerationService = mock<KeyGenerationService>();
+    localStorageService = mock<AbstractStorageService>();
+    sessionStorageService = mock<AbstractMemoryStorageService>();
 
-    sut = new LocalBackedSessionStorageService(encryptService, keyGenerationService);
+    sut = new LocalBackedSessionStorageService(
+      encryptService,
+      keyGenerationService,
+      localStorageService,
+      sessionStorageService,
+      "test",
+    );
 
     cache = sut["cache"];
-    localStorage = sut["localStorage"];
+
-    sessionStorage = sut["sessionStorage"];
+    keyGenerationService.createKeyWithPurpose.mockResolvedValue({
+      derivedKey: key,
+      salt: "bitwarden-ephemeral",
+      material: null, // Not used
+    });
+
     getSessionKeySpy = jest.spyOn(sut, "getSessionEncKey");
     getSessionKeySpy.mockResolvedValue(key);
-  });
 
-  it("should exist", () => {
+    sendUpdateSpy = jest.spyOn(sut, "sendUpdate");
-    expect(sut).toBeInstanceOf(LocalBackedSessionStorageService);
+    sendUpdateSpy.mockReturnValue();
   });
 
   describe("get", () => {
@@ -54,7 +79,7 @@ describe("Browser Session Storage Service", () => {
       const session = { test: testObj };
 
       beforeEach(() => {
-        jest.spyOn(sut, "getSessionEncKey").mockResolvedValue(key);
+        mockExistingSessionKey(key);
       });
 
       describe("no session retrieved", () => {
@@ -62,6 +87,7 @@ describe("Browser Session Storage Service", () => {
         let spy: jest.SpyInstance;
         beforeEach(async () => {
           spy = jest.spyOn(sut, "getLocalSession").mockResolvedValue(null);
+          localStorageService.get.mockResolvedValue(null);
           result = await sut.get("test");
         });
 
@@ -123,31 +149,31 @@ describe("Browser Session Storage Service", () => {
 
   describe("remove", () => {
     it("should save null", async () => {
-      const spy = jest.spyOn(sut, "save");
-      spy.mockResolvedValue(null);
       await sut.remove("test");
-      expect(spy).toHaveBeenCalledWith("test", null);
+      expect(sendUpdateSpy).toHaveBeenCalledWith({ key: "test", updateType: "remove" });
     });
   });
 
   describe("save", () => {
     describe("caching", () => {
       beforeEach(() => {
-        jest.spyOn(localStorage, "get").mockResolvedValue(null);
+        localStorageService.get.mockResolvedValue(null);
-        jest.spyOn(sessionStorage, "get").mockResolvedValue(null);
+        sessionStorageService.get.mockResolvedValue(null);
-        jest.spyOn(localStorage, "save").mockResolvedValue();
+
-        jest.spyOn(sessionStorage, "save").mockResolvedValue();
+        localStorageService.save.mockResolvedValue();
+        sessionStorageService.save.mockResolvedValue();
 
         encryptService.encrypt.mockResolvedValue(mockEnc("{}"));
       });
 
       it("should remove key from cache if value is null", async () => {
         cache.set("test", {});
-        const deleteSpy = jest.spyOn(cache, "delete");
+        const cacheSetSpy = jest.spyOn(cache, "set");
         expect(cache.has("test")).toBe(true);
         await sut.save("test", null);
-        expect(cache.has("test")).toBe(false);
+        // Don't remove from cache, just replace with null
-        expect(deleteSpy).toHaveBeenCalledWith("test");
+        expect(cache.get("test")).toBe(null);
+        expect(cacheSetSpy).toHaveBeenCalledWith("test", null);
       });
 
       it("should set cache if value is non-null", async () => {
@@ -197,7 +223,7 @@ describe("Browser Session Storage Service", () => {
     });
 
     it("should return the stored symmetric crypto key", async () => {
-      jest.spyOn(sessionStorage, "get").mockResolvedValue({ ...key });
+      sessionStorageService.get.mockResolvedValue({ ...key });
       const result = await sut.getSessionEncKey();
 
       expect(result).toStrictEqual(key);
@@ -205,7 +231,6 @@ describe("Browser Session Storage Service", () => {
 
     describe("new key creation", () => {
       beforeEach(() => {
-        jest.spyOn(sessionStorage, "get").mockResolvedValue(null);
         keyGenerationService.createKeyWithPurpose.mockResolvedValue({
           salt: "salt",
           material: null,
@@ -218,25 +243,24 @@ describe("Browser Session Storage Service", () => {
         const result = await sut.getSessionEncKey();
 
         expect(result).toStrictEqual(key);
-        expect(keyGenerationService.createKeyWithPurpose).toBeCalledTimes(1);
+        expect(keyGenerationService.createKeyWithPurpose).toHaveBeenCalledTimes(1);
       });
 
       it("should store a symmetric crypto key if it makes one", async () => {
         const spy = jest.spyOn(sut, "setSessionEncKey").mockResolvedValue();
         await sut.getSessionEncKey();
 
-        expect(spy).toBeCalledWith(key);
+        expect(spy).toHaveBeenCalledWith(key);
       });
     });
   });
 
   describe("getLocalSession", () => {
     it("should return null if session is null", async () => {
-      const spy = jest.spyOn(localStorage, "get").mockResolvedValue(null);
       const result = await sut.getLocalSession(key);
 
       expect(result).toBeNull();
-      expect(spy).toBeCalledWith("session");
+      expect(localStorageService.get).toHaveBeenCalledWith("session_test");
     });
 
     describe("non-null sessions", () => {
@@ -245,7 +269,7 @@ describe("Browser Session Storage Service", () => {
       const decryptedSession = JSON.stringify(session);
 
       beforeEach(() => {
-        jest.spyOn(localStorage, "get").mockResolvedValue(encSession.encryptedString);
+        localStorageService.get.mockResolvedValue(encSession.encryptedString);
       });
 
       it("should decrypt returned sessions", async () => {
@@ -267,13 +291,12 @@ describe("Browser Session Storage Service", () => {
       it("should remove state if decryption fails", async () => {
         encryptService.decryptToUtf8.mockResolvedValue(null);
         const setSessionEncKeySpy = jest.spyOn(sut, "setSessionEncKey").mockResolvedValue();
-        const removeLocalSessionSpy = jest.spyOn(localStorage, "remove").mockResolvedValue();
 
         const result = await sut.getLocalSession(key);
 
         expect(result).toBeNull();
         expect(setSessionEncKeySpy).toHaveBeenCalledWith(null);
-        expect(removeLocalSessionSpy).toHaveBeenCalledWith("session");
+        expect(localStorageService.remove).toHaveBeenCalledWith("session_test");
       });
     });
   });
@@ -284,7 +307,7 @@ describe("Browser Session Storage Service", () => {
 
     it("should encrypt a stringified session", async () => {
       encryptService.encrypt.mockImplementation(mockEnc);
-      jest.spyOn(localStorage, "save").mockResolvedValue();
+      localStorageService.save.mockResolvedValue();
       await sut.setLocalSession(testSession, key);
 
       expect(encryptService.encrypt).toHaveBeenNthCalledWith(1, testJSON, key);
@@ -292,32 +315,31 @@ describe("Browser Session Storage Service", () => {
 
     it("should remove local session if null", async () => {
       encryptService.encrypt.mockResolvedValue(null);
-      const spy = jest.spyOn(localStorage, "remove").mockResolvedValue();
       await sut.setLocalSession(null, key);
 
-      expect(spy).toHaveBeenCalledWith("session");
+      expect(localStorageService.remove).toHaveBeenCalledWith("session_test");
     });
 
     it("should save encrypted string", async () => {
       encryptService.encrypt.mockImplementation(mockEnc);
-      const spy = jest.spyOn(localStorage, "save").mockResolvedValue();
       await sut.setLocalSession(testSession, key);
 
-      expect(spy).toHaveBeenCalledWith("session", (await mockEnc(testJSON)).encryptedString);
+      expect(localStorageService.save).toHaveBeenCalledWith(
+        "session_test",
+        (await mockEnc(testJSON)).encryptedString,
+      );
     });
   });
 
   describe("setSessionKey", () => {
     it("should remove if null", async () => {
-      const spy = jest.spyOn(sessionStorage, "remove").mockResolvedValue();
       await sut.setSessionEncKey(null);
-      expect(spy).toHaveBeenCalledWith("localEncryptionKey");
+      expect(sessionStorageService.remove).toHaveBeenCalledWith("localEncryptionKey_test");
     });
 
     it("should save key when not null", async () => {
-      const spy = jest.spyOn(sessionStorage, "save").mockResolvedValue();
       await sut.setSessionEncKey(key);
-      expect(spy).toHaveBeenCalledWith("localEncryptionKey", key);
+      expect(sessionStorageService.save).toHaveBeenCalledWith("localEncryptionKey_test", key);
     });
   });
 });

apps/browser/src/platform/services/local-backed-session-storage.service.ts

@@ -1,40 +1,60 @@
-import { Subject } from "rxjs";
+import { Observable, Subject, filter, map, merge, share, tap } from "rxjs";
 import { Jsonify } from "type-fest";
 
 import { EncryptService } from "@bitwarden/common/platform/abstractions/encrypt.service";
 import { KeyGenerationService } from "@bitwarden/common/platform/abstractions/key-generation.service";
 import {
   AbstractMemoryStorageService,
+  AbstractStorageService,
+  ObservableStorageService,
   StorageUpdate,
 } from "@bitwarden/common/platform/abstractions/storage.service";
 import { EncString } from "@bitwarden/common/platform/models/domain/enc-string";
 import { MemoryStorageOptions } from "@bitwarden/common/platform/models/domain/storage-options";
 import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
 
+import { fromChromeEvent } from "../browser/from-chrome-event";
 import { devFlag } from "../decorators/dev-flag.decorator";
 import { devFlagEnabled } from "../flags";
 
-import BrowserLocalStorageService from "./browser-local-storage.service";
+export class LocalBackedSessionStorageService
-import BrowserMemoryStorageService from "./browser-memory-storage.service";
+  extends AbstractMemoryStorageService
-
+  implements ObservableStorageService
-const keys = {
+{
-  encKey: "localEncryptionKey",
-  sessionKey: "session",
-};
-
-export class LocalBackedSessionStorageService extends AbstractMemoryStorageService {
   private cache = new Map<string, unknown>();
-  private localStorage = new BrowserLocalStorageService();
-  private sessionStorage = new BrowserMemoryStorageService();
   private updatesSubject = new Subject<StorageUpdate>();
-  updates$;
+
+  private commandName = `localBackedSessionStorage_${this.name}`;
+  private encKey = `localEncryptionKey_${this.name}`;
+  private sessionKey = `session_${this.name}`;
+
+  updates$: Observable<StorageUpdate>;
 
   constructor(
     private encryptService: EncryptService,
     private keyGenerationService: KeyGenerationService,
+    private localStorage: AbstractStorageService,
+    private sessionStorage: AbstractStorageService,
+    private name: string,
   ) {
     super();
-    this.updates$ = this.updatesSubject.asObservable();
+
+    const remoteObservable = fromChromeEvent(chrome.runtime.onMessage).pipe(
+      filter(([msg]) => msg.command === this.commandName),
+      map(([msg]) => msg.update as StorageUpdate),
+      tap((update) => {
+        if (update.updateType === "remove") {
+          this.cache.set(update.key, null);
+        } else {
+          this.cache.delete(update.key);
+        }
+      }),
+      share(),
+    );
+
+    remoteObservable.subscribe();
+
+    this.updates$ = merge(this.updatesSubject.asObservable(), remoteObservable);
   }
 
   get valuesRequireDeserialization(): boolean {
@@ -70,23 +90,37 @@ export class LocalBackedSessionStorageService extends AbstractMemoryStorageServi
 
   async save<T>(key: string, obj: T): Promise<void> {
     if (obj == null) {
-      this.cache.delete(key);
+      return await this.remove(key);
-    } else {
-      this.cache.set(key, obj);
     }
 
+    this.cache.set(key, obj);
+    await this.updateLocalSessionValue(key, obj);
+    this.sendUpdate({ key, updateType: "save" });
+  }
+
+  async remove(key: string): Promise<void> {
+    this.cache.set(key, null);
+    await this.updateLocalSessionValue(key, null);
+    this.sendUpdate({ key, updateType: "remove" });
+  }
+
+  sendUpdate(storageUpdate: StorageUpdate) {
+    this.updatesSubject.next(storageUpdate);
+    void chrome.runtime.sendMessage({
+      command: this.commandName,
+      update: storageUpdate,
+    });
+  }
+
+  private async updateLocalSessionValue<T>(key: string, obj: T) {
     const sessionEncKey = await this.getSessionEncKey();
     const localSession = (await this.getLocalSession(sessionEncKey)) ?? {};
     localSession[key] = obj;
     await this.setLocalSession(localSession, sessionEncKey);
   }
 
-  async remove(key: string): Promise<void> {
-    await this.save(key, null);
-  }
-
   async getLocalSession(encKey: SymmetricCryptoKey): Promise<Record<string, unknown>> {
-    const local = await this.localStorage.get<string>(keys.sessionKey);
+    const local = await this.localStorage.get<string>(this.sessionKey);
 
     if (local == null) {
       return null;
@@ -100,7 +134,7 @@ export class LocalBackedSessionStorageService extends AbstractMemoryStorageServi
     if (sessionJson == null) {
       // Error with decryption -- session is lost, delete state and key and start over
       await this.setSessionEncKey(null);
-      await this.localStorage.remove(keys.sessionKey);
+      await this.localStorage.remove(this.sessionKey);
       return null;
     }
     return JSON.parse(sessionJson);
@@ -119,9 +153,9 @@ export class LocalBackedSessionStorageService extends AbstractMemoryStorageServi
     // Make sure we're storing the jsonified version of the session
     const jsonSession = JSON.parse(JSON.stringify(session));
     if (session == null) {
-      await this.localStorage.remove(keys.sessionKey);
+      await this.localStorage.remove(this.sessionKey);
     } else {
-      await this.localStorage.save(keys.sessionKey, jsonSession);
+      await this.localStorage.save(this.sessionKey, jsonSession);
     }
   }
 
@@ -130,13 +164,13 @@ export class LocalBackedSessionStorageService extends AbstractMemoryStorageServi
     const encSession = await this.encryptService.encrypt(jsonSession, key);
 
     if (encSession == null) {
-      return await this.localStorage.remove(keys.sessionKey);
+      return await this.localStorage.remove(this.sessionKey);
     }
-    await this.localStorage.save(keys.sessionKey, encSession.encryptedString);
+    await this.localStorage.save(this.sessionKey, encSession.encryptedString);
   }
 
   async getSessionEncKey(): Promise<SymmetricCryptoKey> {
-    let storedKey = await this.sessionStorage.get<SymmetricCryptoKey>(keys.encKey);
+    let storedKey = await this.sessionStorage.get<SymmetricCryptoKey>(this.encKey);
     if (storedKey == null || Object.keys(storedKey).length == 0) {
       const generatedKey = await this.keyGenerationService.createKeyWithPurpose(
         128,
@@ -153,9 +187,9 @@ export class LocalBackedSessionStorageService extends AbstractMemoryStorageServi
 
   async setSessionEncKey(input: SymmetricCryptoKey): Promise<void> {
     if (input == null) {
-      await this.sessionStorage.remove(keys.encKey);
+      await this.sessionStorage.remove(this.encKey);
     } else {
-      await this.sessionStorage.save(keys.encKey, input);
+      await this.sessionStorage.save(this.encKey, input);
     }
   }
 }