diff --git a/packages/snap/snap.manifest.json b/packages/snap/snap.manifest.json index 098d9f22..8a28bc15 100644 --- a/packages/snap/snap.manifest.json +++ b/packages/snap/snap.manifest.json @@ -7,7 +7,7 @@ "url": "https://github.com/MetaMask/snap-tron-wallet.git" }, "source": { - "shasum": "tSxhnKoJe+iTI1HJ13ENq81vPK99Q12rj7owcUCpqxQ=", + "shasum": "KWVrILlsFNeR5YoP7RLalzLse8MDPlTzrXn7tihjabA=", "location": { "npm": { "filePath": "dist/bundle.js", diff --git a/packages/snap/src/services/accounts/AccountsService.test.ts b/packages/snap/src/services/accounts/AccountsService.test.ts index 5ecc24f7..9c17e64c 100644 --- a/packages/snap/src/services/accounts/AccountsService.test.ts +++ b/packages/snap/src/services/accounts/AccountsService.test.ts @@ -30,6 +30,8 @@ import { mockLogger } from '../../utils/mockLogger'; import type { AssetsService } from '../assets/AssetsService'; import type { ConfigProvider } from '../config'; import type { Config } from '../config/ConfigProvider'; +import { MigrationStage } from '../migration/stage'; +import type { TronAssetsControllerAdapter } from '../migration/TronAssetsControllerAdapter'; import type { TransactionsService } from '../transactions/TransactionsService'; jest.mock('@metamask/keyring-snap-sdk', () => ({ @@ -120,7 +122,12 @@ type WithAccountsServiceCallback = (payload: { mockConfigProvider: jest.Mocked>; mockLogger: ILogger; mockAssetsService: jest.Mocked< - Pick + Pick< + AssetsService, + | 'fetchAssetsAndBalancesForAccount' + | 'saveMany' + | 'buildControllerSnapshot' + > >; mockSnapClient: jest.Mocked< Pick @@ -128,6 +135,9 @@ type WithAccountsServiceCallback = (payload: { mockTransactionsService: jest.Mocked< Pick >; + mockTronAssetsControllerAdapter: jest.Mocked< + Pick + >; }) => void | Promise; /** @@ -265,10 +275,16 @@ async function withAccountsService( }; const mockAssetsService: jest.Mocked< - Pick + Pick< + AssetsService, + | 'fetchAssetsAndBalancesForAccount' + | 'saveMany' + | 'buildControllerSnapshot' + > > = { fetchAssetsAndBalancesForAccount: jest.fn().mockResolvedValue([]), saveMany: jest.fn().mockResolvedValue(undefined), + buildControllerSnapshot: jest.fn().mockReturnValue([]), }; const mockTransactionsService: jest.Mocked< @@ -278,6 +294,13 @@ async function withAccountsService( saveMany: jest.fn().mockResolvedValue(undefined), }; + const mockTronAssetsControllerAdapter: jest.Mocked< + Pick + > = { + getMigrationStage: jest.fn().mockResolvedValue(MigrationStage.Off), + pushAssetSnapshot: jest.fn().mockResolvedValue(undefined), + }; + const accountsService = new AccountsService({ accountsRepository: mockAccountsRepository, configProvider: mockConfigProvider, @@ -285,6 +308,7 @@ async function withAccountsService( assetsService: mockAssetsService, snapClient: mockSnapClient, transactionsService: mockTransactionsService, + tronAssetsControllerAdapter: mockTronAssetsControllerAdapter, } as unknown as ConstructorParameters[0]); await testFn({ @@ -295,6 +319,7 @@ async function withAccountsService( mockAssetsService, mockSnapClient, mockTransactionsService, + mockTronAssetsControllerAdapter, }); } @@ -1147,6 +1172,153 @@ describe('AccountsService', () => { }, ); }); + + it('does not write controller snapshots at stage Off', async () => { + const account: TronKeyringAccount = { + id: 'off-id', + address: 'TOff123456789012345678901', + type: TrxAccountType.Eoa, + options: {}, + methods: [], + scopes: [], + entropySource: 'e1', + derivationPath: "m/44'/195'/0'/0/0", + index: 0, + }; + await withAccountsService( + async ({ + accountsService, + mockConfigProvider, + mockAssetsService, + mockTronAssetsControllerAdapter, + }) => { + mockConfigProvider.get.mockReturnValue({ + ...MOCK_CONFIG, + activeNetworks: [Network.Mainnet], + }); + mockAssetsService.fetchAssetsAndBalancesForAccount.mockResolvedValue( + [], + ); + + await accountsService.synchronizeAssets([account]); + + expect( + mockTronAssetsControllerAdapter.pushAssetSnapshot, + ).not.toHaveBeenCalled(); + }, + ); + }); + + it('writes snapshots at stage 1 and continues when a write fails', async () => { + const account: TronKeyringAccount = { + id: 'write-id', + address: 'TWrite12345678901234567890', + type: TrxAccountType.Eoa, + options: {}, + methods: [], + scopes: [], + entropySource: 'e1', + derivationPath: "m/44'/195'/0'/0/0", + index: 0, + }; + const snapshot = [ + [ + { + assetId: 'asset', + amount: '1', + metadata: { symbol: 'A', name: 'Asset', decimals: 0 }, + }, + ], + ]; + + await withAccountsService( + async ({ + accountsService, + mockConfigProvider, + mockAssetsService, + mockTronAssetsControllerAdapter, + }) => { + mockConfigProvider.get.mockReturnValue({ + ...MOCK_CONFIG, + activeNetworks: [Network.Mainnet], + }); + mockTronAssetsControllerAdapter.getMigrationStage.mockResolvedValue( + MigrationStage.ReadAssetsControllerWithFallback, + ); + mockAssetsService.buildControllerSnapshot.mockReturnValue(snapshot); + mockTronAssetsControllerAdapter.pushAssetSnapshot.mockRejectedValue( + new Error('controller unavailable'), + ); + + expect( + await accountsService.synchronizeAssets([account]), + ).toBeUndefined(); + + expect( + mockTronAssetsControllerAdapter.pushAssetSnapshot, + ).toHaveBeenCalledWith(account.id, Network.Mainnet, snapshot[0]); + }, + ); + }); + + it('does not persist controller-only assets and resumes after rollback', async () => { + const account: TronKeyringAccount = { + id: 'controller-only-id', + address: 'TControllerOnly1234567890', + type: TrxAccountType.Eoa, + options: {}, + methods: [], + scopes: [], + entropySource: 'e1', + derivationPath: "m/44'/195'/0'/0/0", + index: 0, + }; + + await withAccountsService( + async ({ + accountsService, + mockConfigProvider, + mockAssetsService, + mockTronAssetsControllerAdapter, + }) => { + const assets: NativeAsset[] = [ + { + assetType: `${Network.Mainnet}/slip44:195`, + keyringAccountId: account.id, + network: Network.Mainnet, + symbol: 'TRX', + decimals: 6, + rawAmount: '1000000', + uiAmount: '1', + iconUrl: '', + }, + ]; + mockConfigProvider.get.mockReturnValue({ + ...MOCK_CONFIG, + activeNetworks: [Network.Mainnet], + }); + mockAssetsService.fetchAssetsAndBalancesForAccount.mockResolvedValue( + assets, + ); + mockTronAssetsControllerAdapter.getMigrationStage.mockResolvedValueOnce( + MigrationStage.ReadAssetsControllerOnly, + ); + + await accountsService.synchronizeAssets([account]); + + expect(mockAssetsService.saveMany).not.toHaveBeenCalled(); + + mockAssetsService.saveMany.mockClear(); + mockTronAssetsControllerAdapter.getMigrationStage.mockResolvedValueOnce( + MigrationStage.Off, + ); + + await accountsService.synchronizeAssets([account]); + + expect(mockAssetsService.saveMany).toHaveBeenCalledWith(assets); + }, + ); + }); }); describe('synchronizeTransactions', () => { diff --git a/packages/snap/src/services/accounts/AccountsService.ts b/packages/snap/src/services/accounts/AccountsService.ts index d5d82acf..8cb4b404 100644 --- a/packages/snap/src/services/accounts/AccountsService.ts +++ b/packages/snap/src/services/accounts/AccountsService.ts @@ -35,6 +35,8 @@ import { createPrefixedLogger, type ILogger } from '../../utils/logger'; import { DerivationPathStruct } from '../../validation/structs'; import type { AssetsService } from '../assets/AssetsService'; import type { ConfigProvider } from '../config'; +import { MigrationStage, shouldPushSnapAssets } from '../migration/stage'; +import type { TronAssetsControllerAdapter } from '../migration/TronAssetsControllerAdapter'; import type { TransactionsService } from '../transactions/TransactionsService'; /** @@ -110,6 +112,8 @@ export class AccountsService { readonly #snapClient: SnapClient; + readonly #tronAssetsControllerAdapter: TronAssetsControllerAdapter; + constructor({ accountsRepository, configProvider, @@ -117,6 +121,7 @@ export class AccountsService { assetsService, snapClient, transactionsService, + tronAssetsControllerAdapter, }: { accountsRepository: AccountsRepository; configProvider: ConfigProvider; @@ -124,6 +129,7 @@ export class AccountsService { assetsService: AssetsService; snapClient: SnapClient; transactionsService: TransactionsService; + tronAssetsControllerAdapter: TronAssetsControllerAdapter; }) { this.#logger = createPrefixedLogger(logger, '[🔑 AccountsService]'); this.#configProvider = configProvider; @@ -131,6 +137,7 @@ export class AccountsService { this.#assetsService = assetsService; this.#transactionsService = transactionsService; this.#snapClient = snapClient; + this.#tronAssetsControllerAdapter = tronAssetsControllerAdapter; } /** @@ -509,11 +516,99 @@ export class AccountsService { }), ); - const assets = assetResponses.flatMap((response) => - response.status === 'fulfilled' ? response.value : [], + const synchronizedAssets = await Promise.all( + assetResponses.map(async (response, index) => { + if (response.status !== 'fulfilled') { + return { assets: [], shouldPersist: false, combination: undefined }; + } + + const combination = combinations[index]; + + if (!combination) { + return { assets: [], shouldPersist: false, combination: undefined }; + } + + const { scope } = combination; + const stage = + await this.#tronAssetsControllerAdapter.getMigrationStage(scope); + + if (stage >= MigrationStage.ReadAssetsControllerOnly) { + this.#logger.info( + 'Asset tracking and persistence disabled for controller-only stage', + { scope, stage }, + ); + return { + assets: response.value, + shouldPersist: false, + combination: { ...combination, stage }, + }; + } + + return { + assets: response.value, + shouldPersist: true, + combination: { ...combination, stage }, + }; + }), + ); + const assets = synchronizedAssets.flatMap( + ({ assets: responseAssets }) => responseAssets, ); + const hasPersistableScope = synchronizedAssets.some( + ({ shouldPersist }) => shouldPersist, + ); + + if (hasPersistableScope || combinations.length === 0) { + await this.#assetsService.saveMany(assets); + } - await this.#assetsService.saveMany(assets); + const writes = synchronizedAssets.flatMap( + ({ assets: scopeAssets, combination }) => { + if (!combination || !shouldPushSnapAssets(combination.stage)) { + return []; + } + + return this.#assetsService + .buildControllerSnapshot( + combination.account, + combination.scope, + scopeAssets, + ) + .map((group) => ({ ...combination, group })); + }, + ); + + if (writes.length > 0) { + const startedAt = Date.now(); + this.#logger.info('Assets controller snapshot write attempted', { + attemptCount: writes.length, + }); + const results = await Promise.allSettled( + writes.map(async ({ account, scope, group }) => + this.#tronAssetsControllerAdapter.pushAssetSnapshot( + account.id, + scope, + group, + ), + ), + ); + const failureCount = results.filter( + (result) => result.status === 'rejected', + ).length; + + this.#logger.info('Assets controller snapshot write completed', { + attemptCount: writes.length, + failureCount, + latencyMs: Date.now() - startedAt, + }); + results.forEach((result) => { + if (result.status === 'rejected') { + this.#logger.error('Assets controller snapshot write failed', { + error: result.reason, + }); + } + }); + } } async synchronizeTransactions(accounts: TronKeyringAccount[]): Promise { diff --git a/packages/snap/src/services/migration/TronAssetsControllerAdapter.test.ts b/packages/snap/src/services/migration/TronAssetsControllerAdapter.test.ts new file mode 100644 index 00000000..e6ee14cc --- /dev/null +++ b/packages/snap/src/services/migration/TronAssetsControllerAdapter.test.ts @@ -0,0 +1,94 @@ +import { MigrationStage } from './stage'; +import { TronAssetsControllerAdapter } from './TronAssetsControllerAdapter'; +import { WalletMessengerClient } from '../../clients/wallet/WalletMessengerClient'; +import { KnownCaip19Id, Network } from '../../constants'; + +describe('TronAssetsControllerAdapter', () => { + const resolveStage = jest.fn().mockResolvedValue(MigrationStage.Off); + + it('pushes the complete snapshot through the AssetsController action', async () => { + const messenger = { call: jest.fn().mockResolvedValue(undefined) }; + const adapter = new TronAssetsControllerAdapter( + new WalletMessengerClient(messenger), + resolveStage, + ); + const assets = [ + { + assetId: 'tron:728126428/slip44:energy', + amount: '42', + metadata: { symbol: 'ENERGY', name: 'Energy', decimals: 0 }, + }, + ]; + + await adapter.pushAssetSnapshot('account-id', 'tron:728126428', assets); + + expect(messenger.call).toHaveBeenCalledWith( + 'AssetsController:upsertSnapAssets', + 'account-id', + 'tron:728126428', + assets, + ); + }); + + it('maps controller asset data to AssetEntity', async () => { + const messenger = { + call: jest.fn().mockResolvedValue({ + amount: '1000000', + metadata: { + symbol: 'TRX', + name: 'Tron', + decimals: 6, + }, + }), + }; + const adapter = new TronAssetsControllerAdapter( + new WalletMessengerClient(messenger), + resolveStage, + ); + + const asset = await adapter.getAsset( + 'account-id', + KnownCaip19Id.TrxMainnet, + ); + + expect(messenger.call).toHaveBeenCalledWith( + 'AssetsController:getAsset', + 'account-id', + KnownCaip19Id.TrxMainnet, + ); + expect(asset).toMatchObject({ + assetType: KnownCaip19Id.TrxMainnet, + keyringAccountId: 'account-id', + network: Network.Mainnet, + symbol: 'TRX', + decimals: 6, + rawAmount: '1000000', + uiAmount: '1', + }); + }); + + it('returns null when controller has no asset', async () => { + const messenger = { call: jest.fn().mockResolvedValue(undefined) }; + const adapter = new TronAssetsControllerAdapter( + new WalletMessengerClient(messenger), + resolveStage, + ); + + expect( + await adapter.getAsset('account-id', KnownCaip19Id.TrxMainnet), + ).toBeNull(); + }); + + it('stores the most recently resolved stage', async () => { + const adapter = new TronAssetsControllerAdapter( + new WalletMessengerClient(undefined), + jest.fn().mockResolvedValue(MigrationStage.ReadAssetsControllerOnly), + ); + + await adapter.resolveAndSetStage(Network.Mainnet); + + expect(adapter.getCurrentStage()).toBe( + MigrationStage.ReadAssetsControllerOnly, + ); + }); +}); diff --git a/packages/snap/src/services/migration/stage.test.ts b/packages/snap/src/services/migration/stage.test.ts new file mode 100644 index 00000000..071a2607 --- /dev/null +++ b/packages/snap/src/services/migration/stage.test.ts @@ -0,0 +1,169 @@ +/* eslint-disable no-restricted-globals */ + +import { + MigrationStage, + resetMigrationStageEnvForTests, + resolveStage, + shouldPushSnapAssets, +} from './stage'; +import { WalletMessengerClient } from '../../clients/wallet/WalletMessengerClient'; + +const CHAIN_ID = 'tron:728126428'; + +describe('resolveStage', () => { + const originalEnvironment = process.env.ENVIRONMENT; + + beforeEach(() => { + process.env.ENVIRONMENT = 'local'; + resetMigrationStageEnvForTests(); + }); + + afterEach(() => { + process.env.ENVIRONMENT = originalEnvironment; + resetMigrationStageEnvForTests(); + }); + + it('returns Off by default when messenger is unavailable', async () => { + const client = new WalletMessengerClient(undefined); + + expect(await resolveStage(CHAIN_ID, client)).toBe(MigrationStage.Off); + }); + + it('returns ReadAssetsControllerWithFallback from remote flag', async () => { + const messenger = { + call: jest.fn().mockReturnValue({ + remoteFeatureFlags: { + snapsAssetsMigration: { + stages: { + [CHAIN_ID]: MigrationStage.ReadAssetsControllerWithFallback, + }, + }, + }, + }), + }; + const client = new WalletMessengerClient(messenger); + + expect(await resolveStage(CHAIN_ID, client)).toBe( + MigrationStage.ReadAssetsControllerWithFallback, + ); + }); + + it('returns ReadAssetsControllerWithoutFallback from remote flag', async () => { + const messenger = { + call: jest.fn().mockReturnValue({ + remoteFeatureFlags: { + snapsAssetsMigration: { + stages: { + [CHAIN_ID]: MigrationStage.ReadAssetsControllerWithoutFallback, + }, + }, + }, + }), + }; + const client = new WalletMessengerClient(messenger); + + expect(await resolveStage(CHAIN_ID, client)).toBe( + MigrationStage.ReadAssetsControllerWithoutFallback, + ); + }); + + it('returns ReadAssetsControllerOnly from remote flag', async () => { + const messenger = { + call: jest.fn().mockReturnValue({ + remoteFeatureFlags: { + snapsAssetsMigration: { + stages: { [CHAIN_ID]: MigrationStage.ReadAssetsControllerOnly }, + }, + }, + }), + }; + const client = new WalletMessengerClient(messenger); + + expect(await resolveStage(CHAIN_ID, client)).toBe( + MigrationStage.ReadAssetsControllerOnly, + ); + }); + + it('returns Off when killSwitch is true', async () => { + const messenger = { + call: jest.fn().mockReturnValue({ + remoteFeatureFlags: { + snapsAssetsMigration: { + killSwitch: true, + stages: { [CHAIN_ID]: MigrationStage.ReadAssetsControllerOnly }, + }, + }, + }), + }; + const client = new WalletMessengerClient(messenger); + + expect(await resolveStage(CHAIN_ID, client)).toBe(MigrationStage.Off); + }); + + it('returns Off when remote flag is missing', async () => { + const messenger = { + call: jest.fn().mockReturnValue({ + remoteFeatureFlags: {}, + }), + }; + const client = new WalletMessengerClient(messenger); + + expect(await resolveStage(CHAIN_ID, client)).toBe(MigrationStage.Off); + }); + + it('returns Off for malformed remote flag values', async () => { + const messenger = { + call: jest.fn().mockReturnValue({ + remoteFeatureFlags: { + snapsAssetsMigration: { + stages: { [CHAIN_ID]: 99 }, + }, + }, + }), + }; + const client = new WalletMessengerClient(messenger); + + expect(await resolveStage(CHAIN_ID, client)).toBe(MigrationStage.Off); + }); + + it('falls back to TRON_ASSETS_MIGRATION_STAGE in dev when messenger is unavailable', async () => { + process.env.TRON_ASSETS_MIGRATION_STAGE = + 'read-assets-controller-with-fallback'; + const client = new WalletMessengerClient(undefined); + + expect(await resolveStage(CHAIN_ID, client)).toBe( + MigrationStage.ReadAssetsControllerWithFallback, + ); + }); + + it('falls back to env when messenger read fails in dev', async () => { + process.env.TRON_ASSETS_MIGRATION_STAGE = '2'; + const messenger = { + call: jest.fn().mockImplementation(() => { + throw new Error('messenger unavailable'); + }), + }; + const client = new WalletMessengerClient(messenger); + + expect(await resolveStage(CHAIN_ID, client)).toBe( + MigrationStage.ReadAssetsControllerWithoutFallback, + ); + }); + + it('ignores env override in production', async () => { + process.env.ENVIRONMENT = 'production'; + process.env.TRON_ASSETS_MIGRATION_STAGE = '3'; + const client = new WalletMessengerClient(undefined); + + expect(await resolveStage(CHAIN_ID, client)).toBe(MigrationStage.Off); + }); +}); + +describe('shouldPushSnapAssets', () => { + it('returns false only when migration is off', () => { + expect(shouldPushSnapAssets(MigrationStage.Off)).toBe(false); + expect( + shouldPushSnapAssets(MigrationStage.ReadAssetsControllerWithFallback), + ).toBe(true); + }); +}); diff --git a/packages/snap/src/services/migration/stage.ts b/packages/snap/src/services/migration/stage.ts index ae51f545..d004c5a7 100644 --- a/packages/snap/src/services/migration/stage.ts +++ b/packages/snap/src/services/migration/stage.ts @@ -9,6 +9,19 @@ export enum MigrationStage { ReadAssetsControllerOnly = 3, } +/** + * Returns whether Snap asset snapshots should be pushed to AssetsController. + * + * Writes begin with the first controller migration stage so they can shadow + * the existing Snap store while reads still use their configured behavior. + * + * @param stage - The current assets migration stage. + * @returns Whether controller writes are enabled. + */ +export function shouldPushSnapAssets(stage: MigrationStage): boolean { + return stage >= MigrationStage.ReadAssetsControllerWithFallback; +} + export const SNAPS_ASSETS_MIGRATION_FLAG = 'snapsAssetsMigration'; const STAGE_ENV_VAR = 'TRON_ASSETS_MIGRATION_STAGE';