diff --git a/packages/assets-controller/package.json b/packages/assets-controller/package.json index 17ac0ab59bb..8adad28f50b 100644 --- a/packages/assets-controller/package.json +++ b/packages/assets-controller/package.json @@ -54,12 +54,17 @@ "@metamask/base-controller": "^9.0.0", "@metamask/controller-utils": "^11.18.0", "@metamask/core-backend": "^5.0.0", + "@metamask/keyring-api": "^21.0.0", "@metamask/keyring-controller": "^25.1.0", "@metamask/keyring-internal-api": "^9.0.0", + "@metamask/keyring-snap-client": "^8.0.0", "@metamask/messenger": "^0.3.0", "@metamask/network-controller": "^29.0.0", "@metamask/network-enablement-controller": "^4.1.0", + "@metamask/permission-controller": "^12.2.0", "@metamask/polling-controller": "^16.0.2", + "@metamask/snaps-controllers": "^17.2.0", + "@metamask/snaps-utils": "^11.7.0", "@metamask/utils": "^11.9.0", "async-mutex": "^0.5.0", "bignumber.js": "^9.1.2", diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts index 22639fa3aaf..82fccf47f33 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.test.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.test.ts @@ -6,31 +6,35 @@ import type { MessengerActions, MessengerEvents, } from '@metamask/messenger'; +import type { + PermissionConstraint, + SubjectPermissions, +} from '@metamask/permission-controller'; import type { SnapDataSourceMessenger, SnapDataSourceOptions, - SnapProvider, AccountBalancesUpdatedEventPayload, } from './SnapDataSource'; import { SnapDataSource, createSnapDataSource, - getSnapTypeForChain, - isSnapSupportedChain, extractChainFromAssetId, - isSolanaChain, - isBitcoinChain, - isTronChain, - SOLANA_MAINNET, - SOLANA_SNAP_ID, - BITCOIN_MAINNET, - BITCOIN_SNAP_ID, - TRON_MAINNET, - ALL_DEFAULT_NETWORKS, + getChainIdsCaveat, + KEYRING_PERMISSION, + ASSETS_PERMISSION, } from './SnapDataSource'; import type { ChainId, DataRequest, Context, Caip19AssetId } from '../types'; +// Test chain IDs +const SOLANA_MAINNET = 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; +const BITCOIN_MAINNET = 'bip122:000000000019d6689c085ae165831e93' as ChainId; +const TRON_MAINNET = 'tron:728126428' as ChainId; + +// Test snap IDs +const SOLANA_SNAP_ID = 'npm:@metamask/solana-wallet-snap'; +const BITCOIN_SNAP_ID = 'npm:@metamask/bitcoin-wallet-snap'; + type AllActions = MessengerActions; type AllEvents = MessengerEvents; type RootMessenger = Messenger; @@ -46,7 +50,9 @@ const CHAIN_MAINNET = 'eip155:1' as ChainId; type SetupResult = { controller: SnapDataSource; messenger: RootMessenger; - snapProvider: jest.Mocked; + mockGetRunnableSnaps: jest.Mock; + mockHandleRequest: jest.Mock; + mockGetPermissions: jest.Mock; assetsUpdateHandler: jest.Mock; activeChainsUpdateHandler: jest.Mock; triggerBalancesUpdated: (payload: AccountBalancesUpdatedEventPayload) => void; @@ -68,6 +74,11 @@ function createMockAccount( keyring: { type: 'HD Key Tree' }, importTime: Date.now(), lastSelected: Date.now(), + snap: { + id: SOLANA_SNAP_ID, + name: 'Solana Snap', + enabled: true, + }, }, ...overrides, } as InternalAccount; @@ -91,44 +102,76 @@ function createMiddlewareContext(overrides?: Partial): Context { }; } -function createMockSnapProvider( - installedSnaps: Record = {}, +/** + * Creates mock permissions for PermissionController:getPermissions + * + * @param chainIds - Chain IDs this snap supports + * @returns Mock permissions object with both keyring and assets permissions + */ +function createMockPermissions( + chainIds: ChainId[] = [], +): SubjectPermissions { + if (chainIds.length === 0) { + return {}; + } + + return { + // Keyring permission indicates this is a keyring snap + [KEYRING_PERMISSION]: { + id: 'mock-keyring-permission-id', + parentCapability: KEYRING_PERMISSION, + invoker: 'test', + date: Date.now(), + caveats: null, + }, + // Assets permission contains the chainIds caveat + [ASSETS_PERMISSION]: { + id: 'mock-assets-permission-id', + parentCapability: ASSETS_PERMISSION, + invoker: 'test', + date: Date.now(), + caveats: [ + { + type: 'chainIds', + value: chainIds, + }, + ], + }, + } as unknown as SubjectPermissions; +} + +/** + * Creates a mock handler for SnapController:handleRequest + * + * @param accountAssets - Assets to return for keyring_listAccountAssets + * @param balances - Balances to return for keyring_getAccountBalances + * @returns Mock handler function + */ +function createMockHandleRequest( accountAssets: string[] = [], balances: Record = {}, -): jest.Mocked { - return { - request: jest.fn().mockImplementation(({ method, params }) => { - if (method === 'wallet_getSnaps') { - return Promise.resolve(installedSnaps); - } - if (method === 'wallet_invokeSnap') { - const snapRequest = params?.request; - if (snapRequest?.method === 'keyring_listAccountAssets') { - return Promise.resolve(accountAssets); - } - if (snapRequest?.method === 'keyring_getAccountBalances') { - return Promise.resolve(balances); - } - } - return Promise.resolve(null); - }), - }; +): jest.Mock { + return jest.fn().mockImplementation((params) => { + const { request } = params; + if (request?.method === 'keyring_listAccountAssets') { + return Promise.resolve(accountAssets); + } + if (request?.method === 'keyring_getAccountBalances') { + return Promise.resolve(balances); + } + return Promise.resolve(null); + }); } function setupController( options: { - installedSnaps?: Record; + installedSnaps?: Record; accountAssets?: string[]; balances?: Record; configuredNetworks?: ChainId[]; } = {}, ): SetupResult { - const { - installedSnaps = {}, - accountAssets = [], - balances = {}, - configuredNetworks, - } = options; + const { installedSnaps = {}, accountAssets = [], balances = {} } = options; const rootMessenger = new Messenger({ namespace: MOCK_ANY_NAMESPACE, @@ -149,6 +192,9 @@ function setupController( actions: [ 'AssetsController:assetsUpdate', 'AssetsController:activeChainsUpdate', + 'SnapController:getRunnableSnaps', + 'SnapController:handleRequest', + 'PermissionController:getPermissions', ], events: ['AccountsController:accountBalancesUpdated'], }); @@ -165,21 +211,50 @@ function setupController( activeChainsUpdateHandler, ); - const snapProvider = createMockSnapProvider( - installedSnaps, - accountAssets, - balances, + // Build snaps array for SnapController:getRunnableSnaps + // getRunnableSnaps returns only enabled, non-blocked snaps + const snapsForGetRunnableSnaps = Object.entries(installedSnaps).map( + ([id, { version }]) => ({ + id, + version, + enabled: true, + blocked: false, + }), + ); + + // Register SnapController action handlers + const mockGetRunnableSnaps = jest + .fn() + .mockReturnValue(snapsForGetRunnableSnaps); + rootMessenger.registerActionHandler( + 'SnapController:getRunnableSnaps', + mockGetRunnableSnaps, + ); + + const mockHandleRequest = createMockHandleRequest(accountAssets, balances); + rootMessenger.registerActionHandler( + 'SnapController:handleRequest', + mockHandleRequest, + ); + + // Register PermissionController:getPermissions handler + // Returns permissions with chainIds caveat based on installed snaps config + const mockGetPermissions = jest.fn().mockImplementation((snapId: string) => { + const snapConfig = installedSnaps[snapId]; + if (snapConfig?.chainIds) { + return createMockPermissions(snapConfig.chainIds); + } + return undefined; + }); + rootMessenger.registerActionHandler( + 'PermissionController:getPermissions', + mockGetPermissions, ); const controllerOptions: SnapDataSourceOptions = { messenger: controllerMessenger, - snapProvider, }; - if (configuredNetworks) { - controllerOptions.configuredNetworks = configuredNetworks; - } - const controller = new SnapDataSource(controllerOptions); const triggerBalancesUpdated = ( @@ -196,7 +271,9 @@ function setupController( return { controller, messenger: rootMessenger, - snapProvider, + mockGetRunnableSnaps, + mockHandleRequest, + mockGetPermissions, assetsUpdateHandler, activeChainsUpdateHandler, triggerBalancesUpdated, @@ -205,32 +282,6 @@ function setupController( } describe('SnapDataSource helper functions', () => { - describe('getSnapTypeForChain', () => { - it.each([ - { chainId: SOLANA_MAINNET, expected: 'solana' }, - { chainId: 'solana:devnet' as ChainId, expected: 'solana' }, - { chainId: BITCOIN_MAINNET, expected: 'bitcoin' }, - { chainId: 'bip122:testnet' as ChainId, expected: 'bitcoin' }, - { chainId: TRON_MAINNET, expected: 'tron' }, - { chainId: 'tron:0x2b6653dc' as ChainId, expected: 'tron' }, - { chainId: CHAIN_MAINNET, expected: null }, - { chainId: 'eip155:137' as ChainId, expected: null }, - ])('returns $expected for $chainId', ({ chainId, expected }) => { - expect(getSnapTypeForChain(chainId)).toBe(expected); - }); - }); - - describe('isSnapSupportedChain', () => { - it.each([ - { chainId: SOLANA_MAINNET, expected: true }, - { chainId: BITCOIN_MAINNET, expected: true }, - { chainId: TRON_MAINNET, expected: true }, - { chainId: CHAIN_MAINNET, expected: false }, - ])('returns $expected for $chainId', ({ chainId, expected }) => { - expect(isSnapSupportedChain(chainId)).toBe(expected); - }); - }); - describe('extractChainFromAssetId', () => { it.each([ { assetId: MOCK_SOL_ASSET, expected: SOLANA_MAINNET }, @@ -241,20 +292,42 @@ describe('SnapDataSource helper functions', () => { }); }); - describe('chain type helpers', () => { - it('isSolanaChain returns true for solana chains', () => { - expect(isSolanaChain(SOLANA_MAINNET)).toBe(true); - expect(isSolanaChain(BITCOIN_MAINNET)).toBe(false); + describe('getChainIdsCaveat', () => { + it('returns null when permission has no caveats', () => { + const permission = { + id: 'test', + parentCapability: KEYRING_PERMISSION, + invoker: 'test', + date: Date.now(), + caveats: null, + } as unknown as PermissionConstraint; + + expect(getChainIdsCaveat(permission)).toBeNull(); }); - it('isBitcoinChain returns true for bitcoin chains', () => { - expect(isBitcoinChain(BITCOIN_MAINNET)).toBe(true); - expect(isBitcoinChain(SOLANA_MAINNET)).toBe(false); + it('returns null when no chainIds caveat exists', () => { + const permission = { + id: 'test', + parentCapability: KEYRING_PERMISSION, + invoker: 'test', + date: Date.now(), + caveats: [{ type: 'other', value: [] }], + } as unknown as PermissionConstraint; + + expect(getChainIdsCaveat(permission)).toBeNull(); }); - it('isTronChain returns true for tron chains', () => { - expect(isTronChain(TRON_MAINNET)).toBe(true); - expect(isTronChain(SOLANA_MAINNET)).toBe(false); + it('returns chain IDs from chainIds caveat', () => { + const chainIds = [SOLANA_MAINNET, BITCOIN_MAINNET]; + const permission = { + id: 'test', + parentCapability: KEYRING_PERMISSION, + invoker: 'test', + date: Date.now(), + caveats: [{ type: 'chainIds', value: chainIds }], + } as unknown as PermissionConstraint; + + expect(getChainIdsCaveat(permission)).toStrictEqual(chainIds); }); }); }); @@ -271,24 +344,28 @@ describe('SnapDataSource', () => { cleanup(); }); - it('initializes with default networks', async () => { + it('initializes with empty chains when no keyring snaps discovered', async () => { const { controller, cleanup } = setupController(); await new Promise(process.nextTick); const chains = await controller.getActiveChains(); - expect(chains).toStrictEqual(ALL_DEFAULT_NETWORKS); + expect(chains).toStrictEqual([]); cleanup(); }); - it('initializes with configured networks', async () => { + it('discovers keyring snaps and populates active chains', async () => { const { controller, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET, BITCOIN_MAINNET], + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + [BITCOIN_SNAP_ID]: { version: '2.0.0', chainIds: [BITCOIN_MAINNET] }, + }, }); await new Promise(process.nextTick); const chains = await controller.getActiveChains(); - expect(chains).toStrictEqual([SOLANA_MAINNET, BITCOIN_MAINNET]); + expect(chains).toContain(SOLANA_MAINNET); + expect(chains).toContain(BITCOIN_MAINNET); cleanup(); }); @@ -304,129 +381,48 @@ describe('SnapDataSource', () => { cleanup(); }); - it('checks snap availability on initialization', async () => { - const { controller, snapProvider, cleanup } = setupController({ - installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, - }, - }); + it('checks snap availability on initialization via PermissionController', async () => { + const { mockGetRunnableSnaps, mockGetPermissions, cleanup } = + setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + }); await new Promise(process.nextTick); - expect(snapProvider.request).toHaveBeenCalledWith({ - method: 'wallet_getSnaps', - params: {}, - }); - expect(controller.isSnapAvailable('solana')).toBe(true); - expect(controller.isSnapAvailable('bitcoin')).toBe(false); + expect(mockGetRunnableSnaps).toHaveBeenCalled(); + expect(mockGetPermissions).toHaveBeenCalledWith(SOLANA_SNAP_ID); cleanup(); }); - it('getSnapsInfo returns all snap info', async () => { + it('fetch returns empty response for accounts without snap metadata', async () => { const { controller, cleanup } = setupController({ installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, - [BITCOIN_SNAP_ID]: { version: '2.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, }); await new Promise(process.nextTick); - const info = controller.getSnapsInfo(); - - expect(info.solana).toStrictEqual({ - snapId: SOLANA_SNAP_ID, - chainPrefix: 'solana:', - pollInterval: 30000, - version: '1.0.0', - available: true, - }); - expect(info.bitcoin).toStrictEqual({ - snapId: BITCOIN_SNAP_ID, - chainPrefix: 'bip122:', - pollInterval: 60000, - version: '2.0.0', - available: true, - }); - expect(info.tron.available).toBe(false); - - cleanup(); - }); - - it('refreshSnapsStatus updates availability', async () => { - const { controller, snapProvider, cleanup } = setupController(); - await new Promise(process.nextTick); - - expect(controller.isSnapAvailable('solana')).toBe(false); - - snapProvider.request.mockImplementation(({ method }) => { - if (method === 'wallet_getSnaps') { - return Promise.resolve({ - [SOLANA_SNAP_ID]: { version: '1.0.0' }, - }); - } - return Promise.resolve(null); - }); - - await controller.refreshSnapsStatus(); - - expect(controller.isSnapAvailable('solana')).toBe(true); - - cleanup(); - }); - - it('addNetworks adds snap-supported chains', async () => { - const { controller, activeChainsUpdateHandler, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET], - }); - await new Promise(process.nextTick); - - controller.addNetworks([BITCOIN_MAINNET, CHAIN_MAINNET]); - - expect(activeChainsUpdateHandler).toHaveBeenCalledWith( - 'SnapDataSource', - expect.arrayContaining([SOLANA_MAINNET, BITCOIN_MAINNET]), - ); - - cleanup(); - }); - - it('addNetworks ignores non-snap chains', async () => { - const { controller, activeChainsUpdateHandler, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET], - }); - await new Promise(process.nextTick); - - controller.addNetworks([CHAIN_MAINNET]); - - expect(activeChainsUpdateHandler).not.toHaveBeenCalled(); - - cleanup(); - }); - - it('removeNetworks removes chains', async () => { - const { controller, activeChainsUpdateHandler, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET, BITCOIN_MAINNET], + // Create account without snap metadata (non-snap account) + const nonSnapAccount = createMockAccount({ + metadata: { + name: 'Test Account', + keyring: { type: 'HD Key Tree' }, + importTime: Date.now(), + lastSelected: Date.now(), + // No snap property + }, }); - await new Promise(process.nextTick); - - controller.removeNetworks([SOLANA_MAINNET]); - - expect(activeChainsUpdateHandler).toHaveBeenCalledWith('SnapDataSource', [ - BITCOIN_MAINNET, - ]); - - cleanup(); - }); - - it('fetch returns empty response for non-snap chains', async () => { - const { controller, cleanup } = setupController(); - await new Promise(process.nextTick); const response = await controller.fetch( - createDataRequest({ chainIds: [CHAIN_MAINNET] }), + createDataRequest({ accounts: [nonSnapAccount] }), ); - expect(response).toStrictEqual({}); + expect(response).toStrictEqual({ + assetsBalance: {}, + assetsMetadata: {}, + }); cleanup(); }); @@ -444,23 +440,31 @@ describe('SnapDataSource', () => { cleanup(); }); - it('fetch returns error when snap not available', async () => { + it('fetch returns empty response when accounts array is empty', async () => { const { controller, cleanup } = setupController({ - installedSnaps: {}, + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, }); await new Promise(process.nextTick); - const response = await controller.fetch(createDataRequest()); + const response = await controller.fetch( + createDataRequest({ accounts: [] }), + ); - expect(response.errors?.[SOLANA_MAINNET]).toBe('solana snap not available'); + // No accounts to fetch, so empty balances + expect(response).toStrictEqual({ + assetsBalance: {}, + assetsMetadata: {}, + }); cleanup(); }); - it('fetch calls snap keyring methods when snap available', async () => { - const { controller, snapProvider, cleanup } = setupController({ + it('fetch calls snap keyring methods when snap discovered', async () => { + const { controller, mockHandleRequest, cleanup } = setupController({ installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, accountAssets: [MOCK_SOL_ASSET], balances: { @@ -471,26 +475,30 @@ describe('SnapDataSource', () => { const response = await controller.fetch(createDataRequest()); - expect(snapProvider.request).toHaveBeenCalledWith({ - method: 'wallet_invokeSnap', - params: { + // Verify keyring_listAccountAssets was called + expect(mockHandleRequest).toHaveBeenCalledWith( + expect.objectContaining({ snapId: SOLANA_SNAP_ID, - request: { + origin: 'metamask', + handler: 'onKeyringRequest', + request: expect.objectContaining({ method: 'keyring_listAccountAssets', params: { id: 'mock-account-id' }, - }, - }, - }); - expect(snapProvider.request).toHaveBeenCalledWith({ - method: 'wallet_invokeSnap', - params: { + }), + }), + ); + // Verify keyring_getAccountBalances was called + expect(mockHandleRequest).toHaveBeenCalledWith( + expect.objectContaining({ snapId: SOLANA_SNAP_ID, - request: { + origin: 'metamask', + handler: 'onKeyringRequest', + request: expect.objectContaining({ method: 'keyring_getAccountBalances', params: { id: 'mock-account-id', assets: [MOCK_SOL_ASSET] }, - }, - }, - }); + }), + }), + ); expect( response.assetsBalance?.['mock-account-id']?.[MOCK_SOL_ASSET], ).toStrictEqual({ amount: '1000000000' }); @@ -498,41 +506,10 @@ describe('SnapDataSource', () => { cleanup(); }); - it('fetch skips accounts without supported scopes', async () => { - const { controller, snapProvider, cleanup } = setupController({ - installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, - }, - accountAssets: [MOCK_SOL_ASSET], - balances: { - [MOCK_SOL_ASSET]: { amount: '1000000000', unit: 'SOL' }, - }, - }); - await new Promise(process.nextTick); - - const evmOnlyAccount = createMockAccount({ - scopes: ['eip155:0'], - }); - - await controller.fetch( - createDataRequest({ - accounts: [evmOnlyAccount], - }), - ); - - const invokeSnapCalls = snapProvider.request.mock.calls.filter((call) => { - const arg = call[0] as { method: string }; - return arg.method === 'wallet_invokeSnap'; - }); - expect(invokeSnapCalls).toHaveLength(0); - - cleanup(); - }); - it('fetch handles empty account assets gracefully', async () => { - const { controller, snapProvider, cleanup } = setupController({ + const { controller, mockHandleRequest, cleanup } = setupController({ installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, accountAssets: [], }); @@ -540,15 +517,10 @@ describe('SnapDataSource', () => { const response = await controller.fetch(createDataRequest()); - const getBalancesCalls = snapProvider.request.mock.calls.filter((call) => { - const arg = call[0] as { - method: string; - params?: { request?: { method: string } }; - }; - return ( - arg.method === 'wallet_invokeSnap' && - arg.params?.request?.method === 'keyring_getAccountBalances' - ); + // Check that keyring_getAccountBalances was NOT called (since no assets) + const getBalancesCalls = mockHandleRequest.mock.calls.filter((call) => { + const params = call[0] as { request?: { method: string } }; + return params.request?.method === 'keyring_getAccountBalances'; }); expect(getBalancesCalls).toHaveLength(0); expect(response.assetsBalance).toStrictEqual({}); @@ -559,8 +531,8 @@ describe('SnapDataSource', () => { it('fetch merges results from multiple snaps', async () => { const { controller, cleanup } = setupController({ installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, - [BITCOIN_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + [BITCOIN_SNAP_ID]: { version: '1.0.0', chainIds: [BITCOIN_MAINNET] }, }, accountAssets: [MOCK_SOL_ASSET, MOCK_BTC_ASSET], balances: { @@ -583,7 +555,11 @@ describe('SnapDataSource', () => { it('handles snap balances updated event', async () => { const { triggerBalancesUpdated, assetsUpdateHandler, cleanup } = - setupController(); + setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + }); await new Promise(process.nextTick); triggerBalancesUpdated({ @@ -610,9 +586,13 @@ describe('SnapDataSource', () => { cleanup(); }); - it('filters non-snap assets from balance update event', async () => { + it('filters assets for chains without discovered snaps from balance update event', async () => { const { triggerBalancesUpdated, assetsUpdateHandler, cleanup } = - setupController(); + setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + }); await new Promise(process.nextTick); const evmAsset = 'eip155:1/slip44:60' as Caip19AssetId; @@ -644,11 +624,16 @@ describe('SnapDataSource', () => { it('does not report empty balance updates', async () => { const { triggerBalancesUpdated, assetsUpdateHandler, cleanup } = - setupController(); + setupController({ + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, + }); await new Promise(process.nextTick); const evmAsset = 'eip155:1/slip44:60' as Caip19AssetId; + // Only EVM asset - no discovered snap for EVM triggerBalancesUpdated({ balances: { 'account-1': { @@ -667,7 +652,7 @@ describe('SnapDataSource', () => { it('subscribe performs initial fetch', async () => { const { controller, assetsUpdateHandler, cleanup } = setupController({ installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, accountAssets: [MOCK_SOL_ASSET], balances: { @@ -687,7 +672,7 @@ describe('SnapDataSource', () => { cleanup(); }); - it('subscribe does nothing for non-snap chains', async () => { + it('subscribe does nothing for chains without discovered snaps', async () => { const { controller, assetsUpdateHandler, cleanup } = setupController(); await new Promise(process.nextTick); @@ -705,7 +690,7 @@ describe('SnapDataSource', () => { it('subscribe update fetches data', async () => { const { controller, assetsUpdateHandler, cleanup } = setupController({ installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, accountAssets: [MOCK_SOL_ASSET], balances: { @@ -739,7 +724,9 @@ describe('SnapDataSource', () => { it('middleware passes to next when no supported chains', async () => { const { controller, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET], + installedSnaps: { + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, + }, }); await new Promise(process.nextTick); @@ -757,9 +744,8 @@ describe('SnapDataSource', () => { it('middleware merges response into context', async () => { const { controller, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET], installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, accountAssets: [MOCK_SOL_ASSET], balances: { @@ -783,9 +769,8 @@ describe('SnapDataSource', () => { it('middleware removes handled chains from next request', async () => { const { controller, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET], installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, accountAssets: [MOCK_SOL_ASSET], balances: { @@ -814,9 +799,9 @@ describe('SnapDataSource', () => { cleanup(); }); - it('middleware keeps failed chains in request', async () => { + it('middleware keeps chains without discovered snaps in request', async () => { + // No snaps discovered const { controller, cleanup } = setupController({ - configuredNetworks: [SOLANA_MAINNET], installedSnaps: {}, }); await new Promise(process.nextTick); @@ -830,13 +815,8 @@ describe('SnapDataSource', () => { await controller.assetsMiddleware(context, next); - expect(next).toHaveBeenCalledWith( - expect.objectContaining({ - request: expect.objectContaining({ - chainIds: expect.arrayContaining([SOLANA_MAINNET, CHAIN_MAINNET]), - }), - }), - ); + // All chains passed through since no snaps handle any chains + expect(next).toHaveBeenCalledWith(context); cleanup(); }); @@ -844,7 +824,7 @@ describe('SnapDataSource', () => { it('destroy cleans up subscriptions', async () => { const { controller, cleanup } = setupController({ installedSnaps: { - [SOLANA_SNAP_ID]: { version: '1.0.0' }, + [SOLANA_SNAP_ID]: { version: '1.0.0', chainIds: [SOLANA_MAINNET] }, }, accountAssets: [MOCK_SOL_ASSET], balances: { @@ -889,6 +869,9 @@ describe('SnapDataSource', () => { actions: [ 'AssetsController:assetsUpdate', 'AssetsController:activeChainsUpdate', + 'SnapController:getRunnableSnaps', + 'SnapController:handleRequest', + 'PermissionController:getPermissions', ], events: ['AccountsController:accountBalancesUpdated'], }); @@ -901,12 +884,21 @@ describe('SnapDataSource', () => { 'AssetsController:activeChainsUpdate', jest.fn(), ); - - const snapProvider = createMockSnapProvider(); + rootMessenger.registerActionHandler( + 'SnapController:getRunnableSnaps', + jest.fn().mockReturnValue([]), + ); + rootMessenger.registerActionHandler( + 'SnapController:handleRequest', + jest.fn(), + ); + rootMessenger.registerActionHandler( + 'PermissionController:getPermissions', + jest.fn().mockReturnValue(undefined), + ); const instance = createSnapDataSource({ messenger: controllerMessenger, - snapProvider, }); await new Promise(process.nextTick); diff --git a/packages/assets-controller/src/data-sources/SnapDataSource.ts b/packages/assets-controller/src/data-sources/SnapDataSource.ts index 4fb5f031f10..10a7117de26 100644 --- a/packages/assets-controller/src/data-sources/SnapDataSource.ts +++ b/packages/assets-controller/src/data-sources/SnapDataSource.ts @@ -1,5 +1,19 @@ -import type { InternalAccount } from '@metamask/keyring-internal-api'; +import type { Balance, CaipAssetType } from '@metamask/keyring-api'; +import { KeyringClient } from '@metamask/keyring-snap-client'; import type { Messenger } from '@metamask/messenger'; +import type { + Caveat, + GetPermissions, + PermissionConstraint, + SubjectPermissions, +} from '@metamask/permission-controller'; +import type { + GetRunnableSnaps, + HandleSnapRequest, +} from '@metamask/snaps-controllers'; +import type { Snap, SnapId } from '@metamask/snaps-sdk'; +import { HandlerType, SnapCaveatType } from '@metamask/snaps-utils'; +import type { Json, JsonRpcRequest } from '@metamask/utils'; import { AbstractDataSource } from './AbstractDataSource'; import type { @@ -50,120 +64,38 @@ const log = createModuleLogger(projectLogger, 'SnapDataSource'); export const SNAP_DATA_SOURCE_NAME = 'SnapDataSource'; -// Snap IDs -export const SOLANA_SNAP_ID = 'npm:@metamask/solana-wallet-snap'; -export const BITCOIN_SNAP_ID = 'npm:@metamask/bitcoin-wallet-snap'; -export const TRON_SNAP_ID = 'npm:@metamask/tron-wallet-snap'; - -// Chain prefixes for detection -export const SOLANA_CHAIN_PREFIX = 'solana:'; -export const BITCOIN_CHAIN_PREFIX = 'bip122:'; -export const TRON_CHAIN_PREFIX = 'tron:'; - -// Default networks -export const SOLANA_MAINNET = - 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp' as ChainId; -export const SOLANA_DEVNET = - 'solana:EtWTRABZaYq6iMfeYKouRu166VU2xqa1' as ChainId; -export const SOLANA_TESTNET = - 'solana:4uhcVJyU9pJkvQyS88uRDiswHXSCkY3z' as ChainId; - -export const BITCOIN_MAINNET = - 'bip122:000000000019d6689c085ae165831e93' as ChainId; -export const BITCOIN_TESTNET = - 'bip122:000000000933ea01ad0ee984209779ba' as ChainId; - -export const TRON_MAINNET = 'tron:728126428' as ChainId; -export const TRON_SHASTA = 'tron:2494104990' as ChainId; -export const TRON_NILE = 'tron:3448148188' as ChainId; -// Hex format alternatives for Tron -export const TRON_MAINNET_HEX = 'tron:0x2b6653dc' as ChainId; -export const TRON_SHASTA_HEX = 'tron:0x94a9059e' as ChainId; -export const TRON_NILE_HEX = 'tron:0xcd8690dc' as ChainId; - -// Default poll intervals -export const DEFAULT_SOLANA_POLL_INTERVAL = 30_000; // 30 seconds -export const DEFAULT_BITCOIN_POLL_INTERVAL = 60_000; // 1 minute -export const DEFAULT_TRON_POLL_INTERVAL = 30_000; // 30 seconds -export const DEFAULT_SNAP_POLL_INTERVAL = 30_000; // Default for unknown snaps - -// All default networks -export const ALL_DEFAULT_NETWORKS: ChainId[] = [ - SOLANA_MAINNET, - SOLANA_DEVNET, - SOLANA_TESTNET, - BITCOIN_MAINNET, - BITCOIN_TESTNET, - TRON_MAINNET, - TRON_SHASTA, - TRON_NILE, - TRON_MAINNET_HEX, - TRON_SHASTA_HEX, - TRON_NILE_HEX, -]; +/** The permission name for snap keyring endowment */ +export const KEYRING_PERMISSION = 'endowment:keyring'; + +/** The permission name for snap assets endowment (contains chainIds) */ +export const ASSETS_PERMISSION = 'endowment:assets'; // ============================================================================ -// SNAP ROUTING +// PERMISSION UTILITIES // ============================================================================ -export type SnapType = 'solana' | 'bitcoin' | 'tron'; - -export type SnapInfo = { - snapId: string; - chainPrefix: string; - pollInterval: number; - version: string | null; - available: boolean; -}; - -export const SNAP_REGISTRY: Record< - SnapType, - Omit -> = { - solana: { - snapId: SOLANA_SNAP_ID, - chainPrefix: SOLANA_CHAIN_PREFIX, - pollInterval: DEFAULT_SOLANA_POLL_INTERVAL, - }, - bitcoin: { - snapId: BITCOIN_SNAP_ID, - chainPrefix: BITCOIN_CHAIN_PREFIX, - pollInterval: DEFAULT_BITCOIN_POLL_INTERVAL, - }, - tron: { - snapId: TRON_SNAP_ID, - chainPrefix: TRON_CHAIN_PREFIX, - pollInterval: DEFAULT_TRON_POLL_INTERVAL, - }, -}; - /** - * Get the snap type for a chain ID based on its prefix. + * Getter function to get the chainIds caveat from a permission. + * + * This does basic validation of the caveat, but does not validate the type or + * value of the namespaces object itself, as this is handled by the + * `PermissionsController` when the permission is requested. * - * @param chainId - The CAIP-2 chain ID to check. - * @returns The snap type for the chain, or null if not supported. + * @param permission - The permission to get the `chainIds` caveat from. + * @returns An array of `chainIds` that the snap supports, or null if none. */ -export function getSnapTypeForChain(chainId: ChainId): SnapType | null { - if (chainId.startsWith(SOLANA_CHAIN_PREFIX)) { - return 'solana'; - } - if (chainId.startsWith(BITCOIN_CHAIN_PREFIX)) { - return 'bitcoin'; +export function getChainIdsCaveat( + permission?: PermissionConstraint, +): ChainId[] | null { + if (!permission?.caveats) { + return null; } - if (chainId.startsWith(TRON_CHAIN_PREFIX)) { - return 'tron'; - } - return null; -} -/** - * Check if a chain ID is supported by a snap. - * - * @param chainId - The CAIP-2 chain ID to check. - * @returns True if the chain is supported by a snap. - */ -export function isSnapSupportedChain(chainId: ChainId): boolean { - return getSnapTypeForChain(chainId) !== null; + const caveat = permission.caveats.find( + (permCaveat) => permCaveat.type === SnapCaveatType.ChainIds, + ) as Caveat | undefined; + + return caveat ? (caveat.value as ChainId[]) : null; } /** @@ -178,35 +110,25 @@ export function extractChainFromAssetId(assetId: string): ChainId { return parts[0] as ChainId; } -// Helper functions for specific chain types -export function isSolanaChain(chainId: ChainId): boolean { - return chainId.startsWith(SOLANA_CHAIN_PREFIX); -} - -export function isBitcoinChain(chainId: ChainId): boolean { - return chainId.startsWith(BITCOIN_CHAIN_PREFIX); -} - -export function isTronChain(chainId: ChainId): boolean { - return chainId.startsWith(TRON_CHAIN_PREFIX); -} - // ============================================================================ // STATE // ============================================================================ +/** + * State for the SnapDataSource. + * Uses dynamic snap discovery - chains are populated from PermissionController. + */ export type SnapDataSourceState = { - /** Snap availability and versions */ - snaps: Record; + /** + * Mapping of chain IDs to snap IDs that support them. + * Used to filter which accounts to process for a given chain request. + */ + chainToSnap: Record; } & DataSourceState; const defaultSnapState: SnapDataSourceState = { - activeChains: ALL_DEFAULT_NETWORKS, - snaps: { - solana: { version: null, available: false }, - bitcoin: { version: null, available: false }, - tron: { version: null, available: false }, - }, + activeChains: [], + chainToSnap: {}, }; // ============================================================================ @@ -278,7 +200,10 @@ type AssetsControllerAssetsUpdateAction = { export type SnapDataSourceAllowedActions = | AssetsControllerActiveChainsUpdateAction - | AssetsControllerAssetsUpdateAction; + | AssetsControllerAssetsUpdateAction + | GetRunnableSnaps + | HandleSnapRequest + | GetPermissions; export type SnapDataSourceMessenger = Messenger< typeof SNAP_DATA_SOURCE_NAME, @@ -286,17 +211,6 @@ export type SnapDataSourceMessenger = Messenger< SnapDataSourceEvents | SnapDataSourceAllowedEvents >; -// ============================================================================ -// SNAP PROVIDER INTERFACE -// ============================================================================ - -export type SnapProvider = { - request(args: { - method: string; - params?: unknown; - }): Promise; -}; - // ============================================================================ // OPTIONS // ============================================================================ @@ -304,10 +218,6 @@ export type SnapProvider = { export type SnapDataSourceOptions = { /** Messenger for this data source */ messenger: SnapDataSourceMessenger; - /** - * Snap provider for communicating with snaps. - */ - snapProvider: SnapProvider; /** Configured networks to support (defaults to all snap networks) */ configuredNetworks?: ChainId[]; /** Default polling interval in ms for subscriptions */ @@ -324,16 +234,10 @@ export type SnapDataSourceOptions = { * Unified Snap data source that routes requests to the appropriate wallet snap * based on the chain ID prefix. * - * Supports: - * - Solana chains (solana:*) → @metamask/solana-wallet-snap - * - Bitcoin chains (bip122:*) → @metamask/bitcoin-wallet-snap - * - Tron chains (tron:*) → @metamask/tron-wallet-snap - * * @example * ```typescript * const snapDataSource = new SnapDataSource({ * messenger, - * snapProvider: metamaskProvider, * }); * * // Fetch will automatically route to the correct snap @@ -349,28 +253,29 @@ export class SnapDataSource extends AbstractDataSource< > { readonly #messenger: SnapDataSourceMessenger; - readonly #snapProvider: SnapProvider; + /** Bound handler for snap keyring balance updates, stored for cleanup */ + readonly #handleSnapBalancesUpdatedBound: ( + payload: AccountBalancesUpdatedEventPayload, + ) => void; constructor(options: SnapDataSourceOptions) { - const configuredNetworks = - options.configuredNetworks ?? ALL_DEFAULT_NETWORKS; - super(SNAP_DATA_SOURCE_NAME, { ...defaultSnapState, ...options.state, - activeChains: configuredNetworks, }); this.#messenger = options.messenger; - this.#snapProvider = options.snapProvider; + + // Bind handler for cleanup in destroy() + this.#handleSnapBalancesUpdatedBound = this.#handleSnapBalancesUpdated.bind( + this, + ) as (payload: AccountBalancesUpdatedEventPayload) => void; this.#registerActionHandlers(); this.#subscribeToSnapKeyringEvents(); - // Check availability for all snaps on initialization - this.#checkAllSnapsAvailability().catch(() => { - // Silently ignore availability check failures on init - }); + // Discover keyring-capable snaps and populate activeChains dynamically + this.#discoverKeyringSnaps(); } /** @@ -385,9 +290,7 @@ export class SnapDataSource extends AbstractDataSource< try { messenger.subscribe( 'AccountsController:accountBalancesUpdated', - (payload: AccountBalancesUpdatedEventPayload) => { - this.#handleSnapBalancesUpdated(payload); - }, + this.#handleSnapBalancesUpdatedBound, ); } catch (error) { log('Failed to subscribe to snap keyring events', { error }); @@ -413,8 +316,8 @@ export class SnapDataSource extends AbstractDataSource< response.assetsBalance[accountId] = {}; for (const [assetId, balance] of Object.entries(assets)) { - // Only include snap-supported assets (solana, bitcoin, tron) - if (isSnapSupportedChain(extractChainFromAssetId(assetId))) { + const chainId = extractChainFromAssetId(assetId); + if (this.#isChainSupportedBySnap(chainId)) { response.assetsBalance[accountId][assetId as Caip19AssetId] = { amount: balance.amount, }; @@ -436,6 +339,16 @@ export class SnapDataSource extends AbstractDataSource< } } + /** + * Check if a chain ID is supported by any discovered snap. + * + * @param chainId - The CAIP-2 chain ID to check. + * @returns True if we have a snap that supports this chain. + */ + #isChainSupportedBySnap(chainId: ChainId): boolean { + return this.state.activeChains.includes(chainId); + } + #registerActionHandlers(): void { // eslint-disable-next-line @typescript-eslint/no-explicit-any const messenger = this.#messenger as any; @@ -467,447 +380,183 @@ export class SnapDataSource extends AbstractDataSource< } // ============================================================================ - // SNAP AVAILABILITY + // SNAP DISCOVERY (Dynamic via PermissionController) // ============================================================================ /** - * Get all installed snaps from the snap provider. + * Get all runnable snaps from SnapController. + * Runnable snaps are enabled and not blocked. * - * @returns A map of snap IDs to their versions. + * @returns Array of runnable snaps. */ - async #getInstalledSnaps(): Promise> { + #getRunnableSnaps(): Snap[] { try { - const snaps = await this.#snapProvider.request< - Record - >({ - method: 'wallet_getSnaps', - params: {}, - }); - - return snaps; + return this.#messenger.call('SnapController:getRunnableSnaps'); } catch (error) { - log('Failed to get installed snaps', error); - return {}; + log('Failed to get runnable snaps', error); + return []; } } /** - * Check availability for a single snap type on-demand. - * This is called before each fetch to ensure we have the latest availability status. + * Get permissions for a snap from PermissionController. * - * @param snapType - The snap type to check (solana, bitcoin, tron) - * @returns True if the snap is available, false otherwise + * @param snapId - The snap ID to get permissions for. + * @returns The snap's permissions, or undefined if none. */ - async #checkSnapAvailabilityOnDemand(snapType: SnapType): Promise { - const config = SNAP_REGISTRY[snapType]; - const currentState = this.state.snaps[snapType]; - - // If already marked as available, return true (snap was found previously) - if (currentState.available) { - return true; - } - - // Check if snap is now available (handles timing issues where snap wasn't ready at init) + #getSnapPermissions( + snapId: string, + ): SubjectPermissions | undefined { try { - const snaps = await this.#getInstalledSnaps(); - const snap = snaps[config.snapId]; - - if (snap) { - // Snap is now available - update state - this.state.snaps[snapType] = { - version: snap.version, - available: true, - }; - return true; - } - - return false; - } catch { - return false; - } - } - - async #checkAllSnapsAvailability(): Promise { - try { - const snaps = await this.#getInstalledSnaps(); - - for (const [snapType, config] of Object.entries(SNAP_REGISTRY)) { - const snap = snaps[config.snapId]; - - if (snap) { - this.state.snaps[snapType as SnapType] = { - version: snap.version, - available: true, - }; - } else { - this.state.snaps[snapType as SnapType] = { - version: null, - available: false, - }; - } - } - } catch { - // Mark all snaps as unavailable on error - for (const snapType of Object.keys(SNAP_REGISTRY)) { - this.state.snaps[snapType as SnapType] = { - version: null, - available: false, - }; - } - } - } - - /** - * Get info about all snaps. - * - * @returns Record of snap info keyed by snap type. - */ - getSnapsInfo(): Record { - const result: Record = {} as Record; - - for (const [snapType, config] of Object.entries(SNAP_REGISTRY)) { - const state = this.state.snaps[snapType as SnapType]; - result[snapType as SnapType] = { - ...config, - version: state.version, - available: state.available, - }; - } - - return result; - } - - /** - * Check if a specific snap is available. - * - * @param snapType - The snap type to check (solana, bitcoin, tron). - * @returns True if the snap is available. - */ - isSnapAvailable(snapType: SnapType): boolean { - return this.state.snaps[snapType]?.available ?? false; - } - - /** - * Force refresh snap availability check. - */ - async refreshSnapsStatus(): Promise { - await this.#checkAllSnapsAvailability(); - } - - // ============================================================================ - // CHAIN MANAGEMENT - // ============================================================================ - - addNetworks(chainIds: ChainId[]): void { - const snapChains = chainIds.filter(isSnapSupportedChain); - const newChains = snapChains.filter( - (chain) => !this.state.activeChains.includes(chain), - ); - - if (newChains.length > 0) { - const updated = [...this.state.activeChains, ...newChains]; - this.updateActiveChains(updated, (updatedChains) => - this.#messenger.call( - 'AssetsController:activeChainsUpdate', - SNAP_DATA_SOURCE_NAME, - updatedChains, - ), - ); - } - } - - removeNetworks(chainIds: ChainId[]): void { - const chainSet = new Set(chainIds); - const updated = this.state.activeChains.filter( - (chain) => !chainSet.has(chain), - ); - if (updated.length !== this.state.activeChains.length) { - this.updateActiveChains(updated, (updatedChains) => - this.#messenger.call( - 'AssetsController:activeChainsUpdate', - SNAP_DATA_SOURCE_NAME, - updatedChains, - ), - ); + return this.#messenger.call( + 'PermissionController:getPermissions', + snapId, + ) as SubjectPermissions; + } catch (error) { + log('Failed to get permissions for snap', { snapId, error }); + return undefined; } } - // ============================================================================ - // ACCOUNT SCOPE HELPERS - // ============================================================================ - /** - * Check if an account supports a specific chain based on its scopes. - * For snap chains (Solana, Bitcoin, Tron), we check for the appropriate namespace. - * - * @param account - The account to check - * @param chainId - The chain ID to check (e.g., "solana:...", "bip122:...", "tron:...") - * @returns True if the account supports the chain + * Discover all snaps with keyring capabilities and their supported chains. + * Uses PermissionController to find snaps with endowment:keyring permission. + * Updates chainToSnap mapping and activeChains. */ - #accountSupportsChain(account: InternalAccount, chainId: ChainId): boolean { - const scopes = account.scopes ?? []; - - // If no scopes defined, assume it supports the chain (backward compatibility) - if (scopes.length === 0) { - return true; - } + #discoverKeyringSnaps(): void { + try { + const runnableSnaps = this.#getRunnableSnaps(); + const chainToSnap: Record = {}; + const supportedChains: ChainId[] = []; - // Extract namespace and reference from chainId - const [chainNamespace, chainReference] = chainId.split(':'); + for (const snap of runnableSnaps) { + const permissions = this.#getSnapPermissions(snap.id); + if (!permissions) { + continue; + } - for (const scope of scopes) { - const [scopeNamespace, scopeReference] = (scope as string).split(':'); + // Must have endowment:keyring permission to be a keyring snap + const keyringPermission = permissions[KEYRING_PERMISSION]; + if (!keyringPermission) { + continue; + } - // Check if namespaces match - if (scopeNamespace !== chainNamespace) { - continue; + // Get chainIds caveat from the assets permission (not keyring permission) + // The chainIds are stored in endowment:assets + const assetsPermission = permissions[ASSETS_PERMISSION]; + const chainIds = getChainIdsCaveat(assetsPermission); + + // Map each chain to this snap (first snap wins if multiple support same chain) + if (chainIds) { + for (const chainId of chainIds) { + if (!(chainId in chainToSnap)) { + chainToSnap[chainId] = snap.id; + supportedChains.push(chainId); + } + } + } } - // Wildcard scope (e.g., "solana:0" means all chains in that namespace) - if (scopeReference === '0') { - return true; - } + // Update chainToSnap mapping + this.state.chainToSnap = chainToSnap; - // Exact match check - if (scopeReference === chainReference) { - return true; + // Notify if chains changed (wrapped in try-catch for init timing) + try { + this.updateActiveChains(supportedChains, (updatedChains) => { + this.#messenger.call( + 'AssetsController:activeChainsUpdate', + SNAP_DATA_SOURCE_NAME, + updatedChains, + ); + }); + } catch { + // AssetsController not ready yet - expected during initialization + } + } catch (error) { + log('Keyring snap discovery failed', { error }); + this.state.chainToSnap = {}; + try { + this.updateActiveChains([], (updatedChains) => { + this.#messenger.call( + 'AssetsController:activeChainsUpdate', + SNAP_DATA_SOURCE_NAME, + updatedChains, + ); + }); + } catch { + // AssetsController not ready yet - expected during initialization } } - - return false; } // ============================================================================ - // FETCH - Routes to appropriate snap(s) + // FETCH // ============================================================================ async fetch(request: DataRequest): Promise { - // Guard against undefined request or chainIds - if (!request?.chainIds) { + // Guard against undefined request + // Note: chainIds filtering is done by middleware/subscribe before calling fetch + if (!request?.accounts || !request?.chainIds?.length) { return {}; } - // Filter to only snap-supported chains - const snapChains = request.chainIds.filter(isSnapSupportedChain); - - if (snapChains.length === 0) { - return {}; - } - - // Group chains by snap type - const chainsBySnap = this.#groupChainsBySnap(snapChains); - - // Fetch from each snap in parallel - const results = await Promise.all( - Object.entries(chainsBySnap).map(async ([snapType, chains]) => { - return this.#fetchFromSnap(snapType as SnapType, { - ...request, - chainIds: chains, - }); - }), - ); - - // Merge all results - const mergedResponse: DataResponse = {}; - - for (const result of results) { - if (result.assetsBalance) { - mergedResponse.assetsBalance = { - ...mergedResponse.assetsBalance, - ...result.assetsBalance, - }; - } - if (result.assetsMetadata) { - mergedResponse.assetsMetadata = { - ...mergedResponse.assetsMetadata, - ...result.assetsMetadata, - }; - } - if (result.assetsPrice) { - mergedResponse.assetsPrice = { - ...mergedResponse.assetsPrice, - ...result.assetsPrice, - }; - } - if (result.errors) { - mergedResponse.errors = { ...mergedResponse.errors, ...result.errors }; - } - } - - return mergedResponse; - } - - #groupChainsBySnap( - chainIds: ChainId[], - ): Partial> { - const groups: Partial> = {}; - - for (const chainId of chainIds) { - const snapType = getSnapTypeForChain(chainId); - if (snapType) { - groups[snapType] ??= []; - const snapChains = groups[snapType]; - if (snapChains) { - snapChains.push(chainId); - } - } - } - - return groups; - } - - async #fetchFromSnap( - snapType: SnapType, - request: DataRequest, - ): Promise { - const config = SNAP_REGISTRY[snapType]; - - // Check snap availability on-demand - handles timing issues where snap - // wasn't ready during initialization but is now available - const isAvailable = await this.#checkSnapAvailabilityOnDemand(snapType); - if (!isAvailable) { - log(`${snapType} snap not available, skipping fetch`); - // Return errors for these chains so they can fallback to other data sources - const errors: Record = {}; - for (const chainId of request.chainIds) { - errors[chainId] = `${snapType} snap not available`; - } - return { errors }; - } - const results: DataResponse = { assetsBalance: {}, assetsMetadata: {}, }; - // Fetch balances for each account using Keyring API - // Important: Must first get account assets, then request balances for those specific assets + // Fetch balances for each account using its snap ID from metadata for (const account of request.accounts) { - // Filter to only process accounts that support the chains being fetched - const accountSupportedChains = request.chainIds.filter((chainId) => - this.#accountSupportsChain(account, chainId), - ); + // Skip accounts without snap metadata (non-snap accounts) + const snapId = account.metadata.snap?.id; + if (!snapId) { + continue; + } - // Skip accounts that don't support any of the requested chains - if (accountSupportedChains.length === 0) { + // Skip accounts whose snap doesn't support any of the requested chains + const snapSupportsRequestedChains = request.chainIds.some( + (chainId) => this.state.chainToSnap[chainId] === snapId, + ); + if (!snapSupportsRequestedChains) { continue; } const accountId = account.id; try { - // Step 1: Get the list of assets for this account - log(`${snapType} snap calling keyring_listAccountAssets`, { - snapId: config.snapId, - accountId, - }); + const client = this.#getKeyringClient(snapId); - const accountAssets = await this.#snapProvider.request({ - method: 'wallet_invokeSnap', - params: { - snapId: config.snapId, - request: { - method: 'keyring_listAccountAssets', - params: { - id: accountId, // Account UUID - }, - }, - }, - }); - - log(`${snapType} snap keyring_listAccountAssets response`, { - accountId, - assetCount: accountAssets?.length ?? 0, - assets: accountAssets, - }); + // Step 1: Get the list of assets for this account + const accountAssets = await client.listAccountAssets(accountId); // If no assets, skip to next account if (!accountAssets || accountAssets.length === 0) { - log( - `${snapType} snap: account has no assets, skipping balance fetch`, - { - accountId, - }, - ); continue; } // Step 2: Get balances for those specific assets - log(`${snapType} snap calling keyring_getAccountBalances`, { - snapId: config.snapId, - accountId, - requestedAssets: accountAssets.length, - }); - - const balances = await this.#snapProvider.request< - Record - >({ - method: 'wallet_invokeSnap', - params: { - snapId: config.snapId, - request: { - method: 'keyring_getAccountBalances', - params: { - id: accountId, // Account UUID (the keyring API uses 'id' not 'accountId') - assets: accountAssets, // Must pass specific asset types from listAccountAssets - }, - }, - }, - }); - - log(`${snapType} snap keyring_getAccountBalances response`, { - accountId, - balances, - balancesType: typeof balances, - isNull: balances === null, - isUndefined: balances === undefined, - assetCount: balances ? Object.keys(balances).length : 0, - }); + const balances: Record = + await client.getAccountBalances( + accountId, + accountAssets as CaipAssetType[], + ); // Transform keyring response to DataResponse format - // Note: snap may return null/undefined if account doesn't belong to this snap if (balances && typeof balances === 'object' && results.assetsBalance) { - const balanceEntries = Object.entries(balances); - log( - `${snapType} snap processing ${balanceEntries.length} balances for account ${accountId}`, - ); - - for (const [assetId, balance] of balanceEntries) { - // Initialize account balances if not exists + for (const [assetId, balance] of Object.entries(balances)) { results.assetsBalance[accountId] ??= {}; - // Store raw balance for this asset - // Use rawAmount if available (preferred - smallest unit), fall back to amount - // Note: Snaps should return rawAmount in smallest unit (satoshis, lamports, etc.) const accountBalances = results.assetsBalance[accountId]; if (accountBalances) { (accountBalances as Record)[assetId] = { - amount: balance.rawAmount ?? balance.amount, + amount: balance.amount, }; } } - } else if (!balances) { - log( - `${snapType} snap returned empty/null for account (account may not belong to this snap)`, - { - accountId, - balances, - }, - ); } - } catch (error) { - // This is expected when querying a snap with an account it doesn't manage - log(`${snapType} snap fetch FAILED for account`, { - accountId, - error: error instanceof Error ? error.message : String(error), - errorStack: error instanceof Error ? error.stack : undefined, - }); + } catch { + // Expected when account doesn't belong to this snap } } - log(`${snapType} snap fetch completed`, { - chains: request.chainIds.length, - accountsWithBalances: Object.keys(results.assetsBalance ?? {}).length, - }); - return results; } @@ -1016,21 +665,23 @@ export class SnapDataSource extends AbstractDataSource< return; } - // Filter to only snap-supported chains - const snapChains = request.chainIds.filter(isSnapSupportedChain); + // Filter to chains we have a snap for + const supportedChains = request.chainIds.filter((chainId) => + this.#isChainSupportedBySnap(chainId), + ); - if (snapChains.length === 0) { + if (supportedChains.length === 0) { return; } if (isUpdate) { const existing = this.activeSubscriptions.get(subscriptionId); if (existing) { - existing.chains = snapChains; + existing.chains = supportedChains; // Do a fetch to get latest data on subscription update this.fetch({ ...request, - chainIds: snapChains, + chainIds: supportedChains, }) .then(async (fetchResponse) => { if (Object.keys(fetchResponse.assetsBalance ?? {}).length > 0) { @@ -1059,14 +710,14 @@ export class SnapDataSource extends AbstractDataSource< cleanup: () => { // No timer to clear - we use event-based updates }, - chains: snapChains, + chains: supportedChains, }); // Initial fetch to get current balances try { const fetchResponse = await this.fetch({ ...request, - chainIds: snapChains, + chainIds: supportedChains, }); if (Object.keys(fetchResponse.assetsBalance ?? {}).length > 0) { @@ -1081,11 +732,46 @@ export class SnapDataSource extends AbstractDataSource< } } + // ============================================================================ + // KEYRING CLIENT + // ============================================================================ + + /** + * Gets a `KeyringClient` for a Snap. + * + * @param snapId - ID of the Snap to get the client for. + * @returns A `KeyringClient` for the Snap. + */ + #getKeyringClient(snapId: string): KeyringClient { + return new KeyringClient({ + send: async (request: JsonRpcRequest) => + (await this.#messenger.call('SnapController:handleRequest', { + snapId: snapId as SnapId, + origin: 'metamask', + handler: HandlerType.OnKeyringRequest, + request, + })) as Promise, + }); + } + // ============================================================================ // CLEANUP // ============================================================================ destroy(): void { + // Unsubscribe from snap keyring events + try { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const messenger = this.#messenger as any; + messenger.unsubscribe( + 'AccountsController:accountBalancesUpdated', + this.#handleSnapBalancesUpdatedBound, + ); + } catch (error) { + log('Failed to unsubscribe from snap keyring events', { error }); + } + + // Clean up active subscriptions for (const [subscriptionId] of this.activeSubscriptions) { this.unsubscribe(subscriptionId).catch(() => { // Ignore cleanup errors diff --git a/packages/assets-controller/src/data-sources/index.ts b/packages/assets-controller/src/data-sources/index.ts index 09aa5309d13..08cfc7ca4e3 100644 --- a/packages/assets-controller/src/data-sources/index.ts +++ b/packages/assets-controller/src/data-sources/index.ts @@ -64,51 +64,19 @@ export { type PriceDataSourceAssetsUpdatedEvent, } from './PriceDataSource'; -// Unified Snap Data Source (handles Solana, Bitcoin, Tron snaps) +// Unified Snap Data Source (dynamically discovers keyring snaps via PermissionController) export { SnapDataSource, createSnapDataSource, SNAP_DATA_SOURCE_NAME, - // Snap IDs - SOLANA_SNAP_ID, - BITCOIN_SNAP_ID, - TRON_SNAP_ID, - // Chain prefixes - SOLANA_CHAIN_PREFIX, - BITCOIN_CHAIN_PREFIX, - TRON_CHAIN_PREFIX, - // Networks - SOLANA_MAINNET, - SOLANA_DEVNET, - SOLANA_TESTNET, - BITCOIN_MAINNET, - BITCOIN_TESTNET, - TRON_MAINNET, - TRON_SHASTA, - TRON_NILE, - TRON_MAINNET_HEX, - TRON_SHASTA_HEX, - TRON_NILE_HEX, - ALL_DEFAULT_NETWORKS, - // Poll intervals - DEFAULT_SOLANA_POLL_INTERVAL, - DEFAULT_BITCOIN_POLL_INTERVAL, - DEFAULT_TRON_POLL_INTERVAL, - DEFAULT_SNAP_POLL_INTERVAL, - // Snap registry - SNAP_REGISTRY, - // Helper functions - getSnapTypeForChain, - isSnapSupportedChain, - isSolanaChain, - isBitcoinChain, - isTronChain, + // Constants + KEYRING_PERMISSION, + // Utility functions + getChainIdsCaveat, + extractChainFromAssetId, // Types - type SnapType, - type SnapInfo, type SnapDataSourceState, type SnapDataSourceOptions, - type SnapProvider, type SnapDataSourceActions, type SnapDataSourceEvents, type SnapDataSourceMessenger, diff --git a/packages/assets-controller/src/data-sources/initDataSources.test.ts b/packages/assets-controller/src/data-sources/initDataSources.test.ts index a171af7afb7..acd7e8d3048 100644 --- a/packages/assets-controller/src/data-sources/initDataSources.test.ts +++ b/packages/assets-controller/src/data-sources/initDataSources.test.ts @@ -8,7 +8,6 @@ import type { DataSourceMessengers } from './initDataSources'; import { PriceDataSource } from './PriceDataSource'; import { RpcDataSource } from './RpcDataSource'; import { SnapDataSource } from './SnapDataSource'; -import type { SnapProvider } from './SnapDataSource'; import { TokenDataSource } from './TokenDataSource'; import { DetectionMiddleware } from '../middlewares'; @@ -44,12 +43,6 @@ const MockDetectionMiddleware = DetectionMiddleware as jest.MockedClass< typeof DetectionMiddleware >; -function createMockSnapProvider(): SnapProvider { - return { - request: jest.fn(), - } as unknown as SnapProvider; -} - function createMockQueryApiClient(): ApiPlatformClient { return { fetch: jest.fn(), @@ -243,12 +236,10 @@ describe('initDataSources', () => { it('creates all data source instances', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); const dataSources = initDataSources({ messengers, - snapProvider, queryApiClient, }); @@ -263,12 +254,10 @@ describe('initDataSources', () => { it('creates RpcDataSource with correct messenger', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); initDataSources({ messengers, - snapProvider, queryApiClient, }); @@ -279,12 +268,10 @@ describe('initDataSources', () => { it('creates BackendWebsocketDataSource with correct messenger', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); initDataSources({ messengers, - snapProvider, queryApiClient, }); @@ -295,12 +282,10 @@ describe('initDataSources', () => { it('creates AccountsApiDataSource with correct options', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); initDataSources({ messengers, - snapProvider, queryApiClient, }); @@ -312,29 +297,24 @@ describe('initDataSources', () => { it('creates SnapDataSource with correct options', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); initDataSources({ messengers, - snapProvider, queryApiClient, }); expect(MockSnapDataSource).toHaveBeenCalledWith({ messenger: messengers.snapMessenger, - snapProvider, }); }); it('creates TokenDataSource with correct options', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); initDataSources({ messengers, - snapProvider, queryApiClient, }); @@ -346,12 +326,10 @@ describe('initDataSources', () => { it('creates PriceDataSource with correct options', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); initDataSources({ messengers, - snapProvider, queryApiClient, }); @@ -363,12 +341,10 @@ describe('initDataSources', () => { it('creates DetectionMiddleware with correct messenger', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); initDataSources({ messengers, - snapProvider, queryApiClient, }); @@ -379,12 +355,10 @@ describe('initDataSources', () => { it('returns instances of correct types', () => { const messengers = createMockMessengers(); - const snapProvider = createMockSnapProvider(); const queryApiClient = createMockQueryApiClient(); const dataSources = initDataSources({ messengers, - snapProvider, queryApiClient, }); diff --git a/packages/assets-controller/src/data-sources/initDataSources.ts b/packages/assets-controller/src/data-sources/initDataSources.ts index 916a13509db..32caac277df 100644 --- a/packages/assets-controller/src/data-sources/initDataSources.ts +++ b/packages/assets-controller/src/data-sources/initDataSources.ts @@ -37,7 +37,6 @@ import type { SnapDataSourceActions, SnapDataSourceEvents, SnapDataSourceMessenger, - SnapProvider, } from './SnapDataSource'; import { TokenDataSource } from './TokenDataSource'; import type { @@ -168,9 +167,6 @@ export type InitDataSourcesOptions = { /** Messengers for each data source */ messengers: DataSourceMessengers; - /** Snap provider for communicating with snaps */ - snapProvider: SnapProvider; - /** ApiPlatformClient for cached API calls */ queryApiClient: ApiPlatformClient; @@ -274,6 +270,11 @@ export function initMessengers< actions: [ 'AssetsController:activeChainsUpdate', 'AssetsController:assetsUpdate', + // SnapController actions for direct snap communication + 'SnapController:getRunnableSnaps', + 'SnapController:handleRequest', + // PermissionController action for dynamic snap discovery + 'PermissionController:getPermissions', ], events: [ // Snap keyring balance updates - snaps emit this when balances change @@ -339,7 +340,7 @@ export function initMessengers< * // Then initialize data sources * const dataSources = initDataSources({ * messengers, - * snapProvider: snapController, + * queryApiClient, * }); * ``` * @@ -347,8 +348,7 @@ export function initMessengers< * @returns Object containing all data source instances */ export function initDataSources(options: InitDataSourcesOptions): DataSources { - const { messengers, snapProvider, queryApiClient, rpcDataSourceConfig } = - options; + const { messengers, queryApiClient, rpcDataSourceConfig } = options; // Initialize primary data sources (provide balance data) const rpcDataSource = new RpcDataSource({ @@ -367,7 +367,6 @@ export function initDataSources(options: InitDataSourcesOptions): DataSources { const snapDataSource = new SnapDataSource({ messenger: messengers.snapMessenger, - snapProvider, }); // Initialize middleware data sources (enrich responses) diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index 3a09a8308ae..78983f6b771 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -127,53 +127,21 @@ export type { RpcDataSourceMessenger, } from './data-sources'; -// Data sources - Unified Snap Data Source (handles Solana, Bitcoin, Tron) +// Data sources - Unified Snap Data Source (dynamically discovers keyring snaps) export { SnapDataSource, createSnapDataSource, SNAP_DATA_SOURCE_NAME, - // Snap IDs - SOLANA_SNAP_ID, - BITCOIN_SNAP_ID, - TRON_SNAP_ID, - // Chain prefixes - SOLANA_CHAIN_PREFIX, - BITCOIN_CHAIN_PREFIX, - TRON_CHAIN_PREFIX, - // Networks - SOLANA_MAINNET, - SOLANA_DEVNET, - SOLANA_TESTNET, - BITCOIN_MAINNET, - BITCOIN_TESTNET, - TRON_MAINNET, - TRON_SHASTA, - TRON_NILE, - TRON_MAINNET_HEX, - TRON_SHASTA_HEX, - TRON_NILE_HEX, - ALL_DEFAULT_NETWORKS, - // Poll intervals - DEFAULT_SOLANA_POLL_INTERVAL, - DEFAULT_BITCOIN_POLL_INTERVAL, - DEFAULT_TRON_POLL_INTERVAL, - DEFAULT_SNAP_POLL_INTERVAL, - // Snap registry - SNAP_REGISTRY, - // Helper functions - getSnapTypeForChain, - isSnapSupportedChain, - isSolanaChain, - isBitcoinChain, - isTronChain, + // Constants + KEYRING_PERMISSION, + // Utility functions + getChainIdsCaveat, + extractChainFromAssetId, } from './data-sources'; export type { - SnapType, - SnapInfo, SnapDataSourceState, SnapDataSourceOptions, - SnapProvider, SnapDataSourceActions, SnapDataSourceEvents, SnapDataSourceMessenger, diff --git a/packages/assets-controller/tsconfig.build.json b/packages/assets-controller/tsconfig.build.json index 61b8ef501fc..85e7933fb37 100644 --- a/packages/assets-controller/tsconfig.build.json +++ b/packages/assets-controller/tsconfig.build.json @@ -11,7 +11,8 @@ { "path": "../core-backend/tsconfig.build.json" }, { "path": "../keyring-controller/tsconfig.build.json" }, { "path": "../messenger/tsconfig.build.json" }, - { "path": "../network-enablement-controller/tsconfig.build.json" } + { "path": "../network-enablement-controller/tsconfig.build.json" }, + { "path": "../permission-controller/tsconfig.build.json" } ], "include": ["../../types", "./src"], "exclude": ["**/*.test.ts", "**/__fixtures__/"] diff --git a/yarn.lock b/yarn.lock index e035ad4a55d..dd56bb1bbf9 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2672,12 +2672,17 @@ __metadata: "@metamask/base-controller": "npm:^9.0.0" "@metamask/controller-utils": "npm:^11.18.0" "@metamask/core-backend": "npm:^5.0.0" + "@metamask/keyring-api": "npm:^21.0.0" "@metamask/keyring-controller": "npm:^25.1.0" "@metamask/keyring-internal-api": "npm:^9.0.0" + "@metamask/keyring-snap-client": "npm:^8.0.0" "@metamask/messenger": "npm:^0.3.0" "@metamask/network-controller": "npm:^29.0.0" "@metamask/network-enablement-controller": "npm:^4.1.0" + "@metamask/permission-controller": "npm:^12.2.0" "@metamask/polling-controller": "npm:^16.0.2" + "@metamask/snaps-controllers": "npm:^17.2.0" + "@metamask/snaps-utils": "npm:^11.7.0" "@metamask/utils": "npm:^11.9.0" "@ts-bridge/cli": "npm:^0.6.4" "@types/jest": "npm:^27.4.1"