Auth `UserKeyDefinition` Migration (#8587)

* Migrate DeviceTrustCryptoService

* Migrate SsoLoginService

* Migrate TokenService

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

Co-authored-by: Jared Snider <116684653+JaredSnider-Bitwarden@users.noreply.github.com>

* Fix Test

* Actually Fix Tests

---------

Co-authored-by: Jared Snider <116684653+JaredSnider-Bitwarden@users.noreply.github.com>

libs/common/src/auth/services/device-trust-crypto.service.implementation.ts

@@ -14,7 +14,7 @@ import { StorageLocation } from "../../platform/enums";
 import { EncString } from "../../platform/models/domain/enc-string";
 import { StorageOptions } from "../../platform/models/domain/storage-options";
 import { SymmetricCryptoKey } from "../../platform/models/domain/symmetric-crypto-key";
-import { DEVICE_TRUST_DISK_LOCAL, KeyDefinition, StateProvider } from "../../platform/state";
+import { DEVICE_TRUST_DISK_LOCAL, StateProvider, UserKeyDefinition } from "../../platform/state";
 import { UserId } from "../../types/guid";
 import { UserKey, DeviceKey } from "../../types/key";
 import { DeviceTrustCryptoServiceAbstraction } from "../abstractions/device-trust-crypto.service.abstraction";
@@ -27,16 +27,18 @@ import {
 } from "../models/request/update-devices-trust.request";
 
 /** Uses disk storage so that the device key can persist after log out and tab removal. */
-export const DEVICE_KEY = new KeyDefinition<DeviceKey>(DEVICE_TRUST_DISK_LOCAL, "deviceKey", {
+export const DEVICE_KEY = new UserKeyDefinition<DeviceKey>(DEVICE_TRUST_DISK_LOCAL, "deviceKey", {
   deserializer: (deviceKey) => SymmetricCryptoKey.fromJSON(deviceKey) as DeviceKey,
+  clearOn: [], // Device key is needed to log back into device, so we can't clear it automatically during lock or logout
 });
 
 /** Uses disk storage so that the shouldTrustDevice bool can persist across login. */
-export const SHOULD_TRUST_DEVICE = new KeyDefinition<boolean>(
+export const SHOULD_TRUST_DEVICE = new UserKeyDefinition<boolean>(
   DEVICE_TRUST_DISK_LOCAL,
   "shouldTrustDevice",
   {
     deserializer: (shouldTrustDevice) => shouldTrustDevice,
+    clearOn: [], // Need to preserve the user setting, so we can't clear it automatically during lock or logout
   },
 );
 

libs/common/src/auth/services/sso-login.service.ts

@@ -6,6 +6,7 @@ import {
   KeyDefinition,
   SSO_DISK,
   StateProvider,
+  UserKeyDefinition,
 } from "../../platform/state";
 import { SsoLoginServiceAbstraction } from "../abstractions/sso-login.service.abstraction";
 
@@ -26,7 +27,19 @@ const SSO_STATE = new KeyDefinition<string>(SSO_DISK, "ssoState", {
 /**
  * Uses disk storage so that the organization sso identifier can be persisted across sso redirects.
  */
-const ORGANIZATION_SSO_IDENTIFIER = new KeyDefinition<string>(
+const USER_ORGANIZATION_SSO_IDENTIFIER = new UserKeyDefinition<string>(
+  SSO_DISK,
+  "organizationSsoIdentifier",
+  {
+    deserializer: (organizationIdentifier) => organizationIdentifier,
+    clearOn: ["logout"], // Used for login, so not needed past logout
+  },
+);
+
+/**
+ * Uses disk storage so that the organization sso identifier can be persisted across sso redirects.
+ */
+const GLOBAL_ORGANIZATION_SSO_IDENTIFIER = new KeyDefinition<string>(
   SSO_DISK,
   "organizationSsoIdentifier",
   {
@@ -51,10 +64,10 @@ export class SsoLoginService implements SsoLoginServiceAbstraction {
   constructor(private stateProvider: StateProvider) {
     this.codeVerifierState = this.stateProvider.getGlobal(CODE_VERIFIER);
     this.ssoState = this.stateProvider.getGlobal(SSO_STATE);
-    this.orgSsoIdentifierState = this.stateProvider.getGlobal(ORGANIZATION_SSO_IDENTIFIER);
+    this.orgSsoIdentifierState = this.stateProvider.getGlobal(GLOBAL_ORGANIZATION_SSO_IDENTIFIER);
     this.ssoEmailState = this.stateProvider.getGlobal(SSO_EMAIL);
     this.activeUserOrgSsoIdentifierState = this.stateProvider.getActive(
-      ORGANIZATION_SSO_IDENTIFIER,
+      USER_ORGANIZATION_SSO_IDENTIFIER,
     );
   }
 

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

@@ -15,8 +15,8 @@ import { SymmetricCryptoKey } from "../../platform/models/domain/symmetric-crypt
 import {
   GlobalState,
   GlobalStateProvider,
-  KeyDefinition,
   SingleUserStateProvider,
+  UserKeyDefinition,
 } from "../../platform/state";
 import { UserId } from "../../types/guid";
 import { TokenService as TokenServiceAbstraction } from "../abstractions/token.service";
@@ -863,7 +863,7 @@ export class TokenService implements TokenServiceAbstraction {
 
   private async getStateValueByUserIdAndKeyDef(
     userId: UserId,
-    storageLocation: KeyDefinition<string>,
+    storageLocation: UserKeyDefinition<string>,
   ): Promise<string | undefined> {
     // read from single user state provider
     return await firstValueFrom(this.singleUserStateProvider.get(userId, storageLocation).state$);

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

@@ -1,4 +1,4 @@
-import { KeyDefinition } from "../../platform/state";
+import { KeyDefinition, UserKeyDefinition } from "../../platform/state";
 
 import {
   ACCESS_TOKEN_DISK,
@@ -28,8 +28,8 @@ describe.each([
   "deserializes state key definitions",
   (
     keyDefinition:
-      | KeyDefinition<string>
+      | UserKeyDefinition<string>
-      | KeyDefinition<boolean>
+      | UserKeyDefinition<boolean>
       | KeyDefinition<Record<string, string>>,
     state: string | boolean | Record<string, string>,
   ) => {
@@ -50,7 +50,10 @@ describe.each([
       return typeof value === "object" && value !== null && !Array.isArray(value);
     }
 
-    function testDeserialization<T>(keyDefinition: KeyDefinition<T>, state: T) {
+    function testDeserialization<T>(
+      keyDefinition: KeyDefinition<T> | UserKeyDefinition<T>,
+      state: T,
+    ) {
       const deserialized = keyDefinition.deserializer(JSON.parse(JSON.stringify(state)));
       expect(deserialized).toEqual(state);
     }

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

@@ -1,30 +1,41 @@
-import { KeyDefinition, TOKEN_DISK, TOKEN_DISK_LOCAL, TOKEN_MEMORY } from "../../platform/state";
+import {
+  KeyDefinition,
+  TOKEN_DISK,
+  TOKEN_DISK_LOCAL,
+  TOKEN_MEMORY,
+  UserKeyDefinition,
+} from "../../platform/state";
 
 // Note: all tokens / API key information must be cleared on logout.
 // because we are using secure storage, we must manually call to clean up our tokens.
 // See stateService.deAuthenticateAccount for where we call clearTokens(...)
 
-export const ACCESS_TOKEN_DISK = new KeyDefinition<string>(TOKEN_DISK, "accessToken", {
+export const ACCESS_TOKEN_DISK = new UserKeyDefinition<string>(TOKEN_DISK, "accessToken", {
   deserializer: (accessToken) => accessToken,
+  clearOn: [], // Manually handled
 });
 
-export const ACCESS_TOKEN_MEMORY = new KeyDefinition<string>(TOKEN_MEMORY, "accessToken", {
+export const ACCESS_TOKEN_MEMORY = new UserKeyDefinition<string>(TOKEN_MEMORY, "accessToken", {
   deserializer: (accessToken) => accessToken,
+  clearOn: [], // Manually handled
 });
 
-export const REFRESH_TOKEN_DISK = new KeyDefinition<string>(TOKEN_DISK, "refreshToken", {
+export const REFRESH_TOKEN_DISK = new UserKeyDefinition<string>(TOKEN_DISK, "refreshToken", {
   deserializer: (refreshToken) => refreshToken,
+  clearOn: [], // Manually handled
 });
 
-export const REFRESH_TOKEN_MEMORY = new KeyDefinition<string>(TOKEN_MEMORY, "refreshToken", {
+export const REFRESH_TOKEN_MEMORY = new UserKeyDefinition<string>(TOKEN_MEMORY, "refreshToken", {
   deserializer: (refreshToken) => refreshToken,
+  clearOn: [], // Manually handled
 });
 
-export const REFRESH_TOKEN_MIGRATED_TO_SECURE_STORAGE = new KeyDefinition<boolean>(
+export const REFRESH_TOKEN_MIGRATED_TO_SECURE_STORAGE = new UserKeyDefinition<boolean>(
   TOKEN_DISK,
   "refreshTokenMigratedToSecureStorage",
   {
     deserializer: (refreshTokenMigratedToSecureStorage) => refreshTokenMigratedToSecureStorage,
+    clearOn: [], // Don't clear on lock/logout so that we always check the correct place (secure storage) for the refresh token if it's been migrated
   },
 );
 
@@ -36,26 +47,34 @@ export const EMAIL_TWO_FACTOR_TOKEN_RECORD_DISK_LOCAL = KeyDefinition.record<str
   },
 );
 
-export const API_KEY_CLIENT_ID_DISK = new KeyDefinition<string>(TOKEN_DISK, "apiKeyClientId", {
+export const API_KEY_CLIENT_ID_DISK = new UserKeyDefinition<string>(TOKEN_DISK, "apiKeyClientId", {
   deserializer: (apiKeyClientId) => apiKeyClientId,
+  clearOn: [], // Manually handled
 });
 
-export const API_KEY_CLIENT_ID_MEMORY = new KeyDefinition<string>(TOKEN_MEMORY, "apiKeyClientId", {
+export const API_KEY_CLIENT_ID_MEMORY = new UserKeyDefinition<string>(
-  deserializer: (apiKeyClientId) => apiKeyClientId,
+  TOKEN_MEMORY,
-});
+  "apiKeyClientId",
+  {
+    deserializer: (apiKeyClientId) => apiKeyClientId,
+    clearOn: [], // Manually handled
+  },
+);
 
-export const API_KEY_CLIENT_SECRET_DISK = new KeyDefinition<string>(
+export const API_KEY_CLIENT_SECRET_DISK = new UserKeyDefinition<string>(
   TOKEN_DISK,
   "apiKeyClientSecret",
   {
     deserializer: (apiKeyClientSecret) => apiKeyClientSecret,
+    clearOn: [], // Manually handled
   },
 );
 
-export const API_KEY_CLIENT_SECRET_MEMORY = new KeyDefinition<string>(
+export const API_KEY_CLIENT_SECRET_MEMORY = new UserKeyDefinition<string>(
   TOKEN_MEMORY,
   "apiKeyClientSecret",
   {
     deserializer: (apiKeyClientSecret) => apiKeyClientSecret,
+    clearOn: [], // Manually handled
   },
 );