From 9bb4043e4842fc1d8f903bde28ae7dbcdcc859e5 Mon Sep 17 00:00:00 2001
From: Mark Stacey <markjstacey@gmail.com>
Date: Tue, 4 Jun 2024 09:46:04 -0230
Subject: [PATCH] feat: Allow overwriting built-in keyring builders (#4362)

## Explanation

The KeyringController comes with two built-in keyrings: Simple and HD.
Unfortunately it's impossible to overwrite these because when we build a
new keyring, we look for the first keyring builder in the list that
matches the type we want to build, and the built-in keyrings are always
first.

The order has been switched so that built-in keyrings come second, after
custom keyring builders. This allows them to be overwritten.

This makes it possible to test behavior that is specific to simple or HD
keyrings. This is something that I wanted to do in the mobile
repository.

## References

N/A

## Changelog

### `@metamask/keyring-controller`

#### Added

- Add support for overwriting built-in keyring builders for the Simple
and HD keyring.

## Checklist

- [x] I've updated the test suite for new or updated code as appropriate
- [x] I've updated documentation (JSDoc, Markdown, etc.) for new or
updated code as appropriate
- [x] I've highlighted breaking changes using the "BREAKING" category
above as appropriate
---
 .../src/KeyringController.test.ts             | 45 +++++++++++++++++++
 .../src/KeyringController.ts                  |  2 +-
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/packages/keyring-controller/src/KeyringController.test.ts b/packages/keyring-controller/src/KeyringController.test.ts
index 018c136539..48b1130f61 100644
--- a/packages/keyring-controller/src/KeyringController.test.ts
+++ b/packages/keyring-controller/src/KeyringController.test.ts
@@ -11,6 +11,7 @@ import {
   SignTypedDataVersion,
   encrypt,
 } from '@metamask/eth-sig-util';
+import SimpleKeyring from '@metamask/eth-simple-keyring/dist/simple-keyring';
 import type { EthKeyring } from '@metamask/keyring-api';
 import { wordlist } from '@metamask/scure-bip39/dist/wordlists/english';
 import type { KeyringClass } from '@metamask/utils';
@@ -100,6 +101,32 @@ describe('KeyringController', () => {
           }),
       ).toThrow(KeyringControllerError.UnsupportedEncryptionKeyExport);
     });
+
+    it('allows overwriting the built-in Simple keyring builder', async () => {
+      const mockSimpleKeyringBuilder =
+        // @ts-expect-error The simple keyring doesn't yet conform to the KeyringClass type
+        buildKeyringBuilderWithSpy(SimpleKeyring);
+      await withController(
+        { keyringBuilders: [mockSimpleKeyringBuilder] },
+        async ({ controller }) => {
+          await controller.addNewKeyring(KeyringTypes.simple);
+
+          expect(mockSimpleKeyringBuilder).toHaveBeenCalledTimes(1);
+        },
+      );
+    });
+
+    it('allows overwriting the built-in HD keyring builder', async () => {
+      const mockHdKeyringBuilder = buildKeyringBuilderWithSpy(HDKeyring);
+      await withController(
+        { keyringBuilders: [mockHdKeyringBuilder] },
+        async () => {
+          // This is called as part of initializing the controller
+          // because the first keyring is assumed to always be an HD keyring
+          expect(mockHdKeyringBuilder).toHaveBeenCalledTimes(1);
+        },
+      );
+    });
   });
 
   describe('addNewAccount', () => {
@@ -3538,3 +3565,21 @@ async function withController<ReturnValue>(
     messenger,
   });
 }
+
+/**
+ * Construct a keyring builder with a spy.
+ *
+ * @param KeyringConstructor - The constructor to use for building the keyring.
+ * @returns A keyring builder that uses `jest.fn()` to spy on invocations.
+ */
+function buildKeyringBuilderWithSpy(KeyringConstructor: KeyringClass<Json>): {
+  (): EthKeyring<Json>;
+  type: string;
+} {
+  const keyringBuilderWithSpy: { (): EthKeyring<Json>; type?: string } = jest
+    .fn()
+    .mockImplementation((...args) => new KeyringConstructor(...args));
+  keyringBuilderWithSpy.type = KeyringConstructor.type;
+  // Not sure why TypeScript isn't smart enough to infer that `type` is set here.
+  return keyringBuilderWithSpy as { (): EthKeyring<Json>; type: string };
+}
diff --git a/packages/keyring-controller/src/KeyringController.ts b/packages/keyring-controller/src/KeyringController.ts
index 1d657295ad..4979f4f7e0 100644
--- a/packages/keyring-controller/src/KeyringController.ts
+++ b/packages/keyring-controller/src/KeyringController.ts
@@ -611,7 +611,7 @@ export class KeyringController extends BaseController<
     });
 
     this.#keyringBuilders = keyringBuilders
-      ? defaultKeyringBuilders.concat(keyringBuilders)
+      ? keyringBuilders.concat(defaultKeyringBuilders)
       : defaultKeyringBuilders;
 
     this.#encryptor = encryptor;