Skip to content

Commit

Permalink
feat(keyring-snap-bridge): make account creation async (#207)
Browse files Browse the repository at this point in the history 
We now resume the Snap execution before persisting the account into the
Snap keyring state.

This allows the Snap to have an existing account (once approved or not
by the user, just like before), but this delays the actual account
creation on the MetaMask clients (more specifically, the
`AccountsController` in this case), which allows other controllers to
react to the `AccountsController:accountCreated` event and to directly
fetch account informations (assets/balances) from the Snap.

In the following example, the `SomeOtherController` cannot use the
account because the account wasn't persisted by the Snap yet:
```ts
// Snap keyring methods
snap.createAccount = () => {
  const account = ...;
  
  await emitSnapKeyringEvent(AccountCreated, { account }); // -> This will trigger keyring.handleSnapkeyringEvent
  
  // Now persists the account (since we know MetaMask has accepted it).
  state.accounts.push(account);
}

// Snap keyring (on the MetaMask client)
keyring.handleSnapKeyringEvent = (message) => {
  ...
  switch (message.method) {
    case 'AccountCreated': {
      // Show some dialogs for the user
      if (dialog.accepted) {
        // At this point, Snap execution is still pending...
        await keyring.saveState(); // This will trigger AccountsController:accountAdded (see addNewAccount method in this example) event 
      }
    }
  ...
})

AccountsController.addNewAccount = (account) => {
  this.messenger.publish(`AccountsController:accountAdded`, account); // This will trigger SomeOtherController.onAccountAdded handler
}

SomeOtherController.onAccountAdded = (account) => {
  // !
  // ! Here was the problem, since the account was still not persisted onto the Snap.
  // !
  await this.getSnapClient(account.metadata.snap.id).listAccountAssets();
}
```

---------

Co-authored-by: Daniel Rocha <[email protected]>
  • Loading branch information
@ccharly @danroc
ccharly and danroc authored Feb 14, 2025
1 parent 208be1d commit ecdc034
Show file tree
Hide file tree
Showing 2 changed files with 165 additions and 11 deletions.
135 changes: 126 additions & 9 deletions packages/keyring-snap-bridge/src/SnapKeyring.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,17 @@ import {
KeyringEvent,
BtcScope,
SolScope,
KeyringRpcMethod,
} from '@metamask/keyring-api';
import type { JsonRpcRequest } from '@metamask/keyring-utils';
import type { HandleSnapRequest } from '@metamask/snaps-controllers';
import type { SnapId } from '@metamask/snaps-sdk';
import { KnownCaipNamespace, toCaipChainId } from '@metamask/utils';

import type { KeyringState } from '.';
import { SnapKeyring } from '.';
import type { KeyringAccountV1 } from './account';
import { DeferredPromise } from './DeferredPromise';
import { migrateAccountV1, getScopesForAccountV1 } from './migrations';
import type {
SnapKeyringAllowedActions,
Expand Down Expand Up @@ -200,9 +203,38 @@ describe('SnapKeyring', () => {
allowedActions: ['SnapController:get', 'SnapController:handleRequest'],
});

// Allow to map a mocked value for a given keyring RPC method
const mockMessengerHandleRequest = (
// We're using `string` here instead of `KeyringRpcMethod` to avoid having to map
// every RPC methods
handlers: Record<string, () => unknown>,
): void => {
mockMessenger.handleRequest.mockImplementation(
(request: Parameters<HandleSnapRequest['handler']>[0]) => {
// First layer of transport is a Snap RPC request for 'OnKeyringRequest'.
expect(request.handler).toBe('onKeyringRequest');

// Second one is for the actual keyring request.
const keyringRequest = request.request as JsonRpcRequest;
const requestHandler = handlers[keyringRequest.method];

if (!requestHandler) {
throw new Error(
`Missing handleRequest handler for: ${keyringRequest.method}`,
);
}
return requestHandler();
},
);
};

beforeEach(async () => {
keyring = new SnapKeyring(mockSnapKeyringMessenger, mockCallbacks);

// We do need to return a promise for this method now:
mockCallbacks.saveState.mockImplementation(async () => {
return null;
});
mockCallbacks.addAccount.mockImplementation(
async (
_address,
Expand All @@ -218,7 +250,9 @@ describe('SnapKeyring', () => {
mockMessenger.get.mockReset();
mockMessenger.handleRequest.mockReset();
for (const account of accounts) {
mockMessenger.handleRequest.mockResolvedValue(accounts);
mockMessengerHandleRequest({
[KeyringRpcMethod.ListAccounts]: () => accounts,
});
await keyring.handleKeyringSnapMessage(snapId, {
method: KeyringEvent.AccountCreated,
params: { account: account as unknown as KeyringAccount },
Expand All @@ -227,28 +261,34 @@ describe('SnapKeyring', () => {
});

describe('handleKeyringSnapMessage', () => {
const newEthEoaAccount = {
id: 'bd63063d-ed58-4b9b-b3da-4282ac2208a8',
options: {},
methods: ETH_EOA_METHODS,
scopes: [EthScope.Eoa],
type: EthAccountType.Eoa,
address: '0x6431726eee67570bf6f0cf892ae0a3988f03903f',
};

describe('#handleAccountCreated', () => {
it('creates the account with a lower-cased address for EVM', async () => {
const evmAccount = {
id: 'b05d918a-b37c-497a-bb28-3d15c0d56b7a',
options: {},
methods: ETH_EOA_METHODS,
scopes: [EthScope.Eoa],
type: EthAccountType.Eoa,
const account = {
...newEthEoaAccount,
// Even checksummed address will be lower-cased by the bridge.
address: '0x6431726EEE67570BF6f0Cf892aE0a3988F03903F',
};

await keyring.handleKeyringSnapMessage(snapId, {
method: KeyringEvent.AccountCreated,
params: {
account: {
...(evmAccount as unknown as KeyringAccount),
...(account as unknown as KeyringAccount),
id: '56189183-9b89-4ae6-90d9-99d167b28520',
},
},
});
expect(mockCallbacks.addAccount).toHaveBeenLastCalledWith(
evmAccount.address.toLowerCase(),
account.address.toLowerCase(),
snapId,
expect.any(Function),
undefined,
Expand Down Expand Up @@ -549,6 +589,83 @@ describe('SnapKeyring', () => {
});
expect(mockPublishedEventCallback).toHaveBeenCalledWith(event);
});

it('saves to the state asynchronously', async () => {
// We simulate a long running `saveState`
const deferred = new DeferredPromise<void>();
const saveStateContext = {
called: false,
returned: false,
};
mockCallbacks.saveState.mockImplementation(async () => {
saveStateContext.called = true;
await deferred.promise;
saveStateContext.returned = true;
});

const account = newEthEoaAccount;
const result = await keyring.handleKeyringSnapMessage(snapId, {
method: KeyringEvent.AccountCreated,
params: {
account: {
...(account as unknown as KeyringAccount),
id: account.id,
},
},
});
expect(result).toBeNull(); // Yes the result of `AccountCreated` is `null`.

// After reaching that point, the `AccountCreated` has resumed, so the Snap
// will be resuming its execution. However, the account is still not created
// on the Snap keyring state, since the `saveState` is still "pending".

// Now we can resolve, and finalize the `saveState` call.
expect(saveStateContext.called).toBe(true);
expect(saveStateContext.returned).toBe(false);
deferred.resolve();
await deferred.promise;
expect(saveStateContext.returned).toBe(true);
});

it('deletes the account if we cannot save it on the state', async () => {
mockCallbacks.saveState.mockImplementation(async () => {
return Promise.reject(new Error('Could not persist to the state'));
});
mockMessengerHandleRequest({
[KeyringRpcMethod.DeleteAccount]: () => null,
});

const account = newEthEoaAccount;
await keyring.handleKeyringSnapMessage(snapId, {
method: KeyringEvent.AccountCreated,
params: {
account: {
...(account as unknown as KeyringAccount),
id: account.id,
},
},
});
expect(mockCallbacks.addAccount).toHaveBeenLastCalledWith(
account.address.toLowerCase(),
snapId,
expect.any(Function),
undefined,
undefined,
);
expect(mockMessenger.handleRequest).toHaveBeenLastCalledWith({
handler: 'onKeyringRequest',
origin: 'metamask',
snapId,
request: {
id: expect.any(String),
jsonrpc: '2.0',
method: 'keyring_deleteAccount',
params: {
id: account.id,
},
},
});
});
});

describe('#handleAccountUpdated', () => {
Expand Down
41 changes: 39 additions & 2 deletions packages/keyring-snap-bridge/src/SnapKeyring.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,11 @@ export class SnapKeyring extends EventEmitter {
displayConfirmation,
} = message.params;

// READ THIS CAREFULLY:
// ------------------------------------------------------------------------
// The account creation flow is now asynchronous. We expect the Snap to
// first create the account data and then fire the "AccountCreated" event.

// Potentially migrate the account.
const account = transformAccount(newAccountFromEvent);

Expand All @@ -224,18 +229,40 @@ export class SnapKeyring extends EventEmitter {
throw new Error(`Account '${account.id}' already exists`);
}

// Add the account to the keyring, but wait for the MetaMask client to
// approve the account creation first.
await this.#callbacks.addAccount(
address,
snapId,
// This callback is passed to the MetaMask client, it will be called whenever
// the end user will accept or not the account creation.
async (accepted: boolean) => {
if (accepted) {
// We consider the account to be created on the Snap keyring only if
// the user accepted it. Meaning that the Snap MIGHT HAVE created the
// account on its own state, but the Snap keyring MIGHT NOT HAVE it yet.
//
// e.g The account creation dialog crashed on MetaMask, this callback
// will never be called, but the Snap still has the account.
this.#accounts.set(account.id, { account, snapId });
await this.#callbacks.saveState();

// This is the "true async part". We do not `await` for this call, mainly
// because this callback will persist the account on the client side
// (through the `AccountsController`).
//
// Since this will happen after the Snap account creation and Snap
// event, if anything goes wrong, we will delete the account by
// calling `deleteAccount` on the Snap.
// eslint-disable-next-line no-void
void this.#callbacks.saveState().catch(async () => {
await this.#deleteAccount(snapId, account);
});
}
},
accountNameSuggestion,
displayConfirmation,
);

return null;
}

Expand Down Expand Up @@ -1095,6 +1122,16 @@ export class SnapKeyring extends EventEmitter {
async removeAccount(address: string): Promise<void> {
const { account, snapId } = this.#resolveAddress(address);

await this.#deleteAccount(snapId, account);
}

/**
* Removes an account.
*
* @param snapId - Snap ID.
* @param account - Account to delete.
*/
async #deleteAccount(snapId: SnapId, account: KeyringAccount): Promise<void> {
// Always remove the account from the maps, even if the Snap is going to
// fail to delete it.
this.#accounts.delete(snapId, account.id);
Expand All @@ -1106,7 +1143,7 @@ export class SnapKeyring extends EventEmitter {
// with the account deletion, otherwise the account will be stuck in the
// keyring.
console.error(
`Account '${address}' may not have been removed from snap '${snapId}':`,
`Account '${account.address}' may not have been removed from snap '${snapId}':`,
error,
);
}
Expand Down

0 comments on commit ecdc034

Please sign in to comment.