Separate StateService in background/visualizations

Visualizer state services share storages with background page, which
nicely emulates mv3 synchronization through session/local storage. There
should not be multithreading issues since all of these services are
still running through a single thread, we just now have multiple places
we are reading/writing data from.

Smaller improvements
* Rename browser's state service to BrowserStateService
* Remove unused WithPrototype decorator :celebrate:
* Removed conversion on withPrototypeForArrayMembers. It's reasonable to
think that if the type is maintained, it doesn't need conversion.

Eventually, we should be able to remove the withPrototypeForArrayMembers
decorator as well, but that will require a bit more work on
(de)serialization of the Accounts.data property.

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

@@ -93,6 +93,7 @@ import AutofillService from "../services/autofill.service";
 import { BrowserEnvironmentService } from "../services/browser-environment.service";
 import { BrowserFolderService } from "../services/browser-folder.service";
 import { BrowserPolicyService } from "../services/browser-policy.service";
+import { BrowserStateService } from "../services/browser-state.service";
 import { BrowserCryptoService } from "../services/browserCrypto.service";
 import BrowserLocalStorageService from "../services/browserLocalStorage.service";
 import BrowserMessagingService from "../services/browserMessaging.service";
@@ -101,7 +102,6 @@ import BrowserPlatformUtilsService from "../services/browserPlatformUtils.servic
 import I18nService from "../services/i18n.service";
 import { KeyGenerationService } from "../services/keyGeneration.service";
 import { LocalBackedSessionStorageService } from "../services/localBackedSessionStorage.service";
-import { StateService } from "../services/state.service";
 import { VaultFilterService } from "../services/vaultFilter.service";
 import VaultTimeoutService from "../services/vaultTimeout/vaultTimeout.service";
 
@@ -227,7 +227,7 @@ export default class MainBackground {
       this.secureStorageService,
       new StateFactory(GlobalState, Account)
     );
