[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/common/src/auth/abstractions/webauthn/webauthn-login-prf-key.service.abstraction.ts

@@ -11,7 +11,7 @@ export abstract class WebAuthnLoginPrfKeyServiceAbstraction {
 
   /**
    * Create a symmetric key from the PRF-output by stretching it.
-   * This should be used as `ExternalKey` with `RotateableKeySet`.
+   * This should be used as `UpstreamKey` with `RotateableKeySet`.
    */
   abstract createSymmetricKeyFromPrf(prf: ArrayBuffer): Promise<PrfKey>;
 }

libs/common/src/auth/models/request/webauthn-rotate-credential.request.ts

@@ -1,10 +1,7 @@
 // FIXME: Update this file to be type safe and remove this and next line
 // @ts-strict-ignore
-
-// FIXME: remove `src` and fix import
-// eslint-disable-next-line no-restricted-imports
-import { RotateableKeySet } from "../../../../../auth/src/common/models";
 import { EncString } from "../../../key-management/crypto/models/enc-string";
+import { RotateableKeySet } from "../../../key-management/keys/models/rotateable-key-set";
 
 export class WebauthnRotateCredentialRequest {
   id: string;

libs/common/src/auth/models/response/protected-device.response.ts

@@ -2,12 +2,9 @@
 // @ts-strict-ignore
 import { Jsonify } from "type-fest";
 
-// This import has been flagged as unallowed for this class. It may be involved in a circular dependency loop.
-// eslint-disable-next-line no-restricted-imports
-import { RotateableKeySet } from "@bitwarden/auth/common";
-
 import { DeviceType } from "../../../enums";
 import { EncString } from "../../../key-management/crypto/models/enc-string";
+import { RotateableKeySet } from "../../../key-management/keys/models/rotateable-key-set";
 import { BaseResponse } from "../../../models/response/base.response";
 
 export class ProtectedDeviceResponse extends BaseResponse {

libs/common/src/key-management/device-trust/services/device-trust.service.implementation.ts

@@ -4,7 +4,7 @@ import { firstValueFrom, map, Observable, Subject } from "rxjs";
 
 // This import has been flagged as unallowed for this class. It may be involved in a circular dependency loop.
 // eslint-disable-next-line no-restricted-imports
-import { RotateableKeySet, UserDecryptionOptionsServiceAbstraction } from "@bitwarden/auth/common";
+import { UserDecryptionOptionsServiceAbstraction } from "@bitwarden/auth/common";
 // This import has been flagged as unallowed for this class. It may be involved in a circular dependency loop.
 // eslint-disable-next-line no-restricted-imports
 import { KeyService } from "@bitwarden/key-management";
@@ -33,6 +33,7 @@ import { KeyGenerationService } from "../../crypto";
 import { CryptoFunctionService } from "../../crypto/abstractions/crypto-function.service";
 import { EncryptService } from "../../crypto/abstractions/encrypt.service";
 import { EncString } from "../../crypto/models/enc-string";
+import { RotateableKeySet } from "../../keys/models/rotateable-key-set";
 import { DeviceTrustServiceAbstraction } from "../abstractions/device-trust.service.abstraction";
 
 /** Uses disk storage so that the device key can persist after log out and tab removal. */
@@ -145,7 +146,7 @@ export class DeviceTrustService implements DeviceTrustServiceAbstraction {
     }
 
     // Attempt to get user key
-    const userKey: UserKey = await this.keyService.getUserKey(userId);
+    const userKey = await firstValueFrom(this.keyService.userKey$(userId));
 
     // If user key is not found, throw error
     if (!userKey) {
@@ -240,7 +241,7 @@ export class DeviceTrustService implements DeviceTrustServiceAbstraction {
 
           const request = new OtherDeviceKeysUpdateRequest();
           request.encryptedPublicKey = newRotateableKeySet.encryptedPublicKey.encryptedString;
-          request.encryptedUserKey = newRotateableKeySet.encryptedUserKey.encryptedString;
+          request.encryptedUserKey = newRotateableKeySet.encapsulatedDownstreamKey.encryptedString;
           request.deviceId = device.id;
           return request;
         })

libs/common/src/key-management/device-trust/services/device-trust.service.spec.ts

@@ -366,7 +366,6 @@ describe("deviceTrustService", () => {
 
       let makeDeviceKeySpy: jest.SpyInstance;
       let rsaGenerateKeyPairSpy: jest.SpyInstance;
-      let cryptoSvcGetUserKeySpy: jest.SpyInstance;
       let cryptoSvcRsaEncryptSpy: jest.SpyInstance;
       let encryptServiceWrapDecapsulationKeySpy: jest.SpyInstance;
       let encryptServiceWrapEncapsulationKeySpy: jest.SpyInstance;
@@ -402,6 +401,8 @@ describe("deviceTrustService", () => {
           "mockDeviceKeyEncryptedDevicePrivateKey",
         );
 
+        keyService.userKey$.mockReturnValue(of(mockUserKey));
+
         // TypeScript will allow calling private methods if the object is of type 'any'
         makeDeviceKeySpy = jest
           .spyOn(deviceTrustService as any, "makeDeviceKey")
@@ -411,10 +412,6 @@ describe("deviceTrustService", () => {
           .spyOn(cryptoFunctionService, "rsaGenerateKeyPair")
           .mockResolvedValue(mockDeviceRsaKeyPair);
 
-        cryptoSvcGetUserKeySpy = jest
-          .spyOn(keyService, "getUserKey")
-          .mockResolvedValue(mockUserKey);
-
         cryptoSvcRsaEncryptSpy = jest
           .spyOn(encryptService, "encapsulateKeyUnsigned")
           .mockResolvedValue(mockDevicePublicKeyEncryptedUserKey);
@@ -448,7 +445,7 @@ describe("deviceTrustService", () => {
 
         expect(makeDeviceKeySpy).toHaveBeenCalledTimes(1);
         expect(rsaGenerateKeyPairSpy).toHaveBeenCalledTimes(1);
-        expect(cryptoSvcGetUserKeySpy).toHaveBeenCalledTimes(1);
+        expect(keyService.userKey$).toHaveBeenCalledTimes(1);
 
         expect(cryptoSvcRsaEncryptSpy).toHaveBeenCalledTimes(1);
 
@@ -473,18 +470,13 @@ describe("deviceTrustService", () => {
       });
 
       it("throws specific error if user key is not found", async () => {
-        // setup the spy to return null
-        cryptoSvcGetUserKeySpy.mockResolvedValue(null);
+        keyService.userKey$.mockReturnValueOnce(of(null));
         // check if the expected error is thrown
         await expect(deviceTrustService.trustDevice(mockUserId)).rejects.toThrow(
           "User symmetric key not found",
         );
 
-        // reset the spy
-        cryptoSvcGetUserKeySpy.mockReset();
-
-        // setup the spy to return undefined
-        cryptoSvcGetUserKeySpy.mockResolvedValue(undefined);
+        keyService.userKey$.mockReturnValueOnce(of(undefined));
         // check if the expected error is thrown
         await expect(deviceTrustService.trustDevice(mockUserId)).rejects.toThrow(
           "User symmetric key not found",
@@ -502,11 +494,6 @@ describe("deviceTrustService", () => {
           spy: () => rsaGenerateKeyPairSpy,
           errorText: "rsaGenerateKeyPair error",
         },
-        {
-          method: "getUserKey",
-          spy: () => cryptoSvcGetUserKeySpy,
-          errorText: "getUserKey error",
-        },
         {
           method: "rsaEncrypt",
           spy: () => cryptoSvcRsaEncryptSpy,

libs/common/src/key-management/keys/models/rotateable-key-set.ts

@@ -0,0 +1,34 @@
+import { SymmetricCryptoKey } from "../../../platform/models/domain/symmetric-crypto-key";
+import { PrfKey } from "../../../types/key";
+import { EncString } from "../../crypto/models/enc-string";
+
+declare const tag: unique symbol;
+
+/**
+ * A set of keys where a symmetric `DownstreamKey` is protected by an encrypted public/private key-pair.
+ * The `DownstreamKey` is used to encrypt/decrypt data, while the public/private key-pair is
+ * used to rotate the `DownstreamKey`.
+ *
+ * The `PrivateKey` is protected by an `UpstreamKey`, such as a `DeviceKey`, or `PrfKey`,
+ * and the `PublicKey` is protected by the `DownstreamKey`. This setup allows:
+ *
+ *   - Access to `DownstreamKey` by knowing the `UpstreamKey`
+ *   - Rotation to a new `DownstreamKey` by knowing the current `DownstreamKey`,
+ *     without needing access to the `UpstreamKey`
+ */
+export class RotateableKeySet<UpstreamKey extends SymmetricCryptoKey = SymmetricCryptoKey> {
+  private readonly [tag]!: UpstreamKey;
+
+  constructor(
+    /** `DownstreamKey` protected by publicKey */
+    readonly encapsulatedDownstreamKey: EncString,
+
+    /** DownstreamKey encrypted PublicKey */
+    readonly encryptedPublicKey: EncString,
+
+    /** UpstreamKey encrypted PrivateKey */
+    readonly encryptedPrivateKey?: EncString,
+  ) {}
+}
+
+export type PrfKeySet = RotateableKeySet<PrfKey>;

libs/common/src/key-management/keys/services/abstractions/rotateable-key-set.service.ts

@@ -0,0 +1,30 @@
+import { SymmetricCryptoKey } from "../../../../platform/models/domain/symmetric-crypto-key";
+import { RotateableKeySet } from "../../models/rotateable-key-set";
+
+export abstract class RotateableKeySetService {
+  /**
+   * Create a new rotatable key set for the provided downstreamKey, using the provided upstream key.
+   * For more information on rotatable key sets, see {@link RotateableKeySet}
+   * @param upstreamKey The `UpstreamKey` used to encrypt {@link RotateableKeySet.encryptedPrivateKey}
+   * @param downstreamKey The symmetric key to be contained within the `RotateableKeySet`.
+   * @returns RotateableKeySet containing the provided symmetric downstreamKey.
+   */
+  abstract createKeySet<UpstreamKey extends SymmetricCryptoKey>(
+    upstreamKey: UpstreamKey,
+    downstreamKey: SymmetricCryptoKey,
+  ): Promise<RotateableKeySet<UpstreamKey>>;
+
+  /**
+   * Rotates the provided `RotateableKeySet` with the new key.
+   *
+   * @param keySet The current `RotateableKeySet` to be rotated.
+   * @param oldDownstreamKey The current downstreamKey used to decrypt the `PublicKey`.
+   * @param newDownstreamKey The new downstreamKey to encrypt the `PublicKey`.
+   * @returns The updated `RotateableKeySet` that contains the new downstreamKey.
+   */
+  abstract rotateKeySet<UpstreamKey extends SymmetricCryptoKey>(
+    keySet: RotateableKeySet<UpstreamKey>,
+    oldDownstreamKey: SymmetricCryptoKey,
+    newDownstreamKey: SymmetricCryptoKey,
+  ): Promise<RotateableKeySet<UpstreamKey>>;
+}

libs/common/src/key-management/keys/services/default-rotateable-key-set.service.spec.ts

@@ -0,0 +1,185 @@
+import { mock, MockProxy } from "jest-mock-extended";
+
+// This import has been flagged as unallowed for this class. It may be involved in a circular dependency loop.
+// eslint-disable-next-line no-restricted-imports
+import { KeyService } from "@bitwarden/key-management";
+
+import { Utils } from "../../../platform/misc/utils";
+import { SymmetricCryptoKey } from "../../../platform/models/domain/symmetric-crypto-key";
+import { EncryptService } from "../../crypto/abstractions/encrypt.service";
+import { EncString } from "../../crypto/models/enc-string";
+import { RotateableKeySet } from "../models/rotateable-key-set";
+
+import { DefaultRotateableKeySetService } from "./default-rotateable-key-set.service";
+
+describe("DefaultRotateableKeySetService", () => {
+  let keyService!: MockProxy<KeyService>;
+  let encryptService!: MockProxy<EncryptService>;
+  let service!: DefaultRotateableKeySetService;
+
+  beforeEach(() => {
+    keyService = mock<KeyService>();
+    encryptService = mock<EncryptService>();
+    service = new DefaultRotateableKeySetService(keyService, encryptService);
+  });
+
+  describe("createKeySet", () => {
+    test.each([null, undefined])(
+      "throws error when downstreamKey parameter is %s",
+      async (downstreamKey) => {
+        const externalKey = createSymmetricKey();
+        await expect(service.createKeySet(externalKey, downstreamKey as any)).rejects.toThrow(
+          "failed to create key set: downstreamKey is required",
+        );
+      },
+    );
+
+    test.each([null, undefined])(
+      "throws error when upstreamKey parameter is %s",
+      async (upstreamKey) => {
+        const userKey = createSymmetricKey();
+        await expect(service.createKeySet(upstreamKey as any, userKey)).rejects.toThrow(
+          "failed to create key set: upstreamKey is required",
+        );
+      },
+    );
+
+    it("should create a new key set", async () => {
+      const externalKey = createSymmetricKey();
+      const userKey = createSymmetricKey();
+      const encryptedUserKey = new EncString("encryptedUserKey");
+      const encryptedPublicKey = new EncString("encryptedPublicKey");
+      const encryptedPrivateKey = new EncString("encryptedPrivateKey");
+      keyService.makeKeyPair.mockResolvedValue(["publicKey", encryptedPrivateKey]);
+      encryptService.encapsulateKeyUnsigned.mockResolvedValue(encryptedUserKey);
+      encryptService.wrapEncapsulationKey.mockResolvedValue(encryptedPublicKey);
+
+      const result = await service.createKeySet(externalKey, userKey);
+
+      expect(result).toEqual(
+        new RotateableKeySet(encryptedUserKey, encryptedPublicKey, encryptedPrivateKey),
+      );
+      expect(keyService.makeKeyPair).toHaveBeenCalledWith(externalKey);
+      expect(encryptService.encapsulateKeyUnsigned).toHaveBeenCalledWith(
+        userKey,
+        Utils.fromB64ToArray("publicKey"),
+      );
+      expect(encryptService.wrapEncapsulationKey).toHaveBeenCalledWith(
+        Utils.fromB64ToArray("publicKey"),
+        userKey,
+      );
+    });
+  });
+
+  describe("rotateKeySet", () => {
+    const keySet = new RotateableKeySet(
+      new EncString("encUserKey"),
+      new EncString("encPublicKey"),
+      new EncString("encPrivateKey"),
+    );
+    const dataValidationTests = [
+      {
+        keySet: null as any as RotateableKeySet,
+        oldDownstreamKey: createSymmetricKey(),
+        newDownstreamKey: createSymmetricKey(),
+        expectedError: "failed to rotate key set: keySet is required",
+      },
+      {
+        keySet: undefined as any as RotateableKeySet,
+        oldDownstreamKey: createSymmetricKey(),
+        newDownstreamKey: createSymmetricKey(),
+        expectedError: "failed to rotate key set: keySet is required",
+      },
+      {
+        keySet: keySet,
+        oldDownstreamKey: null,
+        newDownstreamKey: createSymmetricKey(),
+        expectedError: "failed to rotate key set: oldDownstreamKey is required",
+      },
+      {
+        keySet: keySet,
+        oldDownstreamKey: undefined,
+        newDownstreamKey: createSymmetricKey(),
+        expectedError: "failed to rotate key set: oldDownstreamKey is required",
+      },
+      {
+        keySet: keySet,
+        oldDownstreamKey: createSymmetricKey(),
+        newDownstreamKey: null,
+        expectedError: "failed to rotate key set: newDownstreamKey is required",
+      },
+      {
+        keySet: keySet,
+        oldDownstreamKey: createSymmetricKey(),
+        newDownstreamKey: undefined,
+        expectedError: "failed to rotate key set: newDownstreamKey is required",
+      },
+    ];
+
+    test.each(dataValidationTests)(
+      "should throw error when required parameter is missing",
+      async ({ keySet, oldDownstreamKey, newDownstreamKey, expectedError }) => {
+        await expect(
+          service.rotateKeySet(keySet, oldDownstreamKey as any, newDownstreamKey as any),
+        ).rejects.toThrow(expectedError);
+      },
+    );
+
+    it("throws an error if the public key cannot be decrypted", async () => {
+      const oldDownstreamKey = createSymmetricKey();
+      const newDownstreamKey = createSymmetricKey();
+
+      encryptService.unwrapEncapsulationKey.mockResolvedValue(null as any);
+
+      await expect(
+        service.rotateKeySet(keySet, oldDownstreamKey, newDownstreamKey),
+      ).rejects.toThrow("failed to rotate key set: could not decrypt public key");
+
+      expect(encryptService.unwrapEncapsulationKey).toHaveBeenCalledWith(
+        keySet.encryptedPublicKey,
+        oldDownstreamKey,
+      );
+
+      expect(encryptService.wrapEncapsulationKey).not.toHaveBeenCalled();
+      expect(encryptService.encapsulateKeyUnsigned).not.toHaveBeenCalled();
+    });
+
+    it("rotates the key set", async () => {
+      const oldDownstreamKey = createSymmetricKey();
+      const newDownstreamKey = new SymmetricCryptoKey(new Uint8Array(64));
+      const publicKey = Utils.fromB64ToArray("decryptedPublicKey");
+      const newEncryptedPublicKey = new EncString("newEncPublicKey");
+      const newEncryptedRotateableKey = new EncString("newEncUserKey");
+
+      encryptService.unwrapEncapsulationKey.mockResolvedValue(publicKey);
+      encryptService.wrapEncapsulationKey.mockResolvedValue(newEncryptedPublicKey);
+      encryptService.encapsulateKeyUnsigned.mockResolvedValue(newEncryptedRotateableKey);
+
+      const result = await service.rotateKeySet(keySet, oldDownstreamKey, newDownstreamKey);
+
+      expect(result).toEqual(
+        new RotateableKeySet(
+          newEncryptedRotateableKey,
+          newEncryptedPublicKey,
+          keySet.encryptedPrivateKey,
+        ),
+      );
+      expect(encryptService.unwrapEncapsulationKey).toHaveBeenCalledWith(
+        keySet.encryptedPublicKey,
+        oldDownstreamKey,
+      );
+      expect(encryptService.wrapEncapsulationKey).toHaveBeenCalledWith(publicKey, newDownstreamKey);
+      expect(encryptService.encapsulateKeyUnsigned).toHaveBeenCalledWith(
+        newDownstreamKey,
+        publicKey,
+      );
+    });
+  });
+});
+
+function createSymmetricKey() {
+  const key = Utils.fromB64ToArray(
+    "1h-TuPwSbX5qoX0aVgjmda_Lfq85qAcKssBlXZnPIsQC3HNDGIecunYqXhJnp55QpdXRh-egJiLH3a0wqlVQsQ",
+  );
+  return new SymmetricCryptoKey(key);
+}

libs/common/src/key-management/keys/services/default-rotateable-key-set.service.ts

@@ -0,0 +1,83 @@
+// This import has been flagged as unallowed for this class. It may be involved in a circular dependency loop.
+// eslint-disable-next-line no-restricted-imports
+import { KeyService } from "@bitwarden/key-management";
+
+import { Utils } from "../../../platform/misc/utils";
+import { SymmetricCryptoKey } from "../../../platform/models/domain/symmetric-crypto-key";
+import { EncryptService } from "../../crypto/abstractions/encrypt.service";
+import { RotateableKeySet } from "../models/rotateable-key-set";
+
+import { RotateableKeySetService } from "./abstractions/rotateable-key-set.service";
+
+export class DefaultRotateableKeySetService implements RotateableKeySetService {
+  constructor(
+    private keyService: KeyService,
+    private encryptService: EncryptService,
+  ) {}
+
+  async createKeySet<UpstreamKey extends SymmetricCryptoKey>(
+    upstreamKey: UpstreamKey,
+    downstreamKey: SymmetricCryptoKey,
+  ): Promise<RotateableKeySet<UpstreamKey>> {
+    if (!upstreamKey) {
+      throw new Error("failed to create key set: upstreamKey is required");
+    }
+    if (!downstreamKey) {
+      throw new Error("failed to create key set: downstreamKey is required");
+    }
+
+    const [publicKey, encryptedPrivateKey] = await this.keyService.makeKeyPair(upstreamKey);
+
+    const rawPublicKey = Utils.fromB64ToArray(publicKey);
+    const encryptedRotateableKey = await this.encryptService.encapsulateKeyUnsigned(
+      downstreamKey,
+      rawPublicKey,
+    );
+    const encryptedPublicKey = await this.encryptService.wrapEncapsulationKey(
+      rawPublicKey,
+      downstreamKey,
+    );
+    return new RotateableKeySet(encryptedRotateableKey, encryptedPublicKey, encryptedPrivateKey);
+  }
+
+  async rotateKeySet<UpstreamKey extends SymmetricCryptoKey>(
+    keySet: RotateableKeySet<UpstreamKey>,
+    oldDownstreamKey: SymmetricCryptoKey,
+    newDownstreamKey: SymmetricCryptoKey,
+  ): Promise<RotateableKeySet<UpstreamKey>> {
+    // validate parameters
+    if (!keySet) {
+      throw new Error("failed to rotate key set: keySet is required");
+    }
+    if (!oldDownstreamKey) {
+      throw new Error("failed to rotate key set: oldDownstreamKey is required");
+    }
+    if (!newDownstreamKey) {
+      throw new Error("failed to rotate key set: newDownstreamKey is required");
+    }
+
+    const publicKey = await this.encryptService.unwrapEncapsulationKey(
+      keySet.encryptedPublicKey,
+      oldDownstreamKey,
+    );
+    if (publicKey == null) {
+      throw new Error("failed to rotate key set: could not decrypt public key");
+    }
+    const newEncryptedPublicKey = await this.encryptService.wrapEncapsulationKey(
+      publicKey,
+      newDownstreamKey,
+    );
+    const newEncryptedRotateableKey = await this.encryptService.encapsulateKeyUnsigned(
+      newDownstreamKey,
+      publicKey,
+    );
+
+    const newRotateableKeySet = new RotateableKeySet<UpstreamKey>(
+      newEncryptedRotateableKey,
+      newEncryptedPublicKey,
+      keySet.encryptedPrivateKey,
+    );
+
+    return newRotateableKeySet;
+  }
+}