fix seedless-onboarding-controller tests

mikesposito · mikesposito · commit 62318ccef7b6 · 2025-06-12T13:28:02.000+02:00
diff --git a/packages/keyring-controller/src/KeyringController.ts b/packages/keyring-controller/src/KeyringController.ts
@@ -374,94 +374,96 @@ export type GenericEncryptor = {
  * An encryptor interface that supports encrypting and decrypting
  * serializable data with a password, and exporting and importing keys.
  */
-export type ExportableKeyEncryptor<EncryptionKey = unknown> =
-  GenericEncryptor & {
-    /**
-     * Encrypts the given object with the given encryption key.
-     *
-     * @param key - The encryption key to encrypt with.
-     * @param object - The object to encrypt.
-     * @returns The encryption result.
-     */
-    encryptWithKey: (
-      key: EncryptionKey,
-      object: Json,
-    ) => Promise<encryptorUtils.EncryptionResult>;
-    /**
-     * Encrypts the given object with the given password, and returns the
-     * encryption result and the exported key string.
-     *
-     * @param password - The password to encrypt with.
-     * @param object - The object to encrypt.
-     * @param salt - The optional salt to use for encryption.
-     * @returns The encrypted string and the exported key string.
-     */
-    encryptWithDetail: (
-      password: string,
-      object: Json,
-      salt?: string,
-    ) => Promise<encryptorUtils.DetailedEncryptionResult>;
-    /**
-     * Decrypts the given encrypted string with the given encryption key.
-     *
-     * @param key - The encryption key to decrypt with.
-     * @param encryptedString - The encrypted string to decrypt.
-     * @returns The decrypted object.
-     */
-    decryptWithKey: (
-      key: EncryptionKey,
-      encryptedString: string,
-    ) => Promise<unknown>;
-    /**
-     * Decrypts the given encrypted string with the given password, and returns
-     * the decrypted object and the salt and exported key string used for
-     * encryption.
-     *
-     * @param password - The password to decrypt with.
-     * @param encryptedString - The encrypted string to decrypt.
-     * @returns The decrypted object and the salt and exported key string used for
-     * encryption.
-     */
-    decryptWithDetail: (
-      password: string,
-      encryptedString: string,
-    ) => Promise<encryptorUtils.DetailedDecryptResult>;
-    /**
-     * Generates an encryption key from exported key string.
-     *
-     * @param key - The exported key string.
-     * @returns The encryption key.
-     */
-    importKey: (key: string) => Promise<EncryptionKey>;
-    /**
-     * Exports the encryption key as a string.
-     *
-     * @param key - The encryption key to export.
-     * @returns The exported key string.
-     */
-    exportKey: (key: EncryptionKey) => Promise<string>;
-    /**
-     * Derives an encryption key from a password.
-     *
-     * @param password - The password to derive the key from.
-     * @param salt - The salt to use for key derivation.
-     * @param exportable - Whether the key should be exportable or not.
-     * @param options - Optional key derivation options.
-     * @returns The derived encryption key.
-     */
-    keyFromPassword: (
-      password: string,
-      salt: string,
-      exportable?: boolean,
-      // setting this to unknown as currently each client has different
-      // key derivation options
-      keyDerivationOptions?: unknown,
-    ) => Promise<EncryptionKey>;
-    /**
-     * Generates a random salt for key derivation.
-     */
-    generateSalt: typeof encryptorUtils.generateSalt;
-  };
+export type ExportableKeyEncryptor<
+  EncryptionKey = unknown,
+  SupportedKeyDerivationParams = unknown,
+> = GenericEncryptor & {
+  /**
+   * Encrypts the given object with the given encryption key.
+   *
+   * @param key - The encryption key to encrypt with.
+   * @param object - The object to encrypt.
+   * @returns The encryption result.
+   */
+  encryptWithKey: (
+    key: EncryptionKey,
+    object: Json,
+  ) => Promise<encryptorUtils.EncryptionResult>;
+  /**
+   * Encrypts the given object with the given password, and returns the
+   * encryption result and the exported key string.
+   *
+   * @param password - The password to encrypt with.
+   * @param object - The object to encrypt.
+   * @param salt - The optional salt to use for encryption.
+   * @returns The encrypted string and the exported key string.
+   */
+  encryptWithDetail: (
+    password: string,
+    object: Json,
+    salt?: string,
+  ) => Promise<encryptorUtils.DetailedEncryptionResult>;
+  /**
+   * Decrypts the given encrypted string with the given encryption key.
+   *
+   * @param key - The encryption key to decrypt with.
+   * @param encryptedString - The encrypted string to decrypt.
+   * @returns The decrypted object.
+   */
+  decryptWithKey: (
+    key: EncryptionKey,
+    encryptedString: string,
+  ) => Promise<unknown>;
+  /**
+   * Decrypts the given encrypted string with the given password, and returns
+   * the decrypted object and the salt and exported key string used for
+   * encryption.
+   *
+   * @param password - The password to decrypt with.
+   * @param encryptedString - The encrypted string to decrypt.
+   * @returns The decrypted object and the salt and exported key string used for
+   * encryption.
+   */
+  decryptWithDetail: (
+    password: string,
+    encryptedString: string,
+  ) => Promise<encryptorUtils.DetailedDecryptResult>;
+  /**
+   * Generates an encryption key from exported key string.
+   *
+   * @param key - The exported key string.
+   * @returns The encryption key.
+   */
+  importKey: (key: string) => Promise<EncryptionKey>;
+  /**
+   * Exports the encryption key as a string.
+   *
+   * @param key - The encryption key to export.
+   * @returns The exported key string.
+   */
+  exportKey: (key: EncryptionKey) => Promise<string>;
+  /**
+   * Derives an encryption key from a password.
+   *
+   * @param password - The password to derive the key from.
+   * @param salt - The salt to use for key derivation.
+   * @param exportable - Whether the key should be exportable or not.
+   * @param options - Optional key derivation options.
+   * @returns The derived encryption key.
+   */
+  keyFromPassword: (
+    password: string,
+    salt: string,
+    exportable?: boolean,
+    // setting this to unknown as currently each client has different
+    // key derivation options
+    keyDerivationOptions?: SupportedKeyDerivationParams,
+  ) => Promise<EncryptionKey>;
+  /**
+   * Generates a random salt for key derivation.
+   */
+  generateSalt: typeof encryptorUtils.generateSalt;
+};
 
 export type KeyringSelector =
   | {
diff --git a/packages/seedless-onboarding-controller/src/SeedlessOnboardingController.test.ts b/packages/seedless-onboarding-controller/src/SeedlessOnboardingController.test.ts
@@ -1,13 +1,19 @@
 import { keccak256AndHexify } from '@metamask/auth-network-utils';
 import type { Messenger } from '@metamask/base-controller';
-import type { EncryptionKey } from '@metamask/browser-passworder';
+import type {
+  EncryptionKey,
+  KeyDerivationOptions,
+} from '@metamask/browser-passworder';
 import {
   encrypt,
   decrypt,
   decryptWithDetail,
   encryptWithDetail,
   decryptWithKey as decryptWithKeyBrowserPassworder,
   importKey as importKeyBrowserPassworder,
+  exportKey as exportKeyBrowserPassworder,
+  generateSalt as generateSaltBrowserPassworder,
+  keyFromPassword as keyFromPasswordBrowserPassworder,
 } from '@metamask/browser-passworder';
 import {
   TOPRFError,
@@ -99,27 +105,34 @@ const MOCK_AUTH_PUB_KEY = 'A09CwPHdl/qo2AjBOHen5d4QORaLedxOrSdgReq8IhzQ';
 const MOCK_AUTH_PUB_KEY_OUTDATED =
   'Ao2sa8imX7SD4KE4fJLoJ/iBufmaBxSFygG1qUhW2qAb';
 
-type WithControllerCallback<ReturnValue, EKey> = ({
-  controller,
-  initialState,
-  encryptor,
-  messenger,
-}: {
-  controller: SeedlessOnboardingController<EKey>;
-  encryptor: VaultEncryptor<EKey>;
-  initialState: SeedlessOnboardingControllerState;
-  messenger: SeedlessOnboardingControllerMessenger;
-  baseMessenger: Messenger<AllowedActions, AllowedEvents>;
-  toprfClient: ToprfSecureBackup;
-}) => Promise<ReturnValue> | ReturnValue;
-
-type WithControllerOptions<EKey> = Partial<
-  SeedlessOnboardingControllerOptions<EKey>
+type WithControllerCallback<ReturnValue, EKey, SupportedKeyDerivationOptions> =
+  ({
+    controller,
+    initialState,
+    encryptor,
+    messenger,
+  }: {
+    controller: SeedlessOnboardingController<
+      EKey,
+      SupportedKeyDerivationOptions
+    >;
+    encryptor: VaultEncryptor<EKey, KeyDerivationOptions>;
+    initialState: SeedlessOnboardingControllerState;
+    messenger: SeedlessOnboardingControllerMessenger;
+    baseMessenger: Messenger<AllowedActions, AllowedEvents>;
+    toprfClient: ToprfSecureBackup;
+  }) => Promise<ReturnValue> | ReturnValue;
+
+type WithControllerOptions<EKey, SupportedKeyDerivationOptions> = Partial<
+  SeedlessOnboardingControllerOptions<EKey, SupportedKeyDerivationOptions>
 >;
 
-type WithControllerArgs<ReturnValue, EKey> =
-  | [WithControllerCallback<ReturnValue, EKey>]
-  | [WithControllerOptions<EKey>, WithControllerCallback<ReturnValue, EKey>];
+type WithControllerArgs<ReturnValue, EKey, SupportedKeyDerivationOptions> =
+  | [WithControllerCallback<ReturnValue, EKey, SupportedKeyDerivationOptions>]
+  | [
+      WithControllerOptions<EKey, SupportedKeyDerivationOptions>,
+      WithControllerCallback<ReturnValue, EKey, SupportedKeyDerivationOptions>,
+    ];
 
 /**
  * Get the default vault encryptor for the Seedless Onboarding Controller.
@@ -139,6 +152,9 @@ function getDefaultSeedlessOnboardingVaultEncryptor() {
       payload: unknown,
     ) => Promise<unknown>,
     importKey: importKeyBrowserPassworder,
+    exportKey: exportKeyBrowserPassworder,
+    generateSalt: generateSaltBrowserPassworder,
+    keyFromPassword: keyFromPasswordBrowserPassworder,
   };
 }
 
@@ -162,13 +178,20 @@ function createMockVaultEncryptor() {
  * @returns Whatever the callback returns.
  */
 async function withController<ReturnValue>(
-  ...args: WithControllerArgs<ReturnValue, EncryptionKey | webcrypto.CryptoKey>
+  ...args: WithControllerArgs<
+    ReturnValue,
+    EncryptionKey | webcrypto.CryptoKey,
+    KeyDerivationOptions
+  >
 ) {
   const [{ ...rest }, fn] = args.length === 2 ? args : [{}, args[0]];
   const encryptor = new MockVaultEncryptor();
   const { messenger, baseMessenger } = mockSeedlessOnboardingMessenger();
 
-  const controller = new SeedlessOnboardingController({
+  const controller = new SeedlessOnboardingController<
+    EncryptionKey | webcrypto.CryptoKey,
+    KeyDerivationOptions
+  >({
     encryptor,
     messenger,
     network: Web3AuthNetwork.Devnet,
@@ -313,9 +336,12 @@ function mockChangeEncKey(
  * @param seedPhrase - The mock seed phrase.
  * @param keyringId - The mock keyring id.
  */
-async function mockCreateToprfKeyAndBackupSeedPhrase<EKey>(
+async function mockCreateToprfKeyAndBackupSeedPhrase<
+  EKey,
+  SupportedKeyDerivationOptions,
+>(
   toprfClient: ToprfSecureBackup,
-  controller: SeedlessOnboardingController<EKey>,
+  controller: SeedlessOnboardingController<EKey, SupportedKeyDerivationOptions>,
   password: string,
   seedPhrase: Uint8Array,
   keyringId: string,
@@ -462,7 +488,10 @@ describe('SeedlessOnboardingController', () => {
   describe('constructor', () => {
     it('should be able to instantiate', () => {
       const { messenger } = mockSeedlessOnboardingMessenger();
-      const controller = new SeedlessOnboardingController({
+      const controller = new SeedlessOnboardingController<
+        EncryptionKey | webcrypto.CryptoKey,
+        KeyDerivationOptions
+      >({
         messenger,
         encryptor: getDefaultSeedlessOnboardingVaultEncryptor(),
       });
diff --git a/packages/seedless-onboarding-controller/src/SeedlessOnboardingController.ts b/packages/seedless-onboarding-controller/src/SeedlessOnboardingController.ts
@@ -111,12 +111,18 @@ const seedlessOnboardingMetadata: StateMetadata<SeedlessOnboardingControllerStat
     },
   };
 
-export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
+export class SeedlessOnboardingController<
+  EncryptionKey,
+  SupportedKeyDerivationOptions,
+> extends BaseController<
   typeof controllerName,
   SeedlessOnboardingControllerState,
   SeedlessOnboardingControllerMessenger
 > {
-  readonly #vaultEncryptor: VaultEncryptor<EncryptionKey>;
+  readonly #vaultEncryptor: VaultEncryptor<
+    EncryptionKey,
+    SupportedKeyDerivationOptions
+  >;
 
   readonly #controllerOperationMutex = new Mutex();
 
@@ -147,7 +153,10 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
     encryptor,
     toprfKeyDeriver,
     network = Web3AuthNetwork.Mainnet,
-  }: SeedlessOnboardingControllerOptions<EncryptionKey>) {
+  }: SeedlessOnboardingControllerOptions<
+    EncryptionKey,
+    SupportedKeyDerivationOptions
+  >) {
     super({
       name: controllerName,
       metadata: seedlessOnboardingMetadata,
diff --git a/packages/seedless-onboarding-controller/src/types.ts b/packages/seedless-onboarding-controller/src/types.ts
@@ -157,8 +157,8 @@ export type SeedlessOnboardingControllerMessenger = RestrictedMessenger<
 /**
  * Encryptor interface for encrypting and decrypting seedless onboarding vault.
  */
-export type VaultEncryptor<EncryptionKey> = Omit<
-  ExportableKeyEncryptor<EncryptionKey>,
+export type VaultEncryptor<EncryptionKey, KeyDerivationParams> = Omit<
+  ExportableKeyEncryptor<EncryptionKey, KeyDerivationParams>,
   'encryptWithKey'
 >;
 
@@ -189,7 +189,10 @@ export type ToprfKeyDeriver = {
  * @param state - The initial state to set on this controller.
  * @param encryptor - The encryptor to use for encrypting and decrypting seedless onboarding vault.
  */
-export type SeedlessOnboardingControllerOptions<EncryptionKey> = {
+export type SeedlessOnboardingControllerOptions<
+  EncryptionKey,
+  SupportedKeyDerivationOptions,
+> = {
   messenger: SeedlessOnboardingControllerMessenger;
 
   /**
@@ -202,7 +205,7 @@ export type SeedlessOnboardingControllerOptions<EncryptionKey> = {
    *
    * @default browser-passworder @link https://github.com/MetaMask/browser-passworder
    */
-  encryptor: VaultEncryptor<EncryptionKey>;
+  encryptor: VaultEncryptor<EncryptionKey, SupportedKeyDerivationOptions>;
 
   /**
    * Optional key derivation interface for the TOPRF client.
diff --git a/packages/seedless-onboarding-controller/tests/mocks/vaultEncryptor.ts b/packages/seedless-onboarding-controller/tests/mocks/vaultEncryptor.ts