-    this.stateService = new StateService(
+    this.stateService = new BrowserStateService(
       this.storageService,
       this.secureStorageService,
       this.memoryStorageService,

apps/browser/src/background/service_factories/state-service.factory.ts

@@ -2,7 +2,7 @@ import { StateFactory } from "@bitwarden/common/factories/stateFactory";
 import { GlobalState } from "@bitwarden/common/models/domain/global-state";
 
 import { Account } from "../../models/account";
-import { StateService } from "../../services/state.service";
+import { BrowserStateService } from "../../services/browser-state.service";
 
 import { CachedServices, factory, FactoryOptions } from "./factory-options";
 import { logServiceFactory, LogServiceInitOptions } from "./log-service.factory";
@@ -34,15 +34,15 @@ export type StateServiceInitOptions = StateServiceFactoryOptions &
   StateMigrationServiceInitOptions;
 
 export async function stateServiceFactory(
-  cache: { stateService?: StateService } & CachedServices,
+  cache: { stateService?: BrowserStateService } & CachedServices,
   opts: StateServiceInitOptions
-): Promise<StateService> {
+): Promise<BrowserStateService> {
   const service = await factory(
     cache,
     "stateService",
     opts,
     async () =>
-      await new StateService(
+      await new BrowserStateService(
         await diskStorageServiceFactory(cache, opts),
         await secureStorageServiceFactory(cache, opts),
         await memoryStorageServiceFactory(cache, opts),

apps/browser/src/popup/services/services.module.ts

@@ -38,6 +38,7 @@ import { SearchService as SearchServiceAbstraction } from "@bitwarden/common/abs
 import { SendService } from "@bitwarden/common/abstractions/send.service";
 import { SettingsService } from "@bitwarden/common/abstractions/settings.service";
 import { StateService as BaseStateServiceAbstraction } from "@bitwarden/common/abstractions/state.service";
+import { StateMigrationService } from "@bitwarden/common/abstractions/stateMigration.service";
 import { AbstractStorageService } from "@bitwarden/common/abstractions/storage.service";
 import { SyncService } from "@bitwarden/common/abstractions/sync/sync.service.abstraction";
 import { TokenService } from "@bitwarden/common/abstractions/token.service";
@@ -47,6 +48,8 @@ import { UserVerificationService } from "@bitwarden/common/abstractions/userVeri
 import { UsernameGenerationService } from "@bitwarden/common/abstractions/usernameGeneration.service";
 import { VaultTimeoutService } from "@bitwarden/common/abstractions/vaultTimeout/vaultTimeout.service";
 import { VaultTimeoutSettingsService } from "@bitwarden/common/abstractions/vaultTimeout/vaultTimeoutSettings.service";
+import { StateFactory } from "@bitwarden/common/factories/stateFactory";
+import { GlobalState } from "@bitwarden/common/models/domain/global-state";
 import { AuthService } from "@bitwarden/common/services/auth.service";
 import { ConsoleLogService } from "@bitwarden/common/services/consoleLog.service";
 import { LoginService } from "@bitwarden/common/services/login.service";
@@ -54,10 +57,12 @@ import { SearchService } from "@bitwarden/common/services/search.service";
 
 import MainBackground from "../../background/main.background";
 import { BrowserApi } from "../../browser/browserApi";
+import { Account } from "../../models/account";
 import { AutofillService } from "../../services/abstractions/autofill.service";
 import { StateService as StateServiceAbstraction } from "../../services/abstractions/state.service";
 import { BrowserEnvironmentService } from "../../services/browser-environment.service";
 import { BrowserPolicyService } from "../../services/browser-policy.service";
+import { BrowserStateService } from "../../services/browser-state.service";
 import { BrowserFileDownloadService } from "../../services/browserFileDownloadService";
 import BrowserMessagingService from "../../services/browserMessaging.service";
 import BrowserMessagingPrivateModePopupService from "../../services/browserMessagingPrivateModePopup.service";
@@ -299,10 +304,36 @@ function getBgService<T>(service: keyof MainBackground) {
       useFactory: getBgService<AbstractStorageService>("memoryStorageService"),
     },
     {
-      provide: StateServiceAbstraction,
+      provide: StateMigrationService,
-      useFactory: getBgService<StateServiceAbstraction>("stateService"),
+      useFactory: getBgService<StateMigrationService>("stateMigrationService"),
       deps: [],
     },
+    {
+      provide: StateServiceAbstraction,
+      useFactory: (
+        storageService: AbstractStorageService,
+        secureStorageService: AbstractStorageService,
+        memoryStorageService: AbstractStorageService,
+        logService: LogServiceAbstraction,
+        stateMigrationService: StateMigrationService
+      ) => {
+        return new BrowserStateService(
+          storageService,
+          secureStorageService,
+          memoryStorageService,
+          logService,
+          stateMigrationService,
+          new StateFactory(GlobalState, Account)
+        );
+      },
+      deps: [
+        AbstractStorageService,
+        SECURE_STORAGE,
+        MEMORY_STORAGE,
+        LogServiceAbstraction,
+        StateMigrationService,
+      ],
+    },
     {
       provide: UsernameGenerationService,
       useFactory: getBgService<UsernameGenerationService>("usernameGenerationService"),
@@ -323,17 +354,19 @@ function getBgService<T>(service: keyof MainBackground) {
     },
     {
       provide: AbstractThemingService,
-      useFactory: () => {
+      useFactory: (
+        stateService: StateServiceAbstraction,
+        platformUtilsService: PlatformUtilsService
+      ) => {
         return new ThemingService(
-          getBgService<StateServiceAbstraction>("stateService")(),
+          stateService,
           // Safari doesn't properly handle the (prefers-color-scheme) media query in the popup window, it always returns light.
           // In Safari we have to use the background page instead, which comes with limitations like not dynamically changing the extension theme when the system theme is changed.
-          getBgService<PlatformUtilsService>("platformUtilsService")().isSafari()
+          platformUtilsService.isSafari() ? getBgService<Window>("backgroundWindow")() : window,
-            ? getBgService<Window>("backgroundWindow")()
-            : window,
           document
         );
       },
+      deps: [StateServiceAbstraction, PlatformUtilsService],
     },
   ],
 })

apps/browser/src/services/browser-state.service.spec.ts

@@ -1,6 +1,6 @@
-// eslint-disable-next-line no-restricted-imports
+import { mock, MockProxy } from "jest-mock-extended";
-import { Arg, Substitute, SubstituteOf } from "@fluffy-spoon/substitute";
 
+import { EncryptService } from "@bitwarden/common/abstractions/encrypt.service";
 import { LogService } from "@bitwarden/common/abstractions/log.service";
 import {
   MemoryStorageServiceInterface,
@@ -18,28 +18,29 @@ import { BrowserComponentState } from "../models/browserComponentState";
 import { BrowserGroupingsComponentState } from "../models/browserGroupingsComponentState";
 import { BrowserSendComponentState } from "../models/browserSendComponentState";
 
+import { AbstractKeyGenerationService } from "./abstractions/abstractKeyGeneration.service";
+import { BrowserStateService } from "./browser-state.service";
 import { LocalBackedSessionStorageService } from "./localBackedSessionStorage.service";
-import { StateService } from "./state.service";
 
 describe("Browser State Service", () => {
-  let secureStorageService: SubstituteOf<AbstractStorageService>;
+  let secureStorageService: MockProxy<AbstractStorageService>;
-  let diskStorageService: SubstituteOf<AbstractStorageService>;
+  let diskStorageService: MockProxy<AbstractStorageService>;
-  let logService: SubstituteOf<LogService>;
+  let logService: MockProxy<LogService>;
-  let stateMigrationService: SubstituteOf<StateMigrationService>;
+  let stateMigrationService: MockProxy<StateMigrationService>;
-  let stateFactory: SubstituteOf<StateFactory<GlobalState, Account>>;
+  let stateFactory: MockProxy<StateFactory<GlobalState, Account>>;
   let useAccountCache: boolean;
 
   let state: State<GlobalState, Account>;
   const userId = "userId";
 
-  let sut: StateService;
+  let sut: BrowserStateService;
 
   beforeEach(() => {
-    secureStorageService = Substitute.for();
+    secureStorageService = mock();
-    diskStorageService = Substitute.for();
+    diskStorageService = mock();
-    logService = Substitute.for();
+    logService = mock();
-    stateMigrationService = Substitute.for();
+    stateMigrationService = mock();
-    stateFactory = Substitute.for();
+    stateFactory = mock();
     useAccountCache = true;
 
     state = new State(new GlobalState());
@@ -54,9 +55,12 @@ describe("Browser State Service", () => {
 
     beforeEach(() => {
       // We need `AbstractCachedStorageService` in the prototype chain to correctly test cache bypass.
-      memoryStorageService = Object.create(LocalBackedSessionStorageService.prototype);
+      memoryStorageService = new LocalBackedSessionStorageService(
+        mock<EncryptService>(),
+        mock<AbstractKeyGenerationService>()
+      );
 
-      sut = new StateService(
+      sut = new BrowserStateService(
         diskStorageService,
         secureStorageService,
         memoryStorageService,
@@ -80,14 +84,14 @@ describe("Browser State Service", () => {
   });
 
   describe("state methods", () => {
-    let memoryStorageService: SubstituteOf<AbstractStorageService & MemoryStorageServiceInterface>;
+    let memoryStorageService: MockProxy<AbstractStorageService & MemoryStorageServiceInterface>;
 
     beforeEach(() => {
-      memoryStorageService = Substitute.for();
+      memoryStorageService = mock();
-      const stateGetter = (key: string) => Promise.resolve(JSON.parse(JSON.stringify(state)));
+      const stateGetter = (key: string) => Promise.resolve(state);
-      memoryStorageService.get("state", Arg.any()).mimicks(stateGetter);
+      memoryStorageService.get.mockImplementation(stateGetter);
 
-      sut = new StateService(
+      sut = new BrowserStateService(
         diskStorageService,
         secureStorageService,
         memoryStorageService,
@@ -128,6 +132,7 @@ describe("Browser State Service", () => {
           [SendType.Text, 5],
         ]);
         state.accounts[userId].send = sendState;
+        (global as any)["watch"] = state;
 
         const actual = await sut.getBrowserSendComponentState();
         expect(actual).toBeInstanceOf(BrowserSendComponentState);

apps/browser/src/services/browser-state.service.ts

@@ -1,13 +1,12 @@
+import { BehaviorSubject } from "rxjs";
 import { Jsonify } from "type-fest";
 
 import { AbstractCachedStorageService } from "@bitwarden/common/abstractions/storage.service";
 import { GlobalState } from "@bitwarden/common/models/domain/global-state";
 import { StorageOptions } from "@bitwarden/common/models/domain/storage-options";
-import {
+import { StateService as BaseStateService } from "@bitwarden/common/services/state.service";
-  StateService as BaseStateService,
-  withPrototype,
-} from "@bitwarden/common/services/state.service";
 
+import { browserSession, sessionSync } from "../decorators/session-sync-observable";
 import { Account } from "../models/account";
 import { BrowserComponentState } from "../models/browserComponentState";
 import { BrowserGroupingsComponentState } from "../models/browserGroupingsComponentState";
@@ -15,10 +14,27 @@ import { BrowserSendComponentState } from "../models/browserSendComponentState";
 
 import { StateService as StateServiceAbstraction } from "./abstractions/state.service";
 
-export class StateService
+@browserSession
+export class BrowserStateService
   extends BaseStateService<GlobalState, Account>
   implements StateServiceAbstraction
 {
+  @sessionSync({
+    initializer: Account.fromJSON as any, // TODO: Remove this any when all any types are removed from Account
+    initializeAsRecord: true,
+  })
+  protected accountsSubject: BehaviorSubject<{ [userId: string]: Account }>;
+  @sessionSync({ ctor: String })
+  protected activeAccountSubject: BehaviorSubject<string>;
+  @sessionSync({ ctor: Boolean })
+  protected activeAccountUnlockedSubject: BehaviorSubject<boolean>;
+
+  protected accountDeserializer = Account.fromJSON;
+
+  async hasInSessionMemory(key: string): Promise<boolean> {
+    return await this.memoryStorageService.has(key);
+  }
+
   async getFromSessionMemory<T>(key: string, deserializer?: (obj: Jsonify<T>) => T): Promise<T> {
     return this.memoryStorageService instanceof AbstractCachedStorageService
       ? await this.memoryStorageService.getBypassCache<T>(key, { deserializer: deserializer })
@@ -44,7 +60,6 @@ export class StateService
     );
   }
 
-  @withPrototype(BrowserGroupingsComponentState)
   async getBrowserGroupingComponentState(
     options?: StorageOptions
   ): Promise<BrowserGroupingsComponentState> {
@@ -67,7 +82,6 @@ export class StateService
     );
   }
 
-  @withPrototype(BrowserComponentState)
   async getBrowserCipherComponentState(options?: StorageOptions): Promise<BrowserComponentState> {
     return (
       await this.getAccount(this.reconcileOptions(options, await this.defaultInMemoryOptions()))
@@ -88,7 +102,6 @@ export class StateService
     );
   }
 
-  @withPrototype(BrowserSendComponentState)
   async getBrowserSendComponentState(options?: StorageOptions): Promise<BrowserSendComponentState> {
     return (
       await this.getAccount(this.reconcileOptions(options, await this.defaultInMemoryOptions()))
@@ -109,7 +122,6 @@ export class StateService
     );
   }
 
-  @withPrototype(BrowserComponentState)
   async getBrowserSendTypeComponentState(options?: StorageOptions): Promise<BrowserComponentState> {
     return (
       await this.getAccount(this.reconcileOptions(options, await this.defaultInMemoryOptions()))

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

@@ -2772,42 +2772,6 @@ export class StateService<
   }
 }
 
-export function withPrototype<T>(
-  constructor: new (...args: any[]) => T,
-  converter: (input: any) => T = (i) => i
-): (
-  target: any,
-  propertyKey: string | symbol,
-  descriptor: PropertyDescriptor
-) => { value: (...args: any[]) => Promise<T> } {
-  return (target: any, propertyKey: string | symbol, descriptor: PropertyDescriptor) => {
-    const originalMethod = descriptor.value;
-
-    return {
-      value: function (...args: any[]) {
-        const originalResult: Promise<T> = originalMethod.apply(this, args);
-
-        if (!(originalResult instanceof Promise)) {
-          throw new Error(
-            `Error applying prototype to stored value -- result is not a promise for method ${String(
-              propertyKey
-            )}`
-          );
-        }
-
-        return originalResult.then((result) => {
-          return result == null ||
-            result.constructor.name === constructor.prototype.constructor.name
-            ? converter(result as T)
-            : converter(
-                Object.create(constructor.prototype, Object.getOwnPropertyDescriptors(result)) as T
-              );
-        });
-      },
-    };
-  };
-}
-
 function withPrototypeForArrayMembers<T>(
   memberConstructor: new (...args: any[]) => T,
   memberConverter: (input: any) => T = (i) => i
@@ -2844,7 +2808,7 @@ function withPrototypeForArrayMembers<T>(
             return result.map((r) => {
               return r == null ||
                 r.constructor.name === memberConstructor.prototype.constructor.name
-                ? memberConverter(r)
+                ? r
                 : memberConverter(
                     Object.create(memberConstructor.prototype, Object.getOwnPropertyDescriptors(r))
                   );