Ps/pm 2910/handle switch messaging (#6823)

* Handle switch messaging

TODO: handle loading state for account switcher

* Async updates required for state

* Fallback to email for current account avatar

* Await un-awaited promises

* Remove unnecessary Prune

Prune was getting confused in browser and deleting memory in browser on
account switch. This method isn't needed since logout already removes
memory data, which is the condition for pruning

* Fix temp password in browser

* Use direct memory access until data is serializable

Safari uses a different message object extraction than firefox/chrome
and is removing `UInt8Array`s. Until all data passed into StorageService
is guaranteed serializable, we need to use direct access in state
service

* Reload badge and context menu on switch

* Gracefully switch account as they log out.

* Maintain location on account switch

* Remove unused state definitions

* Prefer null for state

undefined can be misinterpreted to indicate a value has not been set.

* Hack: structured clone in memory storage

We are currently getting dead objects on account switch due to updating
the object in the foreground state service. However, the storage service
is owned by the background. This structured clone hack ensures that all
objects stored in memory are owned by the appropriate context

* Null check nullable values

active account can be null, so we should include null safety in the
equality

* Correct background->foreground switch command

* Already providing background memory storage

* Handle connection and clipboard on switch account

* Prefer strict equal

* Ensure structuredClone is available to jsdom

This is a deficiency in jsdom --
jsdom/jsdom#3363 -- structured clone is well
supported.

* Fixup types in faker class

apps/browser/src/auth/popup/account-switching/account-switcher.component.ts

@@ -1,20 +1,25 @@
 import { Component } from "@angular/core";
 import { Router } from "@angular/router";
 
+import { BrowserRouterService } from "../../../platform/popup/services/browser-router.service";
 import { AccountSwitcherService } from "../services/account-switcher.service";
 
 @Component({
   templateUrl: "account-switcher.component.html",
 })
 export class AccountSwitcherComponent {
-  constructor(private accountSwitcherService: AccountSwitcherService, private router: Router) {}
+  constructor(
+    private accountSwitcherService: AccountSwitcherService,
+    private router: Router,
+    private routerService: BrowserRouterService
+  ) {}
 
   get accountOptions$() {
     return this.accountSwitcherService.accountOptions$;
   }
 
   async selectAccount(id: string) {
     await this.accountSwitcherService.selectAccount(id);
-    this.router.navigate(["/home"]);
+    this.router.navigate([this.routerService.getPreviousUrl() ?? "/home"]);
   }
 }

apps/browser/src/auth/popup/account-switching/current-account.component.html

@@ -1,5 +1,5 @@
 <div *ngIf="currentAccount$ | async as currentAccount">
   <div (click)="currentAccountClicked()" class="tw-mr-1 tw-mt-1">
-    <bit-avatar [id]="currentAccount.id" [text]="currentAccount.name"></bit-avatar>
+    <bit-avatar [id]="currentAccount.id" [text]="currentAccountName$ | async"></bit-avatar>
   </div>
 </div>

apps/browser/src/auth/popup/account-switching/current-account.component.ts

@@ -1,7 +1,9 @@
 import { Component } from "@angular/core";
 import { Router } from "@angular/router";
+import { map } from "rxjs";
 
 import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
+import { Utils } from "@bitwarden/common/platform/misc/utils";
 
 @Component({
   selector: "app-current-account",
@@ -14,7 +16,15 @@ export class CurrentAccountComponent {
     return this.accountService.activeAccount$;
   }
 
-  currentAccountClicked() {
-    this.router.navigate(["/account-switcher"]);
+  get currentAccountName$() {
+    return this.currentAccount$.pipe(
+      map((a) => {
+        return Utils.isNullOrWhitespace(a.name) ? a.email : a.name;
+      })
+    );
+  }
+
+  async currentAccountClicked() {
+    await this.router.navigate(["/account-switcher"]);
   }
 }

apps/browser/src/auth/popup/services/account-switcher.service.ts

@@ -54,7 +54,7 @@ export class AccountSwitcherService {
       return;
     }
 
-    this.accountService.switchAccount(id as UserId);
+    await this.accountService.switchAccount(id as UserId);
     this.messagingService.send("switchAccount", { userId: id });
   }
 }

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

@@ -24,6 +24,8 @@ import { TokenService as TokenServiceAbstraction } from "@bitwarden/common/auth/
 import { TwoFactorService as TwoFactorServiceAbstraction } from "@bitwarden/common/auth/abstractions/two-factor.service";
 import { UserVerificationApiServiceAbstraction } from "@bitwarden/common/auth/abstractions/user-verification/user-verification-api.service.abstraction";
 import { UserVerificationService as UserVerificationServiceAbstraction } from "@bitwarden/common/auth/abstractions/user-verification/user-verification.service.abstraction";
+import { AuthenticationStatus } from "@bitwarden/common/auth/enums/authentication-status";
+import { ForceSetPasswordReason } from "@bitwarden/common/auth/models/domain/force-set-password-reason";
 import { AccountServiceImplementation } from "@bitwarden/common/auth/services/account.service";
 import { AuthRequestCryptoServiceImplementation } from "@bitwarden/common/auth/services/auth-request-crypto.service.implementation";
 import { AuthService } from "@bitwarden/common/auth/services/auth.service";
@@ -87,6 +89,7 @@ import {
 import { SendApiService } from "@bitwarden/common/tools/send/services/send-api.service";
 import { SendApiService as SendApiServiceAbstraction } from "@bitwarden/common/tools/send/services/send-api.service.abstraction";
 import { InternalSendService as InternalSendServiceAbstraction } from "@bitwarden/common/tools/send/services/send.service.abstraction";
+import { UserId } from "@bitwarden/common/types/guid";
 import { CipherService as CipherServiceAbstraction } from "@bitwarden/common/vault/abstractions/cipher.service";
 import { CollectionService as CollectionServiceAbstraction } from "@bitwarden/common/vault/abstractions/collection.service";
 import { Fido2AuthenticatorService as Fido2AuthenticatorServiceAbstraction } from "@bitwarden/common/vault/abstractions/fido2/fido2-authenticator.service.abstraction";
@@ -829,6 +832,32 @@ export default class MainBackground {
     }
   }
 
+  async switchAccount(userId: UserId) {
+    if (userId != null) {
+      await this.stateService.setActiveUser(userId);
+    }
+
+    const status = await this.authService.getAuthStatus(userId);
+    const forcePasswordReset =
+      (await this.stateService.getForceSetPasswordReason({ userId: userId })) !=
+      ForceSetPasswordReason.None;
+
+    await this.systemService.clearPendingClipboard();
+    await this.notificationsService.updateConnection(false);
+
+    if (status === AuthenticationStatus.Locked) {
+      this.messagingService.send("locked", { userId: userId });
+    } else if (forcePasswordReset) {
+      this.messagingService.send("update-temp-password", { userId: userId });
+    } else {
+      this.messagingService.send("unlocked", { userId: userId });
+      await this.refreshBadge();
+      await this.refreshMenu();
+      await this.syncService.fullSync(false);
+      this.messagingService.send("switchAccountFinish", { userId: userId });
+    }
+  }
+
   async logout(expired: boolean, userId?: string) {
     await this.eventUploadService.uploadEvents(userId);
 
@@ -849,7 +878,14 @@ export default class MainBackground {
     //Needs to be checked before state is cleaned
     const needStorageReseed = await this.needsStorageReseed();
 
-    await this.stateService.clean({ userId: userId });
+    const newActiveUser = await this.stateService.clean({ userId: userId });
+
+    if (newActiveUser != null) {
+      // we have a new active user, do not continue tearing down application
+      this.switchAccount(newActiveUser as UserId);
+      this.messagingService.send("switchAccountFinish");
+      return;
+    }
 
     if (userId == null || userId === (await this.stateService.getUserId())) {
       this.searchService.clearIndex();

apps/browser/src/background/runtime.background.ts

@@ -294,6 +294,9 @@ export default class RuntimeBackground {
             }
           }
         );
+      case "switchAccount": {
+        await this.main.switchAccount(msg.userId);
+      }
     }
   }
 

