From 88eeebb0844028206dfb57c67226374366085252 Mon Sep 17 00:00:00 2001
From: Matt Gibson <mgibson@bitwarden.com>
Date: Sat, 27 Apr 2024 16:32:34 -0400
Subject: [PATCH] Use a passed in key in derivation so we can validate other
 keys (#8954)

* Use a passed in key in derivation so we can validate other keys

* Fix user key type tests
---
 .../src/platform/services/crypto.service.ts      | 16 +++++++++++++---
 .../services/key-state/user-key.state.spec.ts    | 16 ++++++----------
 .../services/key-state/user-key.state.ts         |  8 ++++----
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/libs/common/src/platform/services/crypto.service.ts b/libs/common/src/platform/services/crypto.service.ts
index 798173f513..713fe7d230 100644
--- a/libs/common/src/platform/services/crypto.service.ts
+++ b/libs/common/src/platform/services/crypto.service.ts
@@ -100,7 +100,7 @@ export class CryptoService implements CryptoServiceAbstraction {
       USER_PRIVATE_KEY,
       {
         encryptService: this.encryptService,
-        cryptoService: this,
+        getUserKey: (userId) => this.getUserKey(userId),
       },
     );
     this.activeUserPrivateKey$ = this.activeUserPrivateKeyState.state$; // may be null
@@ -738,13 +738,23 @@ export class CryptoService implements CryptoServiceAbstraction {
       // Can decrypt private key
       const privateKey = await USER_PRIVATE_KEY.derive([userId, encPrivateKey], {
         encryptService: this.encryptService,
-        cryptoService: this,
+        getUserKey: () => Promise.resolve(key),
       });
 
+      if (privateKey == null) {
+        // failed to decrypt
+        return false;
+      }
+
       // Can successfully derive public key
-      await USER_PUBLIC_KEY.derive(privateKey, {
+      const publicKey = await USER_PUBLIC_KEY.derive(privateKey, {
         cryptoFunctionService: this.cryptoFunctionService,
       });
+
+      if (publicKey == null) {
+        // failed to decrypt
+        return false;
+      }
     } catch (e) {
       return false;
     }
diff --git a/libs/common/src/platform/services/key-state/user-key.state.spec.ts b/libs/common/src/platform/services/key-state/user-key.state.spec.ts
index cb758943e5..5c5c5ac70c 100644
--- a/libs/common/src/platform/services/key-state/user-key.state.spec.ts
+++ b/libs/common/src/platform/services/key-state/user-key.state.spec.ts
@@ -8,7 +8,6 @@ import { EncryptService } from "../../abstractions/encrypt.service";
 import { EncryptionType } from "../../enums";
 import { Utils } from "../../misc/utils";
 import { EncString } from "../../models/domain/enc-string";
-import { CryptoService } from "../crypto.service";
 
 import {
   USER_ENCRYPTED_PRIVATE_KEY,
@@ -89,40 +88,37 @@ describe("Derived decrypted private key", () => {
   });
 
   it("should derive decrypted private key", async () => {
-    const cryptoService = mock<CryptoService>();
-    cryptoService.getUserKey.mockResolvedValue(userKey);
+    const getUserKey = jest.fn(async () => userKey);
     const encryptService = mock<EncryptService>();
     encryptService.decryptToBytes.mockResolvedValue(decryptedPrivateKey);
 
     const result = await sut.derive([userId, encryptedPrivateKey], {
       encryptService,
-      cryptoService,
+      getUserKey,
     });
 
     expect(result).toEqual(decryptedPrivateKey);
   });
 
   it("should handle null input values", async () => {
-    const cryptoService = mock<CryptoService>();
-    cryptoService.getUserKey.mockResolvedValue(userKey);
+    const getUserKey = jest.fn(async () => userKey);
     const encryptService = mock<EncryptService>();
 
     const result = await sut.derive([userId, null], {
       encryptService,
-      cryptoService,
+      getUserKey,
     });
 
     expect(result).toEqual(null);
   });
 
   it("should handle null user key", async () => {
-    const cryptoService = mock<CryptoService>();
-    cryptoService.getUserKey.mockResolvedValue(null);
+    const getUserKey = jest.fn(async () => null);
     const encryptService = mock<EncryptService>();
 
     const result = await sut.derive([userId, encryptedPrivateKey], {
       encryptService,
-      cryptoService,
+      getUserKey,
     });
 
     expect(result).toEqual(null);
diff --git a/libs/common/src/platform/services/key-state/user-key.state.ts b/libs/common/src/platform/services/key-state/user-key.state.ts
index 609525b0ac..3df3b2044b 100644
--- a/libs/common/src/platform/services/key-state/user-key.state.ts
+++ b/libs/common/src/platform/services/key-state/user-key.state.ts
@@ -1,10 +1,10 @@
+import { UserId } from "../../../types/guid";
 import { UserPrivateKey, UserPublicKey, UserKey } from "../../../types/key";
 import { CryptoFunctionService } from "../../abstractions/crypto-function.service";
 import { EncryptService } from "../../abstractions/encrypt.service";
 import { EncString, EncryptedString } from "../../models/domain/enc-string";
 import { SymmetricCryptoKey } from "../../models/domain/symmetric-crypto-key";
 import { CRYPTO_DISK, DeriveDefinition, CRYPTO_MEMORY, UserKeyDefinition } from "../../state";
-import { CryptoService } from "../crypto.service";
 
 export const USER_EVER_HAD_USER_KEY = new UserKeyDefinition<boolean>(
   CRYPTO_DISK,
@@ -28,15 +28,15 @@ export const USER_PRIVATE_KEY = DeriveDefinition.fromWithUserId<
   EncryptedString,
   UserPrivateKey,
   // TODO: update cryptoService to user key directly
-  { encryptService: EncryptService; cryptoService: CryptoService }
+  { encryptService: EncryptService; getUserKey: (userId: UserId) => Promise<UserKey> }
 >(USER_ENCRYPTED_PRIVATE_KEY, {
   deserializer: (obj) => new Uint8Array(Object.values(obj)) as UserPrivateKey,
-  derive: async ([userId, encPrivateKeyString], { encryptService, cryptoService }) => {
+  derive: async ([userId, encPrivateKeyString], { encryptService, getUserKey }) => {
     if (encPrivateKeyString == null) {
       return null;
     }
 
-    const userKey = await cryptoService.getUserKey(userId);
+    const userKey = await getUserKey(userId);
     if (userKey == null) {
       return null;
     }