feat(tron-wallet-snap): add batch proof-of-ownership signing

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))

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",

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<AccountsService>;
 
       mockAssetsService = {} as unknown as jest.Mocked<AssetsService>;
@@ -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<AccountsService>;
 
       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})`,
+            },
+          ],
+        });
+      });
+    });
   });
 });
 

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<Json> {
+    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.
    *