From 9c7240d3ab24bddc2193658192447cfcd694e6b5 Mon Sep 17 00:00:00 2001
From: Maciej Zieniuk <mzieniuk@bitwarden.com>
Date: Thu, 19 Dec 2024 20:32:22 +0000
Subject: [PATCH] PM-16220: Account does not exist during login race condition

Wait for an account to become available from separate observable, instead of blindly accepting that the value is there using `firstValueFrom`, while it's sometimes not there immediately.
---
 libs/common/package.json                      |  3 +-
 .../src/auth/services/account.service.spec.ts | 31 ++++++++++++++++++
 .../src/auth/services/account.service.ts      | 32 ++++++++++++-------
 3 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/libs/common/package.json b/libs/common/package.json
index 5e0f5ae20c6..ad2771e2fff 100644
--- a/libs/common/package.json
+++ b/libs/common/package.json
@@ -15,6 +15,7 @@
   "scripts": {
     "clean": "rimraf dist",
     "build": "npm run clean && tsc",
-    "build:watch": "npm run clean && tsc -watch"
+    "build:watch": "npm run clean && tsc -watch",
+    "test": "jest"
   }
 }
diff --git a/libs/common/src/auth/services/account.service.spec.ts b/libs/common/src/auth/services/account.service.spec.ts
index 227949156ee..2028a7e1fc4 100644
--- a/libs/common/src/auth/services/account.service.spec.ts
+++ b/libs/common/src/auth/services/account.service.spec.ts
@@ -69,6 +69,7 @@ describe("accountService", () => {
   let sut: AccountServiceImplementation;
   let accountsState: FakeGlobalState<Record<UserId, AccountInfo>>;
   let activeAccountIdState: FakeGlobalState<UserId>;
+  let accountActivityState: FakeGlobalState<Record<UserId, Date>>;
   const userId = Utils.newGuid() as UserId;
   const userInfo = { email: "email", name: "name", emailVerified: true };
 
@@ -81,6 +82,7 @@ describe("accountService", () => {
 
     accountsState = globalStateProvider.getFake(ACCOUNT_ACCOUNTS);
     activeAccountIdState = globalStateProvider.getFake(ACCOUNT_ACTIVE_ACCOUNT_ID);
+    accountActivityState = globalStateProvider.getFake(ACCOUNT_ACTIVITY);
   });
 
   afterEach(() => {
@@ -256,6 +258,7 @@ describe("accountService", () => {
     beforeEach(() => {
       accountsState.stateSubject.next({ [userId]: userInfo });
       activeAccountIdState.stateSubject.next(userId);
+      accountActivityState.stateSubject.next({ [userId]: new Date(1) });
     });
 
     it("should emit null if no account is provided", async () => {
@@ -269,6 +272,34 @@ describe("accountService", () => {
       // eslint-disable-next-line @typescript-eslint/no-floating-promises
       expect(sut.switchAccount("unknown" as UserId)).rejects.toThrowError("Account does not exist");
     });
+
+    it("should change active account when switched to the new account", async () => {
+      const newUserId = Utils.newGuid() as UserId;
+      accountsState.stateSubject.next({ [newUserId]: userInfo });
+
+      await sut.switchAccount(newUserId);
+
+      await expect(firstValueFrom(sut.activeAccount$)).resolves.toEqual({
+        id: newUserId,
+        ...userInfo,
+      });
+      await expect(firstValueFrom(sut.accountActivity$)).resolves.toEqual({
+        [userId]: new Date(1),
+        [newUserId]: expect.toAlmostEqual(new Date(), 1000),
+      });
+    });
+
+    it("should not change active account when already switched to the same account", async () => {
+      await sut.switchAccount(userId);
+
+      await expect(firstValueFrom(sut.activeAccount$)).resolves.toEqual({
+        id: userId,
+        ...userInfo,
+      });
+      await expect(firstValueFrom(sut.accountActivity$)).resolves.toEqual({
+        [userId]: new Date(1),
+      });
+    });
   });
 
   describe("account activity", () => {
diff --git a/libs/common/src/auth/services/account.service.ts b/libs/common/src/auth/services/account.service.ts
index d4479815c5d..673d88382fb 100644
--- a/libs/common/src/auth/services/account.service.ts
+++ b/libs/common/src/auth/services/account.service.ts
@@ -7,6 +7,9 @@ import {
   shareReplay,
   combineLatest,
   Observable,
+  filter,
+  timeout,
+  of,
 } from "rxjs";
 
 import {
@@ -149,21 +152,28 @@ export class AccountServiceImplementation implements InternalAccountService {
   async switchAccount(userId: UserId | null): Promise<void> {
     let updateActivity = false;
     await this.activeAccountIdState.update(
-      (_, accounts) => {
-        if (userId == null) {
-          // indicates no account is active
-          return null;
-        }
-
-        if (accounts?.[userId] == null) {
-          throw new Error("Account does not exist");
-        }
+      (_, __) => {
         updateActivity = true;
         return userId;
       },
       {
-        combineLatestWith: this.accounts$,
-        shouldUpdate: (id) => {
+        combineLatestWith: this.accountsState.state$.pipe(
+          filter((accounts) => {
+            if (userId == null) {
+              // Don't worry about accounts when we are about to set active user to null
+              return true;
+            }
+
+            return accounts?.[userId] != null;
+          }),
+          // If we don't get the desired account with enough time, just return empty as that will result in the same error
+          timeout({ first: 1000, with: () => of({} as Record<UserId, AccountInfo>) }),
+        ),
+        shouldUpdate: (id, accounts) => {
+          if (userId != null && accounts?.[userId] == null) {
+            throw new Error("Account does not exist");
+          }
+
           // update only if userId changes
           return id !== userId;
         },