From 7d3db4e2de2977c9edc7d3160d9d07c1a10e1dcb Mon Sep 17 00:00:00 2001
From: Mark Stacey <markjstacey@gmail.com>
Date: Mon, 3 Jun 2024 14:18:22 -0230
Subject: [PATCH] feat: Allow overwriting built-in keyring builders

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.
---
 .../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;