[PM-6727] Implement new setUserKeys function in CryptoService (#10655)

* feat: implement new `setUserKeys` function

* feat: add explicit error for failed decryption

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

@@ -19,6 +19,12 @@ import { EncArrayBuffer } from "../models/domain/enc-array-buffer";
 import { EncString } from "../models/domain/enc-string";
 import { SymmetricCryptoKey } from "../models/domain/symmetric-crypto-key";
 
+export class UserPrivateKeyDecryptionFailedError extends Error {
+  constructor() {
+    super("Failed to decrypt the user's private key.");
+  }
+}
+
 /**
  * An object containing all the users key needed to decrypt a users personal and organization vaults.
  */
@@ -58,6 +64,20 @@ export abstract class CryptoService {
    * @param userId The desired user
    */
   abstract setUserKey(key: UserKey, userId?: string): Promise<void>;
+  /**
+   * Sets the provided user keys and stores any other necessary versions
+   * (such as auto, biometrics, or pin).
+   * Also sets the user's encrypted private key in storage and
+   * clears the decrypted private key from memory
+   * Note: does not clear the private key if null is provided
+   *
+   * @throws Error when userKey, encPrivateKey or userId is null
+   * @throws UserPrivateKeyDecryptionFailedError when the userKey cannot decrypt encPrivateKey
+   * @param userKey The user key to set
+   * @param encPrivateKey An encrypted private key
+   * @param userId The desired user
+   */
+  abstract setUserKeys(userKey: UserKey, encPrivateKey: string, userId: UserId): Promise<void>;
   /**
    * Gets the user key from memory and sets it again,
    * kicking off a refresh of any additional keys

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

@@ -20,6 +20,7 @@ import { OrganizationId, UserId } from "../../types/guid";
 import { UserKey, MasterKey } from "../../types/key";
 import { VaultTimeoutStringType } from "../../types/vault-timeout.type";
 import { CryptoFunctionService } from "../abstractions/crypto-function.service";
+import { UserPrivateKeyDecryptionFailedError } from "../abstractions/crypto.service";
 import { EncryptService } from "../abstractions/encrypt.service";
 import { KeyGenerationService } from "../abstractions/key-generation.service";
 import { LogService } from "../abstractions/log.service";
@@ -318,6 +319,70 @@ describe("cryptoService", () => {
     });
   });
 
+  describe("setUserKeys", () => {
+    let mockUserKey: UserKey;
+    let mockEncPrivateKey: EncryptedString;
+    let everHadUserKeyState: FakeSingleUserState<boolean>;
+
+    beforeEach(() => {
+      const mockRandomBytes = new Uint8Array(64) as CsprngArray;
+      mockUserKey = new SymmetricCryptoKey(mockRandomBytes) as UserKey;
+      mockEncPrivateKey = new SymmetricCryptoKey(mockRandomBytes).toString() as EncryptedString;
+      everHadUserKeyState = stateProvider.singleUser.getFake(mockUserId, USER_EVER_HAD_USER_KEY);
+
+      // Initialize storage
+      everHadUserKeyState.nextState(null);
+
+      // Mock private key decryption
+      encryptService.decryptToBytes.mockResolvedValue(mockRandomBytes);
+    });
+
+    it("throws if userKey is null", async () => {
+      await expect(cryptoService.setUserKeys(null, mockEncPrivateKey, mockUserId)).rejects.toThrow(
+        "No userKey provided.",
+      );
+    });
+
+    it("throws if encPrivateKey is null", async () => {
+      await expect(cryptoService.setUserKeys(mockUserKey, null, mockUserId)).rejects.toThrow(
+        "No encPrivateKey provided.",
+      );
+    });
+
+    it("throws if userId is null", async () => {
+      await expect(cryptoService.setUserKeys(mockUserKey, mockEncPrivateKey, null)).rejects.toThrow(
+        "No userId provided.",
+      );
+    });
+
+    it("throws if encPrivateKey cannot be decrypted with the userKey", async () => {
+      encryptService.decryptToBytes.mockResolvedValue(null);
+
+      await expect(
+        cryptoService.setUserKeys(mockUserKey, mockEncPrivateKey, mockUserId),
+      ).rejects.toThrow(UserPrivateKeyDecryptionFailedError);
+    });
+
+    // We already have tests for setUserKey, so we just need to test that the correct methods are called
+    it("calls setUserKey with the userKey and userId", async () => {
+      const setUserKeySpy = jest.spyOn(cryptoService, "setUserKey");
+
+      await cryptoService.setUserKeys(mockUserKey, mockEncPrivateKey, mockUserId);
+
+      expect(setUserKeySpy).toHaveBeenCalledWith(mockUserKey, mockUserId);
+    });
+
+    // We already have tests for setPrivateKey, so we just need to test that the correct methods are called
+    // TODO: Move those tests into here since `setPrivateKey` will be converted to a private method
+    it("calls setPrivateKey with the encPrivateKey and userId", async () => {
+      const setEncryptedPrivateKeySpy = jest.spyOn(cryptoService, "setPrivateKey");
+
+      await cryptoService.setUserKeys(mockUserKey, mockEncPrivateKey, mockUserId);
+
+      expect(setEncryptedPrivateKeySpy).toHaveBeenCalledWith(mockEncPrivateKey, mockUserId);
+    });
+  });
+
   describe("clearKeys", () => {
     it("resolves active user id when called with no user id", async () => {
       let callCount = 0;

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

@@ -38,6 +38,7 @@ import { CryptoFunctionService } from "../abstractions/crypto-function.service";
 import {
   CipherDecryptionKeys,
   CryptoService as CryptoServiceAbstraction,
+  UserPrivateKeyDecryptionFailedError,
 } from "../abstractions/crypto.service";
 import { EncryptService } from "../abstractions/encrypt.service";
 import { KeyGenerationService } from "../abstractions/key-generation.service";
@@ -104,6 +105,30 @@ export class CryptoService implements CryptoServiceAbstraction {
     await this.storeAdditionalKeys(key, userId);
   }
 
+  async setUserKeys(
+    userKey: UserKey,
+    encPrivateKey: EncryptedString,
+    userId: UserId,
+  ): Promise<void> {
+    if (userKey == null) {
+      throw new Error("No userKey provided. Lock the user to clear the key");
+    }
+    if (encPrivateKey == null) {
+      throw new Error("No encPrivateKey provided.");
+    }
+    if (userId == null) {
+      throw new Error("No userId provided.");
+    }
+
+    const decryptedPrivateKey = await this.decryptPrivateKey(encPrivateKey, userKey);
+    if (decryptedPrivateKey == null) {
+      throw new UserPrivateKeyDecryptionFailedError();
+    }
+
+    await this.setUserKey(userKey, userId);
+    await this.setPrivateKey(encPrivateKey, userId);
+  }
+
   async refreshAdditionalKeys(): Promise<void> {
     const activeUserId = await firstValueFrom(this.stateProvider.activeUserId$);