diff --git a/packages/tron-wallet-snap/CHANGELOG.md b/packages/tron-wallet-snap/CHANGELOG.md index e0ec215c..7e3b3abe 100644 --- a/packages/tron-wallet-snap/CHANGELOG.md +++ b/packages/tron-wallet-snap/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add `signProofOfOwnershipBatch` for signing multiple proof-of-ownership messages in one request. ([#265](https://github.com/MetaMask/internal-snaps/pull/265)) + ### Changed - **BREAKING** Bump `@metamask/keyring-api` from `^23.7.0` to `^24.1.0` ([#214](https://github.com/MetaMask/internal-snaps/pull/214)) diff --git a/packages/tron-wallet-snap/snap.manifest.json b/packages/tron-wallet-snap/snap.manifest.json index feeaa72a..bc08157e 100644 --- a/packages/tron-wallet-snap/snap.manifest.json +++ b/packages/tron-wallet-snap/snap.manifest.json @@ -7,7 +7,7 @@ "url": "https://github.com/MetaMask/internal-snaps.git" }, "source": { - "shasum": "iF+9rnAdk21FSl4s0l1iE9+oTvqg8mAy987x6FDUTvo=", + "shasum": "I0slZAjOKc8RGlYFgPVrQ+Rd0DLgmEU7IKNz/DOZtsk=", "location": { "npm": { "filePath": "dist/bundle.js", diff --git a/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.test.ts b/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.test.ts index 2419adee..323e8b81 100644 --- a/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.test.ts +++ b/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.test.ts @@ -1,4 +1,4 @@ -import { FeeType } from '@metamask/keyring-api'; +import { FeeType, TrxAccountType } from '@metamask/keyring-api'; import type { JsonRpcRequest } from '@metamask/snaps-sdk'; import type { Infer } from '@metamask/superstruct'; import { BigNumber } from 'bignumber.js'; @@ -1092,7 +1092,9 @@ describe('ClientRequestHandler', () => { beforeEach(() => { mockAccountsService = { findById: jest.fn(), + findByIds: jest.fn(), deriveTronKeypair: jest.fn(), + deriveTronKeypairs: jest.fn(), } as unknown as jest.Mocked; mockAssetsService = {} as unknown as jest.Mocked; @@ -1327,7 +1329,9 @@ describe('ClientRequestHandler', () => { beforeEach(() => { mockAccountsService = { findById: jest.fn(), + findByIds: jest.fn(), deriveTronKeypair: jest.fn(), + deriveTronKeypairs: jest.fn(), } as unknown as jest.Mocked; mockTronWeb = { @@ -1453,6 +1457,186 @@ describe('ClientRequestHandler', () => { ), ).rejects.toThrow('does not match signing account address'); }); + + describe('signProofOfOwnershipBatch', () => { + const TEST_ACCOUNT_ID_2 = '123e4567-e89b-42d3-a456-426614174001'; + const TEST_ADDRESS_2 = 'TJRabPrwbZy45sbavfcjinPJC18kjpRTv8'; + + const buildBatchRequest = ( + items: { accountId: string; message: string }[], + ): JsonRpcRequest => ({ + jsonrpc: '2.0' as const, + id: '1', + method: ClientRequestMethod.SignProofOfOwnershipBatch, + params: { items }, + }); + + const account1: TronKeyringAccount = { + id: TEST_ACCOUNT_ID, + address: TEST_ADDRESS, + entropySource: 'test-entropy', + derivationPath: "m/44'/195'/0'/0/0", + index: 0, + type: TrxAccountType.Eoa, + scopes: [Network.Mainnet], + options: {}, + methods: ['signMessage', 'signTransaction'], + }; + const account2: TronKeyringAccount = { + id: TEST_ACCOUNT_ID_2, + address: TEST_ADDRESS_2, + entropySource: 'test-entropy', + derivationPath: "m/44'/195'/0'/0/1", + index: 1, + type: TrxAccountType.Eoa, + scopes: [Network.Mainnet], + options: {}, + methods: ['signMessage', 'signTransaction'], + }; + + it('signs a batch and returns signatures in input order', async () => { + const message1 = buildProofMessage(TEST_ADDRESS); + const message2 = buildProofMessage(TEST_ADDRESS_2); + mockAccountsService.findByIds.mockResolvedValue([account2, account1]); + mockAccountsService.deriveTronKeypairs.mockResolvedValue([ + { + privateKeyBytes: new Uint8Array(), + publicKeyBytes: new Uint8Array(), + privateKeyHex: 'private-key-1', + address: TEST_ADDRESS, + }, + { + privateKeyBytes: new Uint8Array(), + publicKeyBytes: new Uint8Array(), + privateKeyHex: 'private-key-2', + address: TEST_ADDRESS_2, + }, + ]); + mockTronWeb.trx.signMessageV2 + .mockReturnValueOnce('0xsignature1') + .mockReturnValueOnce('0xsignature2'); + + const result = await clientRequestHandler.handle( + buildBatchRequest([ + { accountId: TEST_ACCOUNT_ID, message: message1 }, + { accountId: TEST_ACCOUNT_ID_2, message: message2 }, + ]), + ); + + expect(mockAccountsService.findByIds).toHaveBeenCalledWith([ + TEST_ACCOUNT_ID, + TEST_ACCOUNT_ID_2, + ]); + expect(mockAccountsService.deriveTronKeypairs).toHaveBeenCalledWith([ + account1, + account2, + ]); + expect(mockTronWeb.trx.signMessageV2).toHaveBeenNthCalledWith( + 1, + message1, + 'private-key-1', + ); + expect(mockTronWeb.trx.signMessageV2).toHaveBeenNthCalledWith( + 2, + message2, + 'private-key-2', + ); + expect(result).toStrictEqual({ + results: [ + { accountId: TEST_ACCOUNT_ID, signature: '0xsignature1' }, + { accountId: TEST_ACCOUNT_ID_2, signature: '0xsignature2' }, + ], + }); + }); + + it('returns item-level errors for missing accounts and address mismatches', async () => { + const missingAccountId = '123e4567-e89b-42d3-a456-426614174099'; + const validMessage = buildProofMessage(TEST_ADDRESS); + const mismatchedMessage = buildProofMessage(TEST_ADDRESS_2); + mockAccountsService.findByIds.mockResolvedValue([account1]); + mockAccountsService.deriveTronKeypairs.mockResolvedValue([ + { + privateKeyBytes: new Uint8Array(), + publicKeyBytes: new Uint8Array(), + privateKeyHex: 'private-key-1', + address: TEST_ADDRESS, + }, + ]); + mockTronWeb.trx.signMessageV2.mockReturnValue('0xsignature1'); + + const result = await clientRequestHandler.handle( + buildBatchRequest([ + { accountId: TEST_ACCOUNT_ID, message: validMessage }, + { accountId: missingAccountId, message: validMessage }, + { accountId: TEST_ACCOUNT_ID, message: mismatchedMessage }, + ]), + ); + + expect(mockAccountsService.deriveTronKeypairs).toHaveBeenCalledTimes(1); + expect(result).toStrictEqual({ + results: [ + { accountId: TEST_ACCOUNT_ID, signature: '0xsignature1' }, + { + accountId: missingAccountId, + error: `Account not found: ${missingAccountId}`, + }, + { + accountId: TEST_ACCOUNT_ID, + error: `Address in proof-of-ownership message (${TEST_ADDRESS_2}) does not match signing account address (${TEST_ADDRESS})`, + }, + ], + }); + }); + + it('returns item-level errors from batch key derivation', async () => { + const message = buildProofMessage(TEST_ADDRESS); + mockAccountsService.findByIds.mockResolvedValue([account1]); + mockAccountsService.deriveTronKeypairs.mockResolvedValue([ + { error: 'Unable to derive private key' }, + ]); + + const result = await clientRequestHandler.handle( + buildBatchRequest([{ accountId: TEST_ACCOUNT_ID, message }]), + ); + + expect(result).toStrictEqual({ + results: [ + { + accountId: TEST_ACCOUNT_ID, + error: 'Unable to derive private key', + }, + ], + }); + }); + + it('returns an item-level error when the derived address does not match the account', async () => { + const message = buildProofMessage(TEST_ADDRESS); + mockAccountsService.findByIds.mockResolvedValue([account1]); + mockAccountsService.deriveTronKeypairs.mockResolvedValue([ + { + privateKeyBytes: new Uint8Array(), + publicKeyBytes: new Uint8Array(), + privateKeyHex: 'private-key-1', + address: TEST_ADDRESS_2, + }, + ]); + + const result = await clientRequestHandler.handle( + buildBatchRequest([{ accountId: TEST_ACCOUNT_ID, message }]), + ); + + expect(mockTronWebFactory.createClient).not.toHaveBeenCalled(); + expect(mockTronWeb.trx.signMessageV2).not.toHaveBeenCalled(); + expect(result).toStrictEqual({ + results: [ + { + accountId: TEST_ACCOUNT_ID, + error: `Derived address (${TEST_ADDRESS_2}) does not match signing account address (${TEST_ADDRESS})`, + }, + ], + }); + }); + }); }); }); diff --git a/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.ts b/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.ts index 775ed93a..6d959dc3 100644 --- a/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.ts +++ b/packages/tron-wallet-snap/src/handlers/clientRequest/clientRequest.ts @@ -1,4 +1,5 @@ import { TransactionStatus } from '@metamask/keyring-api'; +import { normalizeError } from '@metamask/snap-networks-utils'; import type { Logger } from '@metamask/snap-networks-utils'; import type { Json, JsonRpcRequest } from '@metamask/snaps-sdk'; import { @@ -63,9 +64,12 @@ import { parseProofOfOwnershipMessage, parseRewardsMessage, SignAndSendTransactionRequestStruct, + SignProofOfOwnershipBatchRequestStruct, + SignProofOfOwnershipBatchResponseStruct, SignProofOfOwnershipRequestStruct, SignRewardsMessageRequestStruct, } from './validation'; +import type { SignProofOfOwnershipBatchResponse } from './validation'; type TransactionRawData = TronwebTypes.Transaction['raw_data'] & { // eslint-disable-next-line @typescript-eslint/naming-convention @@ -193,6 +197,8 @@ export class ClientRequestHandler { */ case ClientRequestMethod.SignProofOfOwnership: return this.#handleSignProofOfOwnership(request); + case ClientRequestMethod.SignProofOfOwnershipBatch: + return this.#handleSignProofOfOwnershipBatch(request); default: throw new MethodNotFoundError() as Error; } @@ -1175,6 +1181,136 @@ export class ClientRequestHandler { return { signature }; } + /** + * Handles silent batch signing of proof-of-ownership messages. + * + * Valid items are signed together so key derivation can be grouped by entropy + * source. Invalid items return per-item errors instead of failing the whole + * batch. + * + * @param request - The JSON-RPC request containing the batch items. + * @returns The response to the JSON-RPC request. + */ + async #handleSignProofOfOwnershipBatch( + request: JsonRpcRequest, + ): Promise { + assertOrThrow( + request, + SignProofOfOwnershipBatchRequestStruct, + new InvalidParamsError(), + ); + + const { + params: { items }, + } = request; + const uniqueAccountIds = [ + ...new Set(items.map(({ accountId }) => accountId)), + ]; + const accounts = await this.#accountsService.findByIds(uniqueAccountIds); + const accountsById = new Map( + accounts.map((account) => [account.id, account]), + ); + const results: SignProofOfOwnershipBatchResponse['results'] = new Array( + items.length, + ); + const signingRequests: { + index: number; + accountId: string; + account: (typeof accounts)[number]; + message: string; + }[] = []; + + items.forEach(({ accountId, message }, index) => { + const account = accountsById.get(accountId); + if (!account) { + results[index] = { + accountId, + error: `Account not found: ${accountId}`, + }; + return; + } + + try { + const { address: messageAddress } = + parseProofOfOwnershipMessage(message); + + if (messageAddress !== account.address) { + results[index] = { + accountId, + error: `Address in proof-of-ownership message (${messageAddress}) does not match signing account address (${account.address})`, + }; + return; + } + + signingRequests.push({ + index, + accountId, + account, + message, + }); + } catch (parseError) { + results[index] = { + accountId, + error: normalizeError(parseError).message, + }; + } + }); + + const derivedKeypairs = await this.#accountsService.deriveTronKeypairs( + signingRequests.map(({ account }) => account), + ); + + derivedKeypairs.forEach((derivedKeypair, signingRequestIndex) => { + const { index, accountId, account, message } = signingRequests[ + signingRequestIndex + ] as (typeof signingRequests)[number]; + const { error } = derivedKeypair as { error?: string }; + + if (error !== undefined) { + results[index] = { accountId, error }; + return; + } + + try { + const { address, privateKeyHex } = derivedKeypair as { + address: string; + privateKeyHex: string; + }; + + if (address !== account.address) { + results[index] = { + accountId, + error: `Derived address (${address}) does not match signing account address (${account.address})`, + }; + return; + } + + const tronWeb = this.#tronWebFactory.createClient( + Network.Mainnet, + privateKeyHex, + ); + const signature = tronWeb.trx.signMessageV2(message, privateKeyHex); + + results[index] = { accountId, signature }; + } catch (signError) { + results[index] = { + accountId, + error: normalizeError(signError).message, + }; + } + }); + + const result: SignProofOfOwnershipBatchResponse = { results }; + + assertOrThrow( + result, + SignProofOfOwnershipBatchResponseStruct, + new InvalidParamsError(), + ); + + return result; + } + /** * Sets the fee limit on a transaction's raw data and re-serializes it to hexadecimal format. * diff --git a/packages/tron-wallet-snap/src/handlers/clientRequest/types.ts b/packages/tron-wallet-snap/src/handlers/clientRequest/types.ts index ad8523f7..2dc96159 100644 --- a/packages/tron-wallet-snap/src/handlers/clientRequest/types.ts +++ b/packages/tron-wallet-snap/src/handlers/clientRequest/types.ts @@ -25,6 +25,11 @@ export const ClientRequestMethod = { * Sign Proof of Ownership */ SignProofOfOwnership: 'signProofOfOwnership', + /** + * Sign multiple proof-of-ownership messages for MetaMask identity + * authentication. + */ + SignProofOfOwnershipBatch: 'signProofOfOwnershipBatch', } as const; export type ClientRequestMethod = diff --git a/packages/tron-wallet-snap/src/handlers/clientRequest/validation.ts b/packages/tron-wallet-snap/src/handlers/clientRequest/validation.ts index db872772..c650a507 100644 --- a/packages/tron-wallet-snap/src/handlers/clientRequest/validation.ts +++ b/packages/tron-wallet-snap/src/handlers/clientRequest/validation.ts @@ -1,5 +1,13 @@ import { AssetStruct, FeeType } from '@metamask/keyring-api'; -import { UuidStruct } from '@metamask/snap-networks-utils'; +import { + parseProofOfOwnershipMessage as parseSharedProofOfOwnershipMessage, + ProofOfOwnershipBatchErrorStruct, + ProofOfOwnershipBatchRequestItemStruct, + ProofOfOwnershipBatchRequestParamsStruct, + ProofOfOwnershipBatchResponseStruct, + UuidStruct, +} from '@metamask/snap-networks-utils'; +import type { ProofOfOwnershipMessage } from '@metamask/snap-networks-utils'; import { literal } from '@metamask/snaps-sdk'; import type { Infer } from '@metamask/superstruct'; import { @@ -12,6 +20,7 @@ import { optional, refine, string, + union, } from '@metamask/superstruct'; import { CaipAssetTypeStruct, @@ -349,8 +358,6 @@ export const SignRewardsMessageRequestStruct = object({ params: SignRewardsMessageRequestParamsStruct, }); -export const PROOF_OF_OWNERSHIP_MESSAGE_PREFIX = 'metamask:proof-of-ownership:'; - /** * Parses a plaintext proof-of-ownership message. * Expected format: 'metamask:proof-of-ownership:{nonce}:{address}' @@ -359,38 +366,17 @@ export const PROOF_OF_OWNERSHIP_MESSAGE_PREFIX = 'metamask:proof-of-ownership:'; * @returns The parsed nonce and address. * @throws Error if the message format is invalid. */ -export function parseProofOfOwnershipMessage(message: string): { - nonce: string; - address: string; -} { - if (!message.startsWith(PROOF_OF_OWNERSHIP_MESSAGE_PREFIX)) { - throw new Error( - `Message must start with "${PROOF_OF_OWNERSHIP_MESSAGE_PREFIX}"`, - ); - } - - const remainder = message.slice(PROOF_OF_OWNERSHIP_MESSAGE_PREFIX.length); - const separatorIdx = remainder.lastIndexOf(':'); - if (separatorIdx === -1) { - throw new Error( - 'Message must follow the format "metamask:proof-of-ownership:{nonce}:{address}"', - ); - } - - const nonce = remainder.slice(0, separatorIdx); - const address = remainder.slice(separatorIdx + 1); - - if (nonce === '') { - throw new Error( - 'Proof-of-ownership message must contain a non-empty nonce', - ); - } +export function parseProofOfOwnershipMessage( + message: string, +): ProofOfOwnershipMessage { + const proofMessage = parseSharedProofOfOwnershipMessage(message); + const { address } = proofMessage; if (!is(address, TronAddressStruct)) { throw new Error('Invalid Tron address in proof-of-ownership message'); } - return { nonce, address }; + return proofMessage; } /** @@ -426,3 +412,63 @@ export const SignProofOfOwnershipRequestStruct = object({ method: literal(ClientRequestMethod.SignProofOfOwnership), params: SignProofOfOwnershipRequestParamsStruct, }); + +/** + * Validates one proof-of-ownership batch request item. + * + * Batch items intentionally validate messages as plain strings so invalid + * proof messages can be reported per item instead of failing the whole batch. + */ +export const SignProofOfOwnershipBatchRequestItemStruct = + ProofOfOwnershipBatchRequestItemStruct; + +/** + * Validates the params object for `signProofOfOwnershipBatch`. + */ +export const SignProofOfOwnershipBatchRequestParamsStruct = + ProofOfOwnershipBatchRequestParamsStruct; + +/** + * Validates a `signProofOfOwnershipBatch` JSON-RPC request. + */ +export const SignProofOfOwnershipBatchRequestStruct = object({ + jsonrpc: JsonRpcVersionStruct, + id: JsonRpcIdStruct, + method: literal(ClientRequestMethod.SignProofOfOwnershipBatch), + params: SignProofOfOwnershipBatchRequestParamsStruct, +}); + +/** + * Validates a successful proof-of-ownership batch item response. + */ +export const SignProofOfOwnershipBatchSuccessStruct = object({ + accountId: string(), + signature: string(), +}); + +/** + * Validates a failed proof-of-ownership batch item response. + */ +export const SignProofOfOwnershipBatchErrorStruct = + ProofOfOwnershipBatchErrorStruct; + +/** + * Validates a proof-of-ownership batch item result. + */ +export const SignProofOfOwnershipBatchItemResponseStruct = union([ + SignProofOfOwnershipBatchSuccessStruct, + SignProofOfOwnershipBatchErrorStruct, +]); + +/** + * Validates a `signProofOfOwnershipBatch` response. + */ +export const SignProofOfOwnershipBatchResponseStruct = + ProofOfOwnershipBatchResponseStruct; + +/** + * Response returned by `signProofOfOwnershipBatch`. + */ +export type SignProofOfOwnershipBatchResponse = Infer< + typeof SignProofOfOwnershipBatchResponseStruct +>; diff --git a/packages/tron-wallet-snap/src/permissions.ts b/packages/tron-wallet-snap/src/permissions.ts index 3dc1a15c..4fff5553 100644 --- a/packages/tron-wallet-snap/src/permissions.ts +++ b/packages/tron-wallet-snap/src/permissions.ts @@ -6,6 +6,7 @@ import { DEFAULT_PROD_ORIGINS, } from '@metamask/snap-networks-utils'; +import { ClientRequestMethod } from './handlers/clientRequest/types'; import { TestDappRpcRequestMethod } from './handlers/rpc/types'; // eslint-disable-next-line no-restricted-globals @@ -51,6 +52,8 @@ const metamaskMethods = [ KeyringRpcMethod.DiscoverAccounts, KeyringRpcMethod.ListAccountTransactions, KeyringRpcMethod.ListAccountAssets, + // Client methods + ClientRequestMethod.SignProofOfOwnershipBatch, ]; export const originPermissions = createOriginPermissions({ diff --git a/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts b/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts index 31879c32..54030d60 100644 --- a/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts +++ b/packages/tron-wallet-snap/src/services/accounts/AccountsRepository.ts @@ -138,9 +138,17 @@ export class AccountsRepository { return accounts.find((account) => account.id === id) ?? null; } + /** + * Finds multiple Tron keyring accounts with a single full account-state read. + * + * @param ids - Account IDs to resolve. + * @returns The matching accounts. Result ordering follows stored account + * ordering, not input ordering. + */ async findByIds(ids: string[]): Promise { + const idSet = new Set(ids); const accounts = await this.getAll(); - return accounts.filter((account) => ids.includes(account.id)); + return accounts.filter((account) => idSet.has(account.id)); } async findByAddress(address: string): Promise { diff --git a/packages/tron-wallet-snap/src/services/accounts/AccountsService.test.ts b/packages/tron-wallet-snap/src/services/accounts/AccountsService.test.ts index 3c708404..395f09de 100644 --- a/packages/tron-wallet-snap/src/services/accounts/AccountsService.test.ts +++ b/packages/tron-wallet-snap/src/services/accounts/AccountsService.test.ts @@ -335,6 +335,67 @@ describe('AccountsService', () => { }); }); + describe('deriveTronKeypairs', () => { + const createAccount = (index: number): TronKeyringAccount => + ({ + id: `account-${index}`, + entropySource: 'test-entropy', + derivationPath: AccountsService.getDefaultDerivationPath(index), + index, + address: `TAccount${index}`, + type: TrxAccountType.Eoa, + scopes: SUPPORTED_SCOPES as unknown as Network[], + options: {}, + methods: ['signMessage', 'signTransaction'], + }) as unknown as TronKeyringAccount; + + it('derives multiple keypairs with one entropy fetch per entropy source', async () => { + const coinJson = await getTronTestCoinTypeJson(); + + await withAccountsService(async ({ accountsService, mockSnapClient }) => { + const result = await accountsService.deriveTronKeypairs([ + createAccount(0), + createAccount(1), + ]); + + expect(mockSnapClient.getBip32Entropy).toHaveBeenCalledTimes(1); + expect(mockSnapClient.getBip32Entropy).toHaveBeenCalledWith({ + entropySource: 'test-entropy', + path: ['m', "44'", "195'"], + curve: 'secp256k1', + }); + expect(result).toHaveLength(2); + expect(result[0]).toMatchObject({ + privateKeyHex: expect.any(String), + address: expect.any(String), + }); + expect(result[1]).toMatchObject({ + privateKeyHex: expect.any(String), + address: expect.any(String), + }); + }, coinJson); + }); + + it('returns an item-level error for unsupported derivation paths', async () => { + const coinJson = await getTronTestCoinTypeJson(); + + await withAccountsService(async ({ accountsService }) => { + const result = await accountsService.deriveTronKeypairs([ + { + ...createAccount(0), + derivationPath: "m/44'/195'/0'", + }, + ]); + + expect(result).toStrictEqual([ + { + error: "Unsupported Tron derivation path: m/44'/195'/0'", + }, + ]); + }, coinJson); + }); + }); + describe('createAccounts', () => { it('persists new accounts with a single merge and one coin-type entropy call', async () => { const coinJson = await getTronTestCoinTypeJson(); diff --git a/packages/tron-wallet-snap/src/services/accounts/AccountsService.ts b/packages/tron-wallet-snap/src/services/accounts/AccountsService.ts index 6c8a0766..d5a0c04a 100644 --- a/packages/tron-wallet-snap/src/services/accounts/AccountsService.ts +++ b/packages/tron-wallet-snap/src/services/accounts/AccountsService.ts @@ -11,7 +11,10 @@ import { } from '@metamask/keyring-api'; import { getSelectedAccounts } from '@metamask/keyring-snap-sdk'; import type { Logger } from '@metamask/snap-networks-utils'; -import { InFlightCoalescer } from '@metamask/snap-networks-utils'; +import { + InFlightCoalescer, + normalizeError, +} from '@metamask/snap-networks-utils'; import { assert } from '@metamask/superstruct'; import { hexToBytes } from '@metamask/utils'; import { computeAddress } from 'ethers'; @@ -22,7 +25,10 @@ import type { SnapClient } from '../../clients/snap/SnapClient'; import { Network } from '../../constants'; import { asStrictKeyringAccount } from '../../entities/keyring-account'; import type { TronKeyringAccount } from '../../entities/keyring-account'; -import { createTronBip44AddressDeriver } from '../../utils/deriveTronFromCoinTypeNode'; +import { + createTronBip44AddressDeriver, + createTronBip44KeypairDeriver, +} from '../../utils/deriveTronFromCoinTypeNode'; import { sanitizeSensitiveError } from '../../utils/errors'; import { DerivationPathStruct } from '../../validation/structs'; import type { AssetsService } from '../assets/AssetsService'; @@ -62,6 +68,60 @@ type TronAddressDeriver = Awaited< ReturnType >; +/** + * A function that derives a TRON keypair from a BIP44 account index. + */ +type TronKeypairDeriver = Awaited< + ReturnType +>; + +/** + * Key material derived for one TRON account. + */ +export type DerivedTronKeypair = { + privateKeyBytes: Uint8Array; + publicKeyBytes: Uint8Array; + privateKeyHex: string; + address: string; +}; + +/** + * Result for one account in a batch TRON keypair derivation. + */ +export type DerivedTronKeypairBatchResult = + | DerivedTronKeypair + | { error: string }; + +const DEFAULT_TRON_DERIVATION_PATH_REGEX = /^m\/44'\/195'\/0'\/0\/([0-9]+)$/u; + +/** + * Extracts the address index from the default TRON BIP-44 derivation path. + * + * Batch derivation starts at the coin-type node (`m/44'/195'`), so it only + * supports the snap's default `m/44'/195'/0'/0/index` path shape. + * + * @param account - The TRON account whose derivation path should be parsed. + * @returns The BIP-44 address index. + */ +function getDefaultTronAddressIndex(account: TronKeyringAccount): number { + const match = DEFAULT_TRON_DERIVATION_PATH_REGEX.exec(account.derivationPath); + + if (!match?.[1]) { + throw new Error( + `Unsupported Tron derivation path: ${account.derivationPath}`, + ); + } + + const addressIndex = Number(match[1]); + if (!Number.isSafeInteger(addressIndex) || addressIndex !== account.index) { + throw new Error( + `Tron derivation path index (${addressIndex}) does not match account index (${account.index})`, + ); + } + + return addressIndex; +} + /** * Validates account creation ranges before any expensive state or entropy work. * @@ -149,12 +209,7 @@ export class AccountsService { }: { entropySource?: EntropySourceId | undefined; derivationPath: string; - }): Promise<{ - privateKeyBytes: Uint8Array; - publicKeyBytes: Uint8Array; - privateKeyHex: string; - address: string; - }> { + }): Promise { try { this.#logger.log({ derivationPath }, 'Generating TRON wallet'); @@ -196,6 +251,59 @@ export class AccountsService { } } + /** + * Derives keypairs for multiple TRON accounts with one coin-type entropy + * fetch per entropy source. + * + * Results are returned in input order. Individual account derivation failures + * are returned as item-level errors so callers can preserve partial success. + * + * @param accounts - The accounts to derive key material for. + * @returns One derivation result per account, in input order. + */ + async deriveTronKeypairs( + accounts: TronKeyringAccount[], + ): Promise { + const results: DerivedTronKeypairBatchResult[] = new Array(accounts.length); + const accountsByEntropySource = new Map< + EntropySourceId, + { index: number; account: TronKeyringAccount }[] + >(); + + accounts.forEach((account, index) => { + const sourceAccounts = + accountsByEntropySource.get(account.entropySource) ?? []; + sourceAccounts.push({ index, account }); + accountsByEntropySource.set(account.entropySource, sourceAccounts); + }); + + await Promise.all( + [...accountsByEntropySource.entries()].map( + async ([entropySource, sourceAccounts]) => { + try { + const keypairDeriver = + await this.#createTronKeypairDeriver(entropySource); + + for (const { index, account } of sourceAccounts) { + try { + const addressIndex = getDefaultTronAddressIndex(account); + results[index] = await keypairDeriver(addressIndex); + } catch (error) { + results[index] = { error: normalizeError(error).message }; + } + } + } catch (error) { + for (const { index } of sourceAccounts) { + results[index] = { error: normalizeError(error).message }; + } + } + }, + ), + ); + + return results; + } + /** * Batch-creates Tron accounts for a BIP-44 index or index range. Existing accounts for the * same entropy source and index are returned without duplicate state writes. @@ -395,6 +503,15 @@ export class AccountsService { return account; } + /** + * Finds multiple TRON keyring accounts. + * + * Missing accounts are logged but not thrown so callers can decide whether + * partial results are acceptable. + * + * @param ids - Account IDs to resolve. + * @returns The matching accounts. + */ async findByIds(ids: string[]): Promise { const accounts = await this.#accountsRepository.findByIds(ids); @@ -493,6 +610,24 @@ export class AccountsService { return createTronBip44AddressDeriver(bip44Node); } + /** + * Creates a TRON keypair deriver from the coin-type node. + * + * @param entropySource - Entropy source used to fetch the coin-type node. + * @returns A deriver for `m/44'/195'/0'/0/index` keypairs. + */ + async #createTronKeypairDeriver( + entropySource: EntropySourceId, + ): Promise { + const bip44Node = (await this.#snapClient.getBip32Entropy({ + entropySource, + path: ['m', "44'", "195'"], + curve: CURVE, + })) as JsonBIP44Node; + + return createTronBip44KeypairDeriver(bip44Node); + } + static getDefaultDerivationPath(index: number): `m/${string}` { return `m/44'/195'/0'/0/${index}`; } diff --git a/packages/tron-wallet-snap/src/utils/deriveTronFromCoinTypeNode.ts b/packages/tron-wallet-snap/src/utils/deriveTronFromCoinTypeNode.ts index 76456336..c6825a08 100644 --- a/packages/tron-wallet-snap/src/utils/deriveTronFromCoinTypeNode.ts +++ b/packages/tron-wallet-snap/src/utils/deriveTronFromCoinTypeNode.ts @@ -8,6 +8,10 @@ import { sanitizeSensitiveError } from './errors'; const DEFAULT_TRON_CHANGE_PATH = [`bip32:0'`, 'bip32:0'] as const; +type TronBip44ChangeNode = Awaited< + ReturnType>['derive']> +>; + /** * Builds a one-segment BIP-32 path for deriving from the cached change node. * @@ -41,6 +45,54 @@ function tronAddressFromPublicKeyHex(publicKey: string): { return { address, publicKeyBytes }; } +/** + * Maps a derived BIP-32 address node to Tron keypair material. + * + * @param addressNode - Node derived at `m/44'/195'/0'/0/i`. + * @param addressNode.privateKey - Private key for the derived address node. + * @param addressNode.publicKey - Public key for the derived address node. + * @returns The Tron keypair material used for signing. + */ +function tronKeypairFromAddressNode(addressNode: { + privateKey?: string; + publicKey?: string; +}): { + address: string; + privateKeyBytes: Uint8Array; + publicKeyBytes: Uint8Array; + privateKeyHex: string; +} { + if (!addressNode.privateKey || !addressNode.publicKey) { + throw new Error('Unable to derive private key'); + } + + const { address, publicKeyBytes } = tronAddressFromPublicKeyHex( + addressNode.publicKey, + ); + const privateKeyBytes = hexToBytes(addressNode.privateKey); + const privateKeyHex = addressNode.privateKey.slice(2); + + return { + address, + privateKeyBytes, + publicKeyBytes, + privateKeyHex, + }; +} + +/** + * Creates the cached Tron change node at `m/44'/195'/0'/0`. + * + * @param coinTypeNodeJson - JSON node from `snap_getBip32Entropy` at path `m/44'/195'`. + * @returns The derived change node used by Tron account/address derivers. + */ +async function createTronBip44ChangeNode( + coinTypeNodeJson: JsonBIP44Node, +): Promise { + const coinTypeNode = await BIP44Node.fromJSON(coinTypeNodeJson); + return coinTypeNode.derive(DEFAULT_TRON_CHANGE_PATH); +} + /** * Builds a reusable deriver for Tron addresses under `m/44'/195'/0'/0/i` from the * coin-type JSON at `m/44'/195'`, caching `0'/0` so each call only derives the @@ -58,8 +110,7 @@ export async function createTronBip44AddressDeriver( }> > { try { - const coinTypeNode = await BIP44Node.fromJSON(coinTypeNodeJson); - const changeNode = await coinTypeNode.derive(DEFAULT_TRON_CHANGE_PATH); + const changeNode = await createTronBip44ChangeNode(coinTypeNodeJson); return async (addressIndex: number) => { try { @@ -80,3 +131,40 @@ export async function createTronBip44AddressDeriver( throw sanitizeSensitiveError(error); } } + +/** + * Builds a reusable deriver for Tron keypairs under `m/44'/195'/0'/0/i` from + * the coin-type JSON at `m/44'/195'`, caching `0'/0` so each call only derives + * the final address index. + * + * @param coinTypeNodeJson - JSON node from `snap_getBip32Entropy` at path `m/44'/195'`. + * @returns A function that derives the Tron keypair for a given BIP-44 `address_index`. + */ +export async function createTronBip44KeypairDeriver( + coinTypeNodeJson: JsonBIP44Node, +): Promise< + (addressIndex: number) => Promise<{ + address: string; + privateKeyBytes: Uint8Array; + publicKeyBytes: Uint8Array; + privateKeyHex: string; + }> +> { + try { + const changeNode = await createTronBip44ChangeNode(coinTypeNodeJson); + + return async (addressIndex: number) => { + try { + const addressNode = await changeNode.derive( + getAddressIndexPath(addressIndex), + ); + + return tronKeypairFromAddressNode(addressNode); + } catch (error) { + throw sanitizeSensitiveError(error); + } + }; + } catch (error) { + throw sanitizeSensitiveError(error); + } +}