Skip to content

Commit ffd16d4

Browse files
authored
[keyring-controller] Add withKeyring method (#4197)
## Explanation Part of `KeyringController` responsibilies is ensuring each operation is [mutually exclusive](#4182) and [atomic](#4192), updating keyring instances and the vault (or rolling them back) in the same mutex lock. However, the ability of clients to have direct access to a keyring instance represents a loophole, as with the current implementation they don’t have to comply with the rules enforced by the controller: we should provide a way for clients to interact with a keyring instance through safeguards provided by KeyringController. The current behavior is this one: 1. Client obtains a keyring instance through `getKeyringForAccount` 2. Client interacts with the instance 3. Client calls `persistAllKeyrings` We should, instead, have something like this: 1. Client calls a `withKeyring` method, passing a _keyring selector_ and a callback 2. KeyringController selects the keyring instance and calls the callback with it, after locking the controller operation mutex 3. Client interacts with the keyring instance safely, inside the callback 4. KeyringController, after the callback execution, internally updates the vault or rolls back changes in case of error, and then releases the mutex lock ## References <!-- Are there any issues that this pull request is tied to? Are there other links that reviewers should consult to understand these changes better? For example: * Fixes #12345 * Related to #67890 --> * Fixes #4198 * Related to #4192 ## Changelog <!-- If you're making any consumer-facing changes, list those changes here as if you were updating a changelog, using the template below as a guide. (CATEGORY is one of BREAKING, ADDED, CHANGED, DEPRECATED, REMOVED, or FIXED. For security-related issues, follow the Security Advisory process.) Please take care to name the exact pieces of the API you've added or changed (e.g. types, interfaces, functions, or methods). If there are any breaking changes, make sure to offer a solution for consumers to follow once they upgrade to the changes. Finally, if you're only making changes to development scripts or tests, you may replace the template below with "None". --> ### `@metamask/keyring-controller` - **ADDED**: Added `withKeyring` method - **DEPRECATED**: Deprecated `persistAllKeyrings` method - Use `withKeyring` instead ## Checklist - [ ] I've updated the test suite for new or updated code as appropriate - [ ] I've updated documentation (JSDoc, Markdown, etc.) for new or updated code as appropriate - [ ] I've highlighted breaking changes using the "BREAKING" category above as appropriate
1 parent 4264100 commit ffd16d4

File tree

4 files changed

+255
-5
lines changed

4 files changed

+255
-5
lines changed

packages/keyring-controller/jest.config.js

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,10 @@ module.exports = merge(baseConfig, {
1717
// An object that configures minimum threshold enforcement for coverage results
1818
coverageThreshold: {
1919
global: {
20-
branches: 94.44,
20+
branches: 94.8,
2121
functions: 100,
22-
lines: 98.84,
23-
statements: 98.85,
22+
lines: 98.87,
23+
statements: 98.88,
2424
},
2525
},
2626

packages/keyring-controller/src/KeyringController.test.ts

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2220,6 +2220,140 @@ describe('KeyringController', () => {
22202220
});
22212221
});
22222222

2223+
describe('withKeyring', () => {
2224+
it('should rollback if an error is thrown', async () => {
2225+
await withController(async ({ controller, initialState }) => {
2226+
const selector = { type: KeyringTypes.hd };
2227+
const fn = async (keyring: EthKeyring<Json>) => {
2228+
await keyring.addAccounts(1);
2229+
throw new Error('Oops');
2230+
};
2231+
2232+
await expect(controller.withKeyring(selector, fn)).rejects.toThrow(
2233+
'Oops',
2234+
);
2235+
expect(controller.state.keyrings[0].accounts).toHaveLength(1);
2236+
expect(await controller.getAccounts()).toStrictEqual(
2237+
initialState.keyrings[0].accounts,
2238+
);
2239+
});
2240+
});
2241+
2242+
describe('when the keyring is selected by type', () => {
2243+
it('should call the given function with the selected keyring', async () => {
2244+
await withController(async ({ controller }) => {
2245+
const fn = jest.fn();
2246+
const selector = { type: KeyringTypes.hd };
2247+
const keyring = controller.getKeyringsByType(KeyringTypes.hd)[0];
2248+
2249+
await controller.withKeyring(selector, fn);
2250+
2251+
expect(fn).toHaveBeenCalledWith(keyring);
2252+
});
2253+
});
2254+
2255+
it('should return the result of the function', async () => {
2256+
await withController(async ({ controller }) => {
2257+
const fn = async () => Promise.resolve('hello');
2258+
const selector = { type: KeyringTypes.hd };
2259+
2260+
expect(await controller.withKeyring(selector, fn)).toBe('hello');
2261+
});
2262+
});
2263+
2264+
it('should throw an error if the callback returns the selected keyring', async () => {
2265+
await withController(async ({ controller }) => {
2266+
await expect(
2267+
controller.withKeyring(
2268+
{ type: KeyringTypes.hd },
2269+
async (keyring) => {
2270+
return keyring;
2271+
},
2272+
),
2273+
).rejects.toThrow(KeyringControllerError.UnsafeDirectKeyringAccess);
2274+
});
2275+
});
2276+
2277+
describe('when the keyring is not found', () => {
2278+
it('should throw an error if the keyring is not found and `createIfMissing` is false', async () => {
2279+
await withController(async ({ controller }) => {
2280+
const selector = { type: 'foo' };
2281+
const fn = jest.fn();
2282+
2283+
await expect(controller.withKeyring(selector, fn)).rejects.toThrow(
2284+
KeyringControllerError.KeyringNotFound,
2285+
);
2286+
expect(fn).not.toHaveBeenCalled();
2287+
});
2288+
});
2289+
2290+
it('should add the keyring if `createIfMissing` is true', async () => {
2291+
await withController(
2292+
{ keyringBuilders: [keyringBuilderFactory(MockKeyring)] },
2293+
async ({ controller }) => {
2294+
const selector = { type: MockKeyring.type };
2295+
const fn = jest.fn();
2296+
2297+
await controller.withKeyring(selector, fn, {
2298+
createIfMissing: true,
2299+
});
2300+
2301+
expect(fn).toHaveBeenCalled();
2302+
expect(controller.state.keyrings).toHaveLength(2);
2303+
},
2304+
);
2305+
});
2306+
});
2307+
});
2308+
2309+
describe('when the keyring is selected by address', () => {
2310+
it('should call the given function with the selected keyring', async () => {
2311+
await withController(async ({ controller, initialState }) => {
2312+
const fn = jest.fn();
2313+
const selector = {
2314+
address: initialState.keyrings[0].accounts[0] as Hex,
2315+
};
2316+
const keyring = controller.getKeyringsByType(KeyringTypes.hd)[0];
2317+
2318+
await controller.withKeyring(selector, fn);
2319+
2320+
expect(fn).toHaveBeenCalledWith(keyring);
2321+
});
2322+
});
2323+
2324+
it('should return the result of the function', async () => {
2325+
await withController(async ({ controller, initialState }) => {
2326+
const fn = async () => Promise.resolve('hello');
2327+
const selector = {
2328+
address: initialState.keyrings[0].accounts[0] as Hex,
2329+
};
2330+
2331+
expect(await controller.withKeyring(selector, fn)).toBe('hello');
2332+
});
2333+
});
2334+
2335+
describe('when the keyring is not found', () => {
2336+
[true, false].forEach((value) =>
2337+
it(`should throw an error if the createIfMissing is ${value}`, async () => {
2338+
await withController(async ({ controller }) => {
2339+
const selector = {
2340+
address: '0x4584d2B4905087A100420AFfCe1b2d73fC69B8E4' as Hex,
2341+
};
2342+
const fn = jest.fn();
2343+
2344+
await expect(
2345+
controller.withKeyring(selector, fn),
2346+
).rejects.toThrow(
2347+
'KeyringController - No keyring found. Error info: There are keyrings, but none match the address',
2348+
);
2349+
expect(fn).not.toHaveBeenCalled();
2350+
});
2351+
}),
2352+
);
2353+
});
2354+
});
2355+
});
2356+
22232357
describe('QR keyring', () => {
22242358
const composeMockSignature = (
22252359
requestId: string,

packages/keyring-controller/src/KeyringController.ts

Lines changed: 116 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -364,6 +364,15 @@ export type ExportableKeyEncryptor = GenericEncryptor & {
364364
importKey: (key: string) => Promise<unknown>;
365365
};
366366

367+
export type KeyringSelector =
368+
| {
369+
type: string;
370+
index?: number;
371+
}
372+
| {
373+
address: Hex;
374+
};
375+
367376
/**
368377
* A function executed within a mutually exclusive lock, with
369378
* a mutex releaser in its option bag.
@@ -845,7 +854,7 @@ export class KeyringController extends BaseController<
845854
*
846855
* @deprecated Use of this method is discouraged as actions executed directly on
847856
* keyrings are not being reflected in the KeyringController state and not
848-
* persisted in the vault.
857+
* persisted in the vault. Use `withKeyring` instead.
849858
* @param account - An account address.
850859
* @returns Promise resolving to keyring of the `account` if one exists.
851860
*/
@@ -888,7 +897,7 @@ export class KeyringController extends BaseController<
888897
*
889898
* @deprecated Use of this method is discouraged as actions executed directly on
890899
* keyrings are not being reflected in the KeyringController state and not
891-
* persisted in the vault.
900+
* persisted in the vault. Use `withKeyring` instead.
892901
* @param type - Keyring type name.
893902
* @returns An array of keyrings of the given type.
894903
*/
@@ -899,6 +908,7 @@ export class KeyringController extends BaseController<
899908
/**
900909
* Persist all serialized keyrings in the vault.
901910
*
911+
* @deprecated This method is being phased out in favor of `withKeyring`.
902912
* @returns Promise resolving with `true` value when the
903913
* operation completes.
904914
*/
@@ -1299,6 +1309,110 @@ export class KeyringController extends BaseController<
12991309
return seedWords;
13001310
}
13011311

1312+
/**
1313+
* Select a keyring and execute the given operation with
1314+
* the selected keyring, as a mutually exclusive atomic
1315+
* operation.
1316+
*
1317+
* The method automatically persists changes at the end of the
1318+
* function execution, or rolls back the changes if an error
1319+
* is thrown.
1320+
*
1321+
* @param selector - Keyring selector object.
1322+
* @param operation - Function to execute with the selected keyring.
1323+
* @param options - Additional options.
1324+
* @param options.createIfMissing - Whether to create a new keyring if the selected one is missing.
1325+
* @param options.createWithData - Optional data to use when creating a new keyring.
1326+
* @returns Promise resolving to the result of the function execution.
1327+
* @template SelectedKeyring - The type of the selected keyring.
1328+
* @template CallbackResult - The type of the value resolved by the callback function.
1329+
* @deprecated This method overload is deprecated. Use `withKeyring` without options instead.
1330+
*/
1331+
async withKeyring<
1332+
SelectedKeyring extends EthKeyring<Json> = EthKeyring<Json>,
1333+
CallbackResult = void,
1334+
>(
1335+
selector: KeyringSelector,
1336+
operation: (keyring: SelectedKeyring) => Promise<CallbackResult>,
1337+
// eslint-disable-next-line @typescript-eslint/unified-signatures
1338+
options:
1339+
| { createIfMissing?: false }
1340+
| { createIfMissing: true; createWithData?: unknown },
1341+
): Promise<CallbackResult>;
1342+
1343+
/**
1344+
* Select a keyring and execute the given operation with
1345+
* the selected keyring, as a mutually exclusive atomic
1346+
* operation.
1347+
*
1348+
* The method automatically persists changes at the end of the
1349+
* function execution, or rolls back the changes if an error
1350+
* is thrown.
1351+
*
1352+
* @param selector - Keyring selector object.
1353+
* @param operation - Function to execute with the selected keyring.
1354+
* @returns Promise resolving to the result of the function execution.
1355+
* @template SelectedKeyring - The type of the selected keyring.
1356+
* @template CallbackResult - The type of the value resolved by the callback function.
1357+
*/
1358+
async withKeyring<
1359+
SelectedKeyring extends EthKeyring<Json> = EthKeyring<Json>,
1360+
CallbackResult = void,
1361+
>(
1362+
selector: KeyringSelector,
1363+
operation: (keyring: SelectedKeyring) => Promise<CallbackResult>,
1364+
): Promise<CallbackResult>;
1365+
1366+
async withKeyring<
1367+
SelectedKeyring extends EthKeyring<Json> = EthKeyring<Json>,
1368+
CallbackResult = void,
1369+
>(
1370+
selector: KeyringSelector,
1371+
operation: (keyring: SelectedKeyring) => Promise<CallbackResult>,
1372+
options:
1373+
| { createIfMissing?: false }
1374+
| { createIfMissing: true; createWithData?: unknown } = {
1375+
createIfMissing: false,
1376+
},
1377+
): Promise<CallbackResult> {
1378+
return this.#persistOrRollback(async () => {
1379+
let keyring: SelectedKeyring | undefined;
1380+
1381+
if ('address' in selector) {
1382+
keyring = (await this.getKeyringForAccount(selector.address)) as
1383+
| SelectedKeyring
1384+
| undefined;
1385+
} else {
1386+
keyring = this.getKeyringsByType(selector.type)[selector.index || 0] as
1387+
| SelectedKeyring
1388+
| undefined;
1389+
1390+
if (!keyring && options.createIfMissing) {
1391+
keyring = (await this.#newKeyring(
1392+
selector.type,
1393+
options.createWithData,
1394+
)) as SelectedKeyring;
1395+
}
1396+
}
1397+
1398+
if (!keyring) {
1399+
throw new Error(KeyringControllerError.KeyringNotFound);
1400+
}
1401+
1402+
const result = await operation(keyring);
1403+
1404+
if (Object.is(result, keyring)) {
1405+
// Access to a keyring instance outside of controller safeguards
1406+
// should be discouraged, as it can lead to unexpected behavior.
1407+
// This error is thrown to prevent consumers using `withKeyring`
1408+
// as a way to get a reference to a keyring instance.
1409+
throw new Error(KeyringControllerError.UnsafeDirectKeyringAccess);
1410+
}
1411+
1412+
return result;
1413+
});
1414+
}
1415+
13021416
// QR Hardware related methods
13031417

13041418
/**

packages/keyring-controller/src/constants.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
export enum KeyringControllerError {
22
NoKeyring = 'KeyringController - No keyring found',
3+
KeyringNotFound = 'KeyringController - Keyring not found.',
4+
UnsafeDirectKeyringAccess = 'KeyringController - Returning keyring instances is unsafe',
35
WrongPasswordType = 'KeyringController - Password must be of type string.',
46
NoFirstAccount = 'KeyringController - First Account not found.',
57
DuplicatedAccount = 'KeyringController - The account you are trying to import is a duplicate',

0 commit comments

Comments
 (0)