apps/browser/src/platform/storage/background-memory-storage.service.ts

@@ -53,7 +53,7 @@ export class BackgroundMemoryStorageService extends MemoryStorageService {
         break;
       }
       case "save":
-        await this.save(message.key, JSON.parse(message.data as string) as unknown);
+        await this.save(message.key, JSON.parse((message.data as string) ?? null) as unknown);
         break;
       case "remove":
         await this.remove(message.key);

apps/browser/src/platform/storage/foreground-memory-storage.service.ts

@@ -78,7 +78,7 @@ export class ForegroundMemoryStorageService extends AbstractMemoryStorageService
     const response = firstValueFrom(
       this._backgroundResponses$.pipe(
         filter((message) => message.id === id),
-        map((message) => JSON.parse(message.data as string) as T)
+        map((message) => JSON.parse((message.data as string) ?? null) as T)
       )
     );
 

apps/browser/src/platform/storage/memory-storage-service-interactions.spec.ts

@@ -1,3 +1,8 @@
+/**
+ * need to update test environment so structuredClone works appropriately
+ * @jest-environment ../../libs/shared/test.environment.ts
+ */
+
 import { trackEmissions } from "@bitwarden/common/../spec/utils";
 
 import { BackgroundMemoryStorageService } from "./background-memory-storage.service";

apps/browser/src/popup/app-routing.module.ts

@@ -366,7 +366,7 @@ const routes: Routes = [
   {
     path: "account-switcher",
     component: AccountSwitcherComponent,
-    data: { state: "account-switcher" },
+    data: { state: "account-switcher", doNotSaveUrl: true },
   },
 ];
 

apps/browser/src/popup/app.component.ts

@@ -100,7 +100,10 @@ export class AppComponent implements OnInit, OnDestroy {
         this.changeDetectorRef.detectChanges();
       } else if (msg.command === "authBlocked") {
         this.router.navigate(["home"]);
-      } else if (msg.command === "locked" && msg.userId == null) {
+      } else if (
+        msg.command === "locked" &&
+        (msg.userId === null || msg.userId == this.activeUserId)
+      ) {
         this.router.navigate(["lock"]);
       } else if (msg.command === "showDialog") {
         this.showDialog(msg);
@@ -123,6 +126,11 @@ export class AppComponent implements OnInit, OnDestroy {
         this.router.navigate(["/"]);
       } else if (msg.command === "convertAccountToKeyConnector") {
         this.router.navigate(["/remove-password"]);
+      } else if (msg.command === "switchAccountFinish") {
+        // TODO: unset loading?
+        // this.loading = false;
+      } else if (msg.command == "update-temp-password") {
+        this.router.navigate(["/update-temp-password"]);
       } else {
         msg.webExtSender = sender;
         this.broadcasterService.send(msg);

apps/desktop/src/app/layout/account-switcher.component.ts

@@ -136,7 +136,7 @@ export class AccountSwitcherComponent implements OnInit, OnDestroy {
     this.close();
     await this.stateService.setActiveUser(null);
     await this.stateService.setRememberedEmail(null);
-    this.router.navigate(["/login"]);
+    await this.router.navigate(["/login"]);
   }
 
   private async createInactiveAccounts(baseAccounts: {

libs/common/spec/fake-state-provider.ts

@@ -0,0 +1,49 @@
+import {
+  GlobalState,
+  GlobalStateProvider,
+  KeyDefinition,
+  UserState,
+  UserStateProvider,
+} from "../src/platform/state";
+
+import { FakeGlobalState, FakeUserState } from "./fake-state";
+
+export class FakeGlobalStateProvider implements GlobalStateProvider {
+  states: Map<KeyDefinition<unknown>, GlobalState<unknown>> = new Map();
+  get<T>(keyDefinition: KeyDefinition<T>): GlobalState<T> {
+    let result = this.states.get(keyDefinition) as GlobalState<T>;
+
+    if (result == null) {
+      result = new FakeGlobalState<T>();
+      this.states.set(keyDefinition, result);
+    }
+    return result;
+  }
+
+  getFake<T>(keyDefinition: KeyDefinition<T>): FakeGlobalState<T> {
+    const key = Array.from(this.states.keys()).find(
+      (k) => k.stateDefinition === keyDefinition.stateDefinition && k.key === keyDefinition.key
+    );
+    return this.get(key) as FakeGlobalState<T>;
+  }
+}
+
+export class FakeUserStateProvider implements UserStateProvider {
+  states: Map<KeyDefinition<unknown>, UserState<unknown>> = new Map();
+  get<T>(keyDefinition: KeyDefinition<T>): UserState<T> {
+    let result = this.states.get(keyDefinition) as UserState<T>;
+
+    if (result == null) {
+      result = new FakeUserState<T>();
+      this.states.set(keyDefinition, result);
+    }
+    return result;
+  }
+
+  getFake<T>(keyDefinition: KeyDefinition<T>): FakeUserState<T> {
+    const key = Array.from(this.states.keys()).find(
+      (k) => k.stateDefinition === keyDefinition.stateDefinition && k.key === keyDefinition.key
+    );
+    return this.get(key) as FakeUserState<T>;
+  }
+}

libs/common/spec/fake-state.ts

@@ -0,0 +1,99 @@
+import { ReplaySubject, firstValueFrom, timeout } from "rxjs";
+
+import { DerivedUserState, GlobalState, UserState } from "../src/platform/state";
+// eslint-disable-next-line import/no-restricted-paths -- using unexposed options for clean typing in test class
+import { StateUpdateOptions } from "../src/platform/state/state-update-options";
+import { UserId } from "../src/types/guid";
+
+const DEFAULT_TEST_OPTIONS: StateUpdateOptions<any, any> = {
+  shouldUpdate: () => true,
+  combineLatestWith: null,
+  msTimeout: 10,
+};
+
+function populateOptionsWithDefault(
+  options: StateUpdateOptions<any, any>
+): StateUpdateOptions<any, any> {
+  return {
+    ...DEFAULT_TEST_OPTIONS,
+    ...options,
+  };
+}
+
+export class FakeGlobalState<T> implements GlobalState<T> {
+  // eslint-disable-next-line rxjs/no-exposed-subjects -- exposed for testing setup
+  stateSubject = new ReplaySubject<T>(1);
+
+  update: <TCombine>(
+    configureState: (state: T, dependency: TCombine) => T,
+    options?: StateUpdateOptions<T, TCombine>
+  ) => Promise<T> = jest.fn(async (configureState, options) => {
+    options = populateOptionsWithDefault(options);
+    if (this.stateSubject["_buffer"].length == 0) {
+      // throw a more helpful not initialized error
+      throw new Error(
+        "You must initialize the state with a value before calling update. Try calling `stateSubject.next(initialState)` before calling update"
+      );
+    }
+    const current = await firstValueFrom(this.state$.pipe(timeout(100)));
+    const combinedDependencies =
+      options.combineLatestWith != null
+        ? await firstValueFrom(options.combineLatestWith.pipe(timeout(options.msTimeout)))
+        : null;
+    if (!options.shouldUpdate(current, combinedDependencies)) {
+      return;
+    }
+    const newState = configureState(current, combinedDependencies);
+    this.stateSubject.next(newState);
+    return newState;
+  });
+
+  updateMock = this.update as jest.MockedFunction<typeof this.update>;
+
+  get state$() {
+    return this.stateSubject.asObservable();
+  }
+}
+
+export class FakeUserState<T> implements UserState<T> {
+  // eslint-disable-next-line rxjs/no-exposed-subjects -- exposed for testing setup
+  stateSubject = new ReplaySubject<T>(1);
+
+  update: <TCombine>(
+    configureState: (state: T, dependency: TCombine) => T,
+    options?: StateUpdateOptions<T, TCombine>
+  ) => Promise<T> = jest.fn(async (configureState, options) => {
+    options = populateOptionsWithDefault(options);
+    const current = await firstValueFrom(this.state$.pipe(timeout(options.msTimeout)));
+    const combinedDependencies =
+      options.combineLatestWith != null
+        ? await firstValueFrom(options.combineLatestWith.pipe(timeout(options.msTimeout)))
+        : null;
+    if (!options.shouldUpdate(current, combinedDependencies)) {
+      return;
+    }
+    const newState = configureState(current, combinedDependencies);
+    this.stateSubject.next(newState);
+    return newState;
+  });
+
+  updateMock = this.update as jest.MockedFunction<typeof this.update>;
+
+  updateFor: <TCombine>(
+    userId: UserId,
+    configureState: (state: T, dependency: TCombine) => T,
+    options?: StateUpdateOptions<T, TCombine>
+  ) => Promise<T> = jest.fn();
+
+  createDerived: <TTo>(
+    converter: (data: T, context: any) => Promise<TTo>
+  ) => DerivedUserState<TTo> = jest.fn();
+
+  getFromState: () => Promise<T> = jest.fn(async () => {
+    return await firstValueFrom(this.state$.pipe(timeout(10)));
+  });
+
+  get state$() {
+    return this.stateSubject.asObservable();
+  }
+}