[PM-8933] Require userId on setUserKey (#9675)

* Updated all sets of user key to pass in userId

* Added userId on auth request login.

* Fixed tests.

* Fixed tests to pass in UserId

* Added parameter to tests.

* Addressed PR feedback.

* Merged main

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

@@ -1067,7 +1067,6 @@ export default class MainBackground {
         this.messagingService,
         this.appIdService,
         this.platformUtilsService,
-        this.stateService,
         this.logService,
         this.authService,
         this.biometricStateService,

apps/browser/src/background/nativeMessaging.background.ts

@@ -9,7 +9,6 @@ import { CryptoService } from "@bitwarden/common/platform/abstractions/crypto.se
 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 { BiometricStateService } from "@bitwarden/common/platform/biometrics/biometric-state.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
 import { EncString } from "@bitwarden/common/platform/models/domain/enc-string";
@@ -79,7 +78,6 @@ export class NativeMessagingBackground {
     private messagingService: MessagingService,
     private appIdService: AppIdService,
     private platformUtilsService: PlatformUtilsService,
-    private stateService: StateService,
     private logService: LogService,
     private authService: AuthService,
     private biometricStateService: BiometricStateService,
@@ -214,7 +212,7 @@ export class NativeMessagingBackground {
       await this.connect();
     }
 
-    message.userId = await this.stateService.getUserId();
+    message.userId = (await firstValueFrom(this.accountService.activeAccount$))?.id;
     message.timestamp = Date.now();
 
     if (this.platformUtilsService.isSafari()) {
@@ -367,13 +365,14 @@ export class NativeMessagingBackground {
     const [publicKey, privateKey] = await this.cryptoFunctionService.rsaGenerateKeyPair(2048);
     this.publicKey = publicKey;
     this.privateKey = privateKey;
+    const userId = (await firstValueFrom(this.accountService.activeAccount$))?.id;
 
     // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
     // eslint-disable-next-line @typescript-eslint/no-floating-promises
     this.sendUnencrypted({
       command: "setupEncryption",
       publicKey: Utils.fromBufferToB64(publicKey),
-      userId: await this.stateService.getUserId(),
+      userId: userId,
     });
 
     return new Promise((resolve, reject) => (this.secureSetupResolve = resolve));
@@ -391,7 +390,7 @@ export class NativeMessagingBackground {
 
   private async showFingerprintDialog() {
     const fingerprint = await this.cryptoService.getFingerprint(
-      await this.stateService.getUserId(),
+      (await firstValueFrom(this.accountService.activeAccount$))?.id,
       this.publicKey,
     );
 

apps/cli/src/auth/commands/unlock.command.ts

@@ -69,7 +69,7 @@ export class UnlockCommand {
     }
 
     const userKey = await this.masterPasswordService.decryptUserKeyWithMasterKey(masterKey);
-    await this.cryptoService.setUserKey(userKey);
+    await this.cryptoService.setUserKey(userKey, userId);
 
     if (await this.keyConnectorService.getConvertAccountRequired()) {
       const convertToKeyConnectorCommand = new ConvertToKeyConnectorCommand(

apps/desktop/src/platform/services/electron-crypto.service.ts

@@ -69,7 +69,7 @@ export class ElectronCryptoService extends CryptoService {
     await super.clearStoredUserKey(keySuffix, userId);
   }
 
-  protected override async storeAdditionalKeys(key: UserKey, userId?: UserId) {
+  protected override async storeAdditionalKeys(key: UserKey, userId: UserId) {
     await super.storeAdditionalKeys(key, userId);
 
     const storeBiometricKey = await this.shouldStoreKey(KeySuffixOptions.Biometric, userId);

libs/angular/src/auth/components/lock.component.ts

@@ -134,10 +134,13 @@ export class LockComponent implements OnInit, OnDestroy {
     }
 
     await this.biometricStateService.setUserPromptCancelled();
-    const userKey = await this.cryptoService.getUserKeyFromStorage(KeySuffixOptions.Biometric);
+    const userKey = await this.cryptoService.getUserKeyFromStorage(
+      KeySuffixOptions.Biometric,
+      this.activeUserId,
+    );
 
     if (userKey) {
-      await this.setUserKeyAndContinue(userKey, false);
+      await this.setUserKeyAndContinue(userKey, this.activeUserId, false);
     }
 
     return !!userKey;
@@ -174,7 +177,7 @@ export class LockComponent implements OnInit, OnDestroy {
       const userKey = await this.pinService.decryptUserKeyWithPin(this.pin, userId);
 
       if (userKey) {
-        await this.setUserKeyAndContinue(userKey);
+        await this.setUserKeyAndContinue(userKey, userId);
         return; // successfully unlocked
       }
 
@@ -257,11 +260,15 @@ export class LockComponent implements OnInit, OnDestroy {
     const userKey = await this.masterPasswordService.decryptUserKeyWithMasterKey(
       response.masterKey,
     );
-    await this.setUserKeyAndContinue(userKey, true);
+    await this.setUserKeyAndContinue(userKey, userId, true);
   }
 
-  private async setUserKeyAndContinue(key: UserKey, evaluatePasswordAfterUnlock = false) {
+  private async setUserKeyAndContinue(
-    await this.cryptoService.setUserKey(key);
+    key: UserKey,
+    userId: UserId,
+    evaluatePasswordAfterUnlock = false,
+  ) {
+    await this.cryptoService.setUserKey(key, userId);
 
     // Now that we have a decrypted user key in memory, we can check if we
     // need to establish trust on the current device

libs/angular/src/auth/components/login-via-auth-request.component.ts

@@ -386,6 +386,7 @@ export class LoginViaAuthRequestComponent
       await this.authRequestService.setKeysAfterDecryptingSharedMasterKeyAndHash(
         adminAuthReqResponse,
         privateKey,
+        userId,
       );
     } else {
       // Flow 3: masterPasswordHash is null
@@ -393,6 +394,7 @@ export class LoginViaAuthRequestComponent
       await this.authRequestService.setUserKeyAfterDecryptingSharedUserKey(
         adminAuthReqResponse,
         privateKey,
+        userId,
       );
     }
 

libs/auth/src/common/abstractions/auth-request.service.abstraction.ts

@@ -49,19 +49,23 @@ export abstract class AuthRequestServiceAbstraction {
    * Sets the `UserKey` from an auth request. Auth request must have a `UserKey`.
    * @param authReqResponse The auth request.
    * @param authReqPrivateKey The private key corresponding to the public key sent in the auth request.
+   * @param userId The ID of the user for whose account we will set the key.
    */
   abstract setUserKeyAfterDecryptingSharedUserKey: (
     authReqResponse: AuthRequestResponse,
     authReqPrivateKey: ArrayBuffer,
+    userId: UserId,
   ) => Promise<void>;
   /**
    * Sets the `MasterKey` and `MasterKeyHash` from an auth request. Auth request must have a `MasterKey` and `MasterKeyHash`.
    * @param authReqResponse The auth request.
    * @param authReqPrivateKey The private key corresponding to the public key sent in the auth request.
+   * @param userId The ID of the user for whose account we will set the keys.
    */
   abstract setKeysAfterDecryptingSharedMasterKeyAndHash: (
     authReqResponse: AuthRequestResponse,
     authReqPrivateKey: ArrayBuffer,
+    userId: UserId,
   ) => Promise<void>;
   /**
    * Decrypts a `UserKey` from a public key encrypted `UserKey`.

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

@@ -159,7 +159,7 @@ describe("AuthRequestLoginStrategy", () => {
       mockUserId,
     );
     expect(cryptoService.setMasterKeyEncryptedUserKey).toHaveBeenCalledWith(tokenResponse.key);
-    expect(cryptoService.setUserKey).toHaveBeenCalledWith(userKey);
+    expect(cryptoService.setUserKey).toHaveBeenCalledWith(userKey, mockUserId);
     expect(deviceTrustService.trustDeviceIfRequired).toHaveBeenCalled();
     expect(cryptoService.setPrivateKey).toHaveBeenCalledWith(tokenResponse.privateKey, mockUserId);
   });
@@ -184,7 +184,7 @@ describe("AuthRequestLoginStrategy", () => {
 
     // setMasterKeyEncryptedUserKey, setUserKey, and setPrivateKey should still be called
     expect(cryptoService.setMasterKeyEncryptedUserKey).toHaveBeenCalledWith(tokenResponse.key);
-    expect(cryptoService.setUserKey).toHaveBeenCalledWith(decUserKey);
+    expect(cryptoService.setUserKey).toHaveBeenCalledWith(decUserKey, mockUserId);
     expect(cryptoService.setPrivateKey).toHaveBeenCalledWith(tokenResponse.privateKey, mockUserId);
 
     // trustDeviceIfRequired should be called

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

@@ -102,7 +102,7 @@ export class AuthRequestLoginStrategy extends LoginStrategy {
     await this.cryptoService.setMasterKeyEncryptedUserKey(response.key);
 
     if (authRequestCredentials.decryptedUserKey) {
-      await this.cryptoService.setUserKey(authRequestCredentials.decryptedUserKey);
+      await this.cryptoService.setUserKey(authRequestCredentials.decryptedUserKey, userId);
     } else {
       await this.trySetUserKeyWithMasterKey(userId);
 
@@ -115,7 +115,7 @@ export class AuthRequestLoginStrategy extends LoginStrategy {
     const masterKey = await firstValueFrom(this.masterPasswordService.masterKey$(userId));
     if (masterKey) {
       const userKey = await this.masterPasswordService.decryptUserKeyWithMasterKey(masterKey);
-      await this.cryptoService.setUserKey(userKey);
+      await this.cryptoService.setUserKey(userKey, userId);
     }
   }
 

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

@@ -252,7 +252,7 @@ describe("SsoLoginStrategy", () => {
       expect(deviceTrustService.getDeviceKey).toHaveBeenCalledTimes(1);
       expect(deviceTrustService.decryptUserKeyWithDeviceKey).toHaveBeenCalledTimes(1);
       expect(cryptoSvcSetUserKeySpy).toHaveBeenCalledTimes(1);
-      expect(cryptoSvcSetUserKeySpy).toHaveBeenCalledWith(mockUserKey);
+      expect(cryptoSvcSetUserKeySpy).toHaveBeenCalledWith(mockUserKey, userId);
     });
 
     it("does not set the user key when deviceKey is missing", async () => {
@@ -498,7 +498,7 @@ describe("SsoLoginStrategy", () => {
         undefined,
         undefined,
       );
-      expect(cryptoService.setUserKey).toHaveBeenCalledWith(userKey);
+      expect(cryptoService.setUserKey).toHaveBeenCalledWith(userKey, userId);
     });
   });
 
@@ -554,7 +554,7 @@ describe("SsoLoginStrategy", () => {
         undefined,
         undefined,
       );
-      expect(cryptoService.setUserKey).toHaveBeenCalledWith(userKey);
+      expect(cryptoService.setUserKey).toHaveBeenCalledWith(userKey, userId);
     });
   });
 });

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

@@ -255,6 +255,7 @@ export class SsoLoginStrategy extends LoginStrategy {
         await this.authRequestService.setKeysAfterDecryptingSharedMasterKeyAndHash(
           adminAuthReqResponse,
           adminAuthReqStorable.privateKey,
+          userId,
         );
       } else {
         // if masterPasswordHash is null, we will always receive authReqResponse.key
@@ -262,6 +263,7 @@ export class SsoLoginStrategy extends LoginStrategy {
         await this.authRequestService.setUserKeyAfterDecryptingSharedUserKey(
           adminAuthReqResponse,
           adminAuthReqStorable.privateKey,
+          userId,
         );
       }
 
@@ -321,7 +323,7 @@ export class SsoLoginStrategy extends LoginStrategy {
     );
 
     if (userKey) {
-      await this.cryptoService.setUserKey(userKey);
+      await this.cryptoService.setUserKey(userKey, userId);
     }
   }
 
@@ -337,7 +339,7 @@ export class SsoLoginStrategy extends LoginStrategy {
     }
 
     const userKey = await this.masterPasswordService.decryptUserKeyWithMasterKey(masterKey);
-    await this.cryptoService.setUserKey(userKey);
+    await this.cryptoService.setUserKey(userKey, userId);
   }
 
   protected override async setPrivateKey(

libs/auth/src/common/services/auth-request/auth-request.service.spec.ts

@@ -133,14 +133,18 @@ describe("AuthRequestService", () => {
       cryptoService.setUserKey.mockResolvedValueOnce(undefined);
 
       // Act
-      await sut.setUserKeyAfterDecryptingSharedUserKey(mockAuthReqResponse, mockPrivateKey);
+      await sut.setUserKeyAfterDecryptingSharedUserKey(
+        mockAuthReqResponse,
+        mockPrivateKey,
+        mockUserId,
+      );
 
       // Assert
       expect(sut.decryptPubKeyEncryptedUserKey).toBeCalledWith(
         mockAuthReqResponse.key,
         mockPrivateKey,
       );
-      expect(cryptoService.setUserKey).toBeCalledWith(mockDecryptedUserKey);
+      expect(cryptoService.setUserKey).toBeCalledWith(mockDecryptedUserKey, mockUserId);
     });
   });
 
@@ -169,7 +173,11 @@ describe("AuthRequestService", () => {
       cryptoService.setUserKey.mockResolvedValueOnce(undefined);
 
       // Act
-      await sut.setKeysAfterDecryptingSharedMasterKeyAndHash(mockAuthReqResponse, mockPrivateKey);
+      await sut.setKeysAfterDecryptingSharedMasterKeyAndHash(
+        mockAuthReqResponse,
+        mockPrivateKey,
+        mockUserId,
+      );
 
       // Assert
       expect(sut.decryptPubKeyEncryptedMasterKeyAndHash).toBeCalledWith(
@@ -190,7 +198,7 @@ describe("AuthRequestService", () => {
         undefined,
         undefined,
       );
-      expect(cryptoService.setUserKey).toHaveBeenCalledWith(mockDecryptedUserKey);
+      expect(cryptoService.setUserKey).toHaveBeenCalledWith(mockDecryptedUserKey, mockUserId);
     });
   });
 

libs/auth/src/common/services/auth-request/auth-request.service.ts

@@ -126,17 +126,19 @@ export class AuthRequestService implements AuthRequestServiceAbstraction {
   async setUserKeyAfterDecryptingSharedUserKey(
     authReqResponse: AuthRequestResponse,
     authReqPrivateKey: Uint8Array,
+    userId: UserId,
   ) {
     const userKey = await this.decryptPubKeyEncryptedUserKey(
       authReqResponse.key,
       authReqPrivateKey,
     );
-    await this.cryptoService.setUserKey(userKey);
+    await this.cryptoService.setUserKey(userKey, userId);
   }
 
   async setKeysAfterDecryptingSharedMasterKeyAndHash(
     authReqResponse: AuthRequestResponse,
     authReqPrivateKey: Uint8Array,
+    userId: UserId,
   ) {
     const { masterKey, masterKeyHash } = await this.decryptPubKeyEncryptedMasterKeyAndHash(
       authReqResponse.key,
@@ -148,11 +150,10 @@ export class AuthRequestService implements AuthRequestServiceAbstraction {
     const userKey = await this.masterPasswordService.decryptUserKeyWithMasterKey(masterKey);
 
     // Set masterKey + masterKeyHash in state after decryption (in case decryption fails)
-    const userId = (await firstValueFrom(this.accountService.activeAccount$)).id;
     await this.masterPasswordService.setMasterKey(masterKey, userId);
     await this.masterPasswordService.setMasterKeyHash(masterKeyHash, userId);
 
-    await this.cryptoService.setUserKey(userKey);
+    await this.cryptoService.setUserKey(userKey, userId);
   }
 
   // Decryption helpers

libs/common/src/platform/services/crypto.service.spec.ts

@@ -263,6 +263,12 @@ describe("cryptoService", () => {
       await expect(cryptoService.setUserKey(null, mockUserId)).rejects.toThrow("No key provided.");
     });
 
+    it("throws if userId is null", async () => {
+      await expect(cryptoService.setUserKey(mockUserKey, null)).rejects.toThrow(
+        "No userId provided.",
+      );
+    });
+
     describe("Pin Key refresh", () => {
       const mockPinKeyEncryptedUserKey = new EncString(
         "2.AAAw2vTUePO+CCyokcIfVw==|DTBNlJ5yVsV2Bsk3UU3H6Q==|YvFBff5gxWqM+UsFB6BKimKxhC32AtjF3IStpU1Ijwg=",

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

@@ -89,12 +89,16 @@ export class CryptoService implements CryptoServiceAbstraction {
     );
   }
 
-  async setUserKey(key: UserKey, userId?: UserId): Promise<void> {
+  async setUserKey(key: UserKey, userId: UserId): Promise<void> {
     if (key == null) {
       throw new Error("No key provided. Lock the user to clear the key");
     }
+    if (userId == null) {
+      throw new Error("No userId provided.");
+    }
+
     // Set userId to ensure we have one for the account status update
-    [userId, key] = await this.stateProvider.setUserState(USER_KEY, key, userId);
+    await this.stateProvider.setUserState(USER_KEY, key, userId);
     await this.stateProvider.setUserState(USER_EVER_HAD_USER_KEY, true, userId);
 
     await this.storeAdditionalKeys(key, userId);
@@ -701,13 +705,7 @@ export class CryptoService implements CryptoServiceAbstraction {
    * @param key The user key
    * @param userId The desired user
    */
-  protected async storeAdditionalKeys(key: UserKey, userId?: UserId) {
+  protected async storeAdditionalKeys(key: UserKey, userId: UserId) {
-    userId ??= await firstValueFrom(this.stateProvider.activeUserId$);
-
-    if (userId == null) {
-      throw new Error("Cannot store additional keys, no user Id resolved.");
-    }
-
     const storeAuto = await this.shouldStoreKey(KeySuffixOptions.Auto, userId);
     if (storeAuto) {
       await this.stateService.setUserKeyAutoUnlock(key.keyB64, { userId: userId });