[PM-24107] Migrate KM's usage of getUserKey from the key service (#17117)

* Remove internal use of getUserKey in the key service

* Move ownership of RotateableKeySet and remove usage of getUserKey

* Add input validation to createKeySet

libs/key-management/src/abstractions/key.service.ts

@@ -166,16 +166,16 @@ export abstract class KeyService {
    */
   abstract makeMasterKey(password: string, email: string, kdfConfig: KdfConfig): Promise<MasterKey>;
   /**
-   * Encrypts the existing (or provided) user key with the
-   * provided master key
+   * Encrypts the provided user key with the provided master key.
    * @deprecated Interacting with the master key directly is prohibited. Use a high level function from MasterPasswordService instead.
    * @param masterKey The user's master key
    * @param userKey The user key
+   * @throws Error when userKey or masterKey is null/undefined.
    * @returns The user key and the master key protected version of it
    */
   abstract encryptUserKeyWithMasterKey(
     masterKey: MasterKey,
-    userKey?: UserKey,
+    userKey: UserKey,
   ): Promise<[UserKey, EncString]>;
   /**
    * Creates a master password hash from the user's master password. Can

libs/key-management/src/key.service.spec.ts

@@ -1357,6 +1357,51 @@ describe("keyService", () => {
     });
   });
 
+  describe("encryptUserKeyWithMasterKey", () => {
+    const mockMasterKey = makeSymmetricCryptoKey<MasterKey>(32);
+    const mockUserKey = makeSymmetricCryptoKey<UserKey>(64);
+
+    test.each([null as unknown as MasterKey, undefined as unknown as MasterKey])(
+      "throws when the provided master key is %s",
+      async (key) => {
+        await expect(keyService.encryptUserKeyWithMasterKey(key, mockUserKey)).rejects.toThrow(
+          "masterKey is required.",
+        );
+      },
+    );
+
+    test.each([null as unknown as UserKey, undefined as unknown as UserKey])(
+      "throws when the provided userKey key is %s",
+      async (key) => {
+        await expect(keyService.encryptUserKeyWithMasterKey(mockMasterKey, key)).rejects.toThrow(
+          "userKey is required.",
+        );
+      },
+    );
+
+    it("throws with invalid master key size", async () => {
+      const invalidMasterKey = new SymmetricCryptoKey(new Uint8Array(78)) as MasterKey;
+
+      await expect(
+        keyService.encryptUserKeyWithMasterKey(invalidMasterKey, mockUserKey),
+      ).rejects.toThrow("Invalid key size.");
+    });
+
+    it("encrypts the user key with the master key", async () => {
+      const mockEncryptedUserKey = makeEncString("encryptedUserKey");
+
+      encryptService.wrapSymmetricKey.mockResolvedValue(mockEncryptedUserKey);
+      const stretchedMasterKey = new SymmetricCryptoKey(new Uint8Array(64));
+      keyGenerationService.stretchKey.mockResolvedValue(stretchedMasterKey);
+
+      const result = await keyService.encryptUserKeyWithMasterKey(mockMasterKey, mockUserKey);
+
+      expect(encryptService.wrapSymmetricKey).toHaveBeenCalledWith(mockUserKey, stretchedMasterKey);
+      expect(result[0]).toBe(mockUserKey);
+      expect(result[1]).toBe(mockEncryptedUserKey);
+    });
+  });
+
   describe("makeKeyPair", () => {
     test.each([null as unknown as SymmetricCryptoKey, undefined as unknown as SymmetricCryptoKey])(
       "throws when the provided key is %s",

libs/key-management/src/key.service.ts

@@ -166,6 +166,9 @@ export class DefaultKeyService implements KeyServiceAbstraction {
     return this.stateProvider.getUserState$(USER_KEY, userId);
   }
 
+  /**
+   * @deprecated Use {@link userKey$} with a required {@link UserId} instead.
+   */
   async getUserKey(userId?: UserId): Promise<UserKey> {
     const userKey = await firstValueFrom(this.stateProvider.getUserState$(USER_KEY, userId));
     return userKey;
@@ -298,9 +301,15 @@ export class DefaultKeyService implements KeyServiceAbstraction {
    */
   async encryptUserKeyWithMasterKey(
     masterKey: MasterKey,
-    userKey?: UserKey,
+    userKey: UserKey,
   ): Promise<[UserKey, EncString]> {
-    userKey ||= await this.getUserKey();
+    if (masterKey == null) {
+      throw new Error("masterKey is required.");
+    }
+    if (userKey == null) {
+      throw new Error("userKey is required.");
+    }
+
     return await this.buildProtectedSymmetricKey(masterKey, userKey);
   }
 
@@ -630,7 +639,7 @@ export class DefaultKeyService implements KeyServiceAbstraction {
     }
 
     // Verify user key doesn't exist
-    const existingUserKey = await this.getUserKey(userId);
+    const existingUserKey = await firstValueFrom(this.userKey$(userId));
 
     if (existingUserKey != null) {
       this.logService.error("Tried to initialize account with existing user key.");