Auth/PM-6689 - Migrate Security Stamp to Token Service and State Provider (#8792)

* PM-6689 - Add security stamp to Token state

* PM-6689 - Remove Security Stamp from account and state service

* PM-6689 - Add security stamp get and set to token service + abstraction + tests

* PM-6689 - Add migration for security stamp, test it, and register it with migrator

* PM-6689 - Update sync service + deps to use token service.

* PM-6689 - Cleanup missed usages of account tokens which has been removed.

* PM-6689 - Per PR feedback, remove unnecessary data migration as the security stamp is only in memory and doesn't need to be migrated.

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

@@ -788,6 +788,7 @@ export default class MainBackground {
       this.avatarService,
       logoutCallback,
       this.billingAccountProfileStateService,
+      this.tokenService,
     );
     this.eventUploadService = new EventUploadService(
       this.apiService,

apps/cli/src/bw.ts

@@ -633,6 +633,7 @@ export class Main {
       this.avatarService,
       async (expired: boolean) => await this.logout(),
       this.billingAccountProfileStateService,
+      this.tokenService,
     );
 
     this.totpService = new TotpService(this.cryptoFunctionService, this.logService);

libs/angular/src/services/jslib-services.module.ts

@@ -628,6 +628,7 @@ const safeProviders: SafeProvider[] = [
       AvatarServiceAbstraction,
       LOGOUT_CALLBACK,
       BillingAccountProfileStateService,
+      TokenServiceAbstraction,
     ],
   }),
   safeProvider({

libs/auth/src/common/login-strategies/login.strategy.spec.ts

@@ -27,7 +27,6 @@ import { Utils } from "@bitwarden/common/platform/misc/utils";
 import {
   Account,
   AccountProfile,
-  AccountTokens,
   AccountKeys,
 } from "@bitwarden/common/platform/models/domain/account";
 import { EncString } from "@bitwarden/common/platform/models/domain/enc-string";
@@ -213,9 +212,6 @@ describe("LoginStrategy", () => {
               kdfType: kdf,
             },
           },
-          tokens: {
-            ...new AccountTokens(),
-          },
           keys: new AccountKeys(),
         }),
       );

libs/auth/src/common/login-strategies/login.strategy.ts

@@ -27,11 +27,7 @@ import { LogService } from "@bitwarden/common/platform/abstractions/log.service"
 import { MessagingService } from "@bitwarden/common/platform/abstractions/messaging.service";
 import { PlatformUtilsService } from "@bitwarden/common/platform/abstractions/platform-utils.service";
 import { StateService } from "@bitwarden/common/platform/abstractions/state.service";
-import {
-  Account,
-  AccountProfile,
-  AccountTokens,
-} from "@bitwarden/common/platform/models/domain/account";
+import { Account, AccountProfile } from "@bitwarden/common/platform/models/domain/account";
 import { UserId } from "@bitwarden/common/types/guid";
 
 import { InternalUserDecryptionOptionsServiceAbstraction } from "../abstractions/user-decryption-options.service.abstraction";
@@ -192,9 +188,6 @@ export abstract class LoginStrategy {
             kdfType: tokenResponse.kdf,
           },
         },
-        tokens: {
-          ...new AccountTokens(),
-        },
       }),
     );
 

libs/common/src/auth/abstractions/token.service.ts

@@ -213,4 +213,10 @@ export abstract class TokenService {
    * @returns A promise that resolves with a boolean representing the user's external authN status.
    */
   getIsExternal: () => Promise<boolean>;
+
+  /** Gets the active or passed in user's security stamp */
+  getSecurityStamp: (userId?: UserId) => Promise<string | null>;
+
+  /** Sets the security stamp for the active or passed in user */
+  setSecurityStamp: (securityStamp: string, userId?: UserId) => Promise<void>;
 }

libs/common/src/auth/services/token.service.spec.ts

@@ -23,6 +23,7 @@ import {
   EMAIL_TWO_FACTOR_TOKEN_RECORD_DISK_LOCAL,
   REFRESH_TOKEN_DISK,
   REFRESH_TOKEN_MEMORY,
+  SECURITY_STAMP_MEMORY,
 } from "./token.state";
 
 describe("TokenService", () => {
@@ -2191,6 +2192,84 @@ describe("TokenService", () => {
     });
   });
 
+  describe("Security Stamp methods", () => {
+    const mockSecurityStamp = "securityStamp";
+
+    describe("setSecurityStamp", () => {
+      it("should throw an error if no user id is provided and there is no active user in global state", async () => {
+        // Act
+        // note: don't await here because we want to test the error
+        const result = tokenService.setSecurityStamp(mockSecurityStamp);
+        // Assert
+        await expect(result).rejects.toThrow("User id not found. Cannot set security stamp.");
+      });
+
+      it("should set the security stamp in memory when there is an active user in global state", async () => {
+        // Arrange
+        globalStateProvider
+          .getFake(ACCOUNT_ACTIVE_ACCOUNT_ID)
+          .stateSubject.next(userIdFromAccessToken);
+
+        // Act
+        await tokenService.setSecurityStamp(mockSecurityStamp);
+
+        // Assert
+        expect(
+          singleUserStateProvider.getFake(userIdFromAccessToken, SECURITY_STAMP_MEMORY).nextMock,
+        ).toHaveBeenCalledWith(mockSecurityStamp);
+      });
+
+      it("should set the security stamp in memory for the specified user id", async () => {
+        // Act
+        await tokenService.setSecurityStamp(mockSecurityStamp, userIdFromAccessToken);
+
+        // Assert
+        expect(
+          singleUserStateProvider.getFake(userIdFromAccessToken, SECURITY_STAMP_MEMORY).nextMock,
+        ).toHaveBeenCalledWith(mockSecurityStamp);
+      });
+    });
+
+    describe("getSecurityStamp", () => {
+      it("should throw an error if no user id is provided and there is no active user in global state", async () => {
+        // Act
+        // note: don't await here because we want to test the error
+        const result = tokenService.getSecurityStamp();
+        // Assert
+        await expect(result).rejects.toThrow("User id not found. Cannot get security stamp.");
+      });
+
+      it("should return the security stamp from memory with no user id specified (uses global active user)", async () => {
+        // Arrange
+        globalStateProvider
+          .getFake(ACCOUNT_ACTIVE_ACCOUNT_ID)
+          .stateSubject.next(userIdFromAccessToken);
+
+        singleUserStateProvider
+          .getFake(userIdFromAccessToken, SECURITY_STAMP_MEMORY)
+          .stateSubject.next([userIdFromAccessToken, mockSecurityStamp]);
+
+        // Act
+        const result = await tokenService.getSecurityStamp();
+
+        // Assert
+        expect(result).toEqual(mockSecurityStamp);
+      });
+
+      it("should return the security stamp from memory for the specified user id", async () => {
+        // Arrange
+        singleUserStateProvider
+          .getFake(userIdFromAccessToken, SECURITY_STAMP_MEMORY)
+          .stateSubject.next([userIdFromAccessToken, mockSecurityStamp]);
+
+        // Act
+        const result = await tokenService.getSecurityStamp(userIdFromAccessToken);
+        // Assert
+        expect(result).toEqual(mockSecurityStamp);
+      });
+    });
+  });
+
   // Helpers
   function createTokenService(supportsSecureStorage: boolean) {
     return new TokenService(

libs/common/src/auth/services/token.service.ts

@@ -32,6 +32,7 @@ import {
   EMAIL_TWO_FACTOR_TOKEN_RECORD_DISK_LOCAL,
   REFRESH_TOKEN_DISK,
   REFRESH_TOKEN_MEMORY,
+  SECURITY_STAMP_MEMORY,
 } from "./token.state";
 
 export enum TokenStorageLocation {
@@ -850,6 +851,30 @@ export class TokenService implements TokenServiceAbstraction {
     return Array.isArray(decoded.amr) && decoded.amr.includes("external");
   }
 
+  async getSecurityStamp(userId?: UserId): Promise<string | null> {
+    userId ??= await firstValueFrom(this.activeUserIdGlobalState.state$);
+
+    if (!userId) {
+      throw new Error("User id not found. Cannot get security stamp.");
+    }
+
+    const securityStamp = await this.getStateValueByUserIdAndKeyDef(userId, SECURITY_STAMP_MEMORY);
+
+    return securityStamp;
+  }
+
+  async setSecurityStamp(securityStamp: string, userId?: UserId): Promise<void> {
+    userId ??= await firstValueFrom(this.activeUserIdGlobalState.state$);
+
+    if (!userId) {
+      throw new Error("User id not found. Cannot set security stamp.");
+    }
+
+    await this.singleUserStateProvider
+      .get(userId, SECURITY_STAMP_MEMORY)
+      .update((_) => securityStamp);
+  }
+
   private async getStateValueByUserIdAndKeyDef(
     userId: UserId,
     storageLocation: UserKeyDefinition<string>,

libs/common/src/auth/services/token.state.spec.ts

@@ -10,6 +10,7 @@ import {
   EMAIL_TWO_FACTOR_TOKEN_RECORD_DISK_LOCAL,
   REFRESH_TOKEN_DISK,
   REFRESH_TOKEN_MEMORY,
+  SECURITY_STAMP_MEMORY,
 } from "./token.state";
 
 describe.each([
@@ -22,6 +23,7 @@ describe.each([
   [API_KEY_CLIENT_ID_MEMORY, "apiKeyClientIdMemory"],
   [API_KEY_CLIENT_SECRET_DISK, "apiKeyClientSecretDisk"],
   [API_KEY_CLIENT_SECRET_MEMORY, "apiKeyClientSecretMemory"],
+  [SECURITY_STAMP_MEMORY, "securityStamp"],
 ])(
   "deserializes state key definitions",
   (

libs/common/src/auth/services/token.state.ts

@@ -69,3 +69,8 @@ export const API_KEY_CLIENT_SECRET_MEMORY = new UserKeyDefinition<string>(
     clearOn: [], // Manually handled
   },
 );
+
+export const SECURITY_STAMP_MEMORY = new UserKeyDefinition<string>(TOKEN_MEMORY, "securityStamp", {
+  deserializer: (securityStamp) => securityStamp,
+  clearOn: ["logout"],
+});

libs/common/src/platform/abstractions/state.service.ts

@@ -181,8 +181,6 @@ export abstract class StateService<T extends Account = Account> {
    * Sets the user's Pin, encrypted by the user key
    */
   setProtectedPin: (value: string, options?: StorageOptions) => Promise<void>;
-  getSecurityStamp: (options?: StorageOptions) => Promise<string>;
-  setSecurityStamp: (value: string, options?: StorageOptions) => Promise<void>;
   getUserId: (options?: StorageOptions) => Promise<string>;
   getVaultTimeout: (options?: StorageOptions) => Promise<number>;
   setVaultTimeout: (value: number, options?: StorageOptions) => Promise<void>;

libs/common/src/platform/models/domain/account-tokens.spec.ts

@@ -1,9 +0,0 @@
-import { AccountTokens } from "./account";
-
-describe("AccountTokens", () => {
-  describe("fromJSON", () => {
-    it("should deserialize to an instance of itself", () => {
-      expect(AccountTokens.fromJSON({})).toBeInstanceOf(AccountTokens);
-    });
-  });
-});

libs/common/src/platform/models/domain/account.spec.ts

@@ -1,4 +1,4 @@
-import { Account, AccountKeys, AccountProfile, AccountSettings, AccountTokens } from "./account";
+import { Account, AccountKeys, AccountProfile, AccountSettings } from "./account";
 
 describe("Account", () => {
   describe("fromJSON", () => {
@@ -10,14 +10,12 @@ describe("Account", () => {
       const keysSpy = jest.spyOn(AccountKeys, "fromJSON");
       const profileSpy = jest.spyOn(AccountProfile, "fromJSON");
       const settingsSpy = jest.spyOn(AccountSettings, "fromJSON");
-      const tokensSpy = jest.spyOn(AccountTokens, "fromJSON");
 
       Account.fromJSON({});
 
       expect(keysSpy).toHaveBeenCalled();
       expect(profileSpy).toHaveBeenCalled();
       expect(settingsSpy).toHaveBeenCalled();
-      expect(tokensSpy).toHaveBeenCalled();
     });
   });
 });

libs/common/src/platform/models/domain/account.ts

@@ -171,24 +171,11 @@ export class AccountSettings {
   }
 }
 
-export class AccountTokens {
-  securityStamp?: string;
-
-  static fromJSON(obj: Jsonify<AccountTokens>): AccountTokens {
-    if (obj == null) {
-      return null;
-    }
-
-    return Object.assign(new AccountTokens(), obj);
-  }
-}
-
 export class Account {
   data?: AccountData = new AccountData();
   keys?: AccountKeys = new AccountKeys();
   profile?: AccountProfile = new AccountProfile();
   settings?: AccountSettings = new AccountSettings();
-  tokens?: AccountTokens = new AccountTokens();
 
   constructor(init: Partial<Account>) {
     Object.assign(this, {
@@ -208,10 +195,6 @@ export class Account {
         ...new AccountSettings(),
         ...init?.settings,
       },
-      tokens: {
-        ...new AccountTokens(),
-        ...init?.tokens,
-      },
     });
   }
 
@@ -225,7 +208,6 @@ export class Account {
       data: AccountData.fromJSON(json?.data),
       profile: AccountProfile.fromJSON(json?.profile),
       settings: AccountSettings.fromJSON(json?.settings),
-      tokens: AccountTokens.fromJSON(json?.tokens),
     });
   }
 }

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

@@ -839,23 +839,6 @@ export class StateService<
     );
   }
 
-  async getSecurityStamp(options?: StorageOptions): Promise<string> {
-    return (
-      await this.getAccount(this.reconcileOptions(options, await this.defaultInMemoryOptions()))
-    )?.tokens?.securityStamp;
-  }
-
-  async setSecurityStamp(value: string, options?: StorageOptions): Promise<void> {
-    const account = await this.getAccount(
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-    account.tokens.securityStamp = value;
-    await this.saveAccount(
-      account,
-      this.reconcileOptions(options, await this.defaultInMemoryOptions()),
-    );
-  }
-
   async getUserId(options?: StorageOptions): Promise<string> {
     return (
       await this.getAccount(this.reconcileOptions(options, await this.defaultOnDiskOptions()))

libs/common/src/vault/services/sync/sync.service.ts

@@ -15,6 +15,7 @@ import { AccountService } from "../../../auth/abstractions/account.service";
 import { AvatarService } from "../../../auth/abstractions/avatar.service";
 import { KeyConnectorService } from "../../../auth/abstractions/key-connector.service";
 import { InternalMasterPasswordServiceAbstraction } from "../../../auth/abstractions/master-password.service.abstraction";
+import { TokenService } from "../../../auth/abstractions/token.service";
 import { ForceSetPasswordReason } from "../../../auth/models/domain/force-set-password-reason";
 import { DomainSettingsService } from "../../../autofill/services/domain-settings.service";
 import { BillingAccountProfileStateService } from "../../../billing/abstractions/account/billing-account-profile-state.service";
@@ -73,6 +74,7 @@ export class SyncService implements SyncServiceAbstraction {
     private avatarService: AvatarService,
     private logoutCallback: (expired: boolean) => Promise<void>,
     private billingAccountProfileStateService: BillingAccountProfileStateService,
+    private tokenService: TokenService,
   ) {}
 
   async getLastSync(): Promise<Date> {
@@ -309,7 +311,7 @@ export class SyncService implements SyncServiceAbstraction {
   }
 
   private async syncProfile(response: ProfileResponse) {
-    const stamp = await this.stateService.getSecurityStamp();
+    const stamp = await this.tokenService.getSecurityStamp(response.id as UserId);
     if (stamp != null && stamp !== response.securityStamp) {
       if (this.logoutCallback != null) {
         await this.logoutCallback(true);
@@ -323,7 +325,7 @@ export class SyncService implements SyncServiceAbstraction {
     await this.cryptoService.setProviderKeys(response.providers);
     await this.cryptoService.setOrgKeys(response.organizations, response.providerOrganizations);
     await this.avatarService.setSyncAvatarColor(response.id as UserId, response.avatarColor);
-    await this.stateService.setSecurityStamp(response.securityStamp);
+    await this.tokenService.setSecurityStamp(response.securityStamp, response.id as UserId);
     await this.stateService.setEmailVerified(response.emailVerified);
 
     await this.billingAccountProfileStateService.setHasPremium(