fix(svm): H-01 optimize claiming relayer refunds

programs/svm-spoke/src/instructions/refund_claims.rs

@@ -9,16 +9,21 @@ use crate::{
 };
 
 #[derive(Accounts)]
-#[instruction(mint: Pubkey, refund_address: Pubkey)]
 pub struct InitializeClaimAccount<'info> {
     #[account(mut)]
     pub signer: Signer<'info>,
 
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub mint: UncheckedAccount<'info>,
+
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub refund_address: UncheckedAccount<'info>,
+
     #[account(
         init,
         payer = signer,
         space = DISCRIMINATOR_SIZE + ClaimAccount::INIT_SPACE,
-        seeds = [b"claim_account", mint.as_ref(), refund_address.as_ref()],
+        seeds = [b"claim_account", mint.key().as_ref(), refund_address.key().as_ref()],
         bump
     )]
     pub claim_account: Account<'info, ClaimAccount>,
@@ -107,7 +112,6 @@ pub fn claim_relayer_refund(ctx: Context<ClaimRelayerRefund>) -> Result<()> {
 
 #[event_cpi]
 #[derive(Accounts)]
-#[instruction(refund_address: Pubkey)]
 pub struct ClaimRelayerRefundFor<'info> {
     pub signer: Signer<'info>,
 
@@ -130,6 +134,9 @@ pub struct ClaimRelayerRefundFor<'info> {
     #[account(mint::token_program = token_program)]
     pub mint: InterfaceAccount<'info, Mint>,
 
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub refund_address: UncheckedAccount<'info>,
+
     #[account(
         mut,
         associated_token::mint = mint,
@@ -141,15 +148,15 @@ pub struct ClaimRelayerRefundFor<'info> {
     #[account(
         mut,
         close = initializer,
-        seeds = [b"claim_account", mint.key().as_ref(), refund_address.as_ref()],
+        seeds = [b"claim_account", mint.key().as_ref(), refund_address.key().as_ref()],
         bump
     )]
     pub claim_account: Account<'info, ClaimAccount>,
 
     pub token_program: Interface<'info, TokenInterface>,
 }
 
-pub fn claim_relayer_refund_for(ctx: Context<ClaimRelayerRefundFor>, refund_address: Pubkey) -> Result<()> {
+pub fn claim_relayer_refund_for(ctx: Context<ClaimRelayerRefundFor>) -> Result<()> {
     // Ensure the claim account holds a non-zero amount.
     let claim_amount = ctx.accounts.claim_account.amount;
     if claim_amount == 0 {
@@ -172,7 +179,11 @@ pub fn claim_relayer_refund_for(ctx: Context<ClaimRelayerRefundFor>, refund_addr
         CpiContext::new_with_signer(ctx.accounts.token_program.to_account_info(), transfer_accounts, signer_seeds);
     transfer_checked(cpi_context, claim_amount, ctx.accounts.mint.decimals)?;
 
-    emit_cpi!(ClaimedRelayerRefund { l2_token_address: ctx.accounts.mint.key(), claim_amount, refund_address });
+    emit_cpi!(ClaimedRelayerRefund {
+        l2_token_address: ctx.accounts.mint.key(),
+        claim_amount,
+        refund_address: ctx.accounts.refund_address.key(),
+    });
 
     Ok(()) // There is no need to reset the claim amount as the account will be closed at the end of instruction.
 }
@@ -181,11 +192,16 @@ pub fn claim_relayer_refund_for(ctx: Context<ClaimRelayerRefundFor>, refund_addr
 // relayer refunds were executed with ATA after initializing the claim account. In such cases, the initializer should be
 // able to close the claim account manually.
 #[derive(Accounts)]
-#[instruction(mint: Pubkey, refund_address: Pubkey)]
 pub struct CloseClaimAccount<'info> {
     #[account(mut, address = claim_account.initializer @ SvmError::InvalidClaimInitializer)]
     pub signer: Signer<'info>,
 
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub mint: UncheckedAccount<'info>,
+
+    /// CHECK: This is only used for claim_account PDA derivation and it is up to the caller to ensure it is valid.
+    pub refund_address: UncheckedAccount<'info>,
+
     #[account(
         mut,
         close = signer,

programs/svm-spoke/src/lib.rs

@@ -463,8 +463,8 @@ pub mod svm_spoke {
     }
 
     /// Functionally identical to claim_relayer_refund() except the refund is sent to a specified refund address.
-    pub fn claim_relayer_refund_for(ctx: Context<ClaimRelayerRefundFor>, refund_address: Pubkey) -> Result<()> {
-        instructions::claim_relayer_refund_for(ctx, refund_address)
+    pub fn claim_relayer_refund_for(ctx: Context<ClaimRelayerRefundFor>) -> Result<()> {
+        instructions::claim_relayer_refund_for(ctx)
     }
 
     /// Creates token accounts in batch for a set of addresses.
@@ -632,18 +632,12 @@ pub mod svm_spoke {
     ///
     /// ### Required Accounts:
     /// - signer (Signer): The account that pays for the transaction and initializes the claim account.
+    /// - mint: The mint associated with the claim account.
+    /// - refund_address: The refund address associated with the claim account.
     /// - claim_account (Writable): The newly created claim account PDA to store claim data for this associated mint.
     ///   Seed: ["claim_account",mint,refund_address].
     /// - system_program: The system program required for account creation.
-    ///
-    /// ### Parameters:
-    /// - _mint: The public key of the mint associated with the claim account.
-    /// - _refund_address: The public key of the refund address associated with the claim account.
-    pub fn initialize_claim_account(
-        ctx: Context<InitializeClaimAccount>,
-        _mint: Pubkey,
-        _refund_address: Pubkey,
-    ) -> Result<()> {
+    pub fn initialize_claim_account(ctx: Context<InitializeClaimAccount>) -> Result<()> {
         instructions::initialize_claim_account(ctx)
     }
 
@@ -655,16 +649,10 @@ pub mod svm_spoke {
     ///
     /// ### Required Accounts:
     /// - signer (Signer): The account that authorizes the closure. Must be the initializer of the claim account.
+    /// - mint: The mint associated with the claim account.
+    /// - refund_address: The refund address associated with the claim account.
     /// - claim_account (Writable): The claim account PDA to be closed. Seed: ["claim_account",mint,refund_address].
-    ///
-    /// ### Parameters:
-    /// - _mint: The public key of the mint associated with the claim account.
-    /// - _refund_address: The public key of the refund address associated with the claim account.
-    pub fn close_claim_account(
-        ctx: Context<CloseClaimAccount>,
-        _mint: Pubkey,           // Only used in account constraints.
-        _refund_address: Pubkey, // Only used in account constraints.
-    ) -> Result<()> {
+    pub fn close_claim_account(ctx: Context<CloseClaimAccount>) -> Result<()> {
         instructions::close_claim_account(ctx)
     }
 

test/svm/SvmSpoke.Bundle.ts

@@ -7,7 +7,7 @@ import {
   mintTo,
   TOKEN_PROGRAM_ID,
 } from "@solana/spl-token";
-import { ComputeBudgetProgram, Keypair, PublicKey } from "@solana/web3.js";
+import { ComputeBudgetProgram, Keypair, PublicKey, TransactionInstruction } from "@solana/web3.js";
 import { assert } from "chai";
 import * as crypto from "crypto";
 import { ethers } from "ethers";
@@ -1015,7 +1015,7 @@ describe("svm_spoke.bundle", () => {
         } else if (!testConfig.deferredRefunds && testConfig.atomicAccountCreation) {
           refundAccounts.push(tokenAccount);
         } else {
-          await program.methods.initializeClaimAccount(mint, tokenOwner).rpc();
+          await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: tokenOwner }).rpc();
           refundAccounts.push(claimAccount);
         }
 
@@ -1397,8 +1397,8 @@ describe("svm_spoke.bundle", () => {
           [Buffer.from("claim_account"), mint.toBuffer(), relayerB.publicKey.toBuffer()],
           program.programId
         );
-        await program.methods.initializeClaimAccount(mint, relayerA.publicKey).rpc();
-        await program.methods.initializeClaimAccount(mint, relayerB.publicKey).rpc();
+        await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: relayerA.publicKey }).rpc();
+        await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: relayerB.publicKey }).rpc();
       }
 
       // Prepare leaf using token accounts.
@@ -1732,4 +1732,220 @@ describe("svm_spoke.bundle", () => {
       assert.include(err.toString(), "Invalid Merkle proof", "Expected merkle verification to fail");
     }
   });
+
+  describe("Execute Max multiple refunds with claims", async () => {
+    const executeMaxRefundClaims = async (testConfig: {
+      solanaDistributions: number;
+      useAddressLookup: boolean;
+      separatePhases: boolean;
+    }) => {
+      // Add leaves for other EVM chains to have non-empty proofs array to ensure we don't run out of memory when processing.
+      const evmDistributions = 100; // This would fit in 7 proof array elements.
+
+      const refundAddresses: web3.PublicKey[] = []; // These are relayer authority addresses used in leaf building.
+      const claimAccounts: web3.PublicKey[] = [];
+      const tokenAccounts: web3.PublicKey[] = [];
+      const refundAmounts: BN[] = [];
+      const initializeInstructions: TransactionInstruction[] = [];
+      const claimInstructions: TransactionInstruction[] = [];
+
+      for (let i = 0; i < testConfig.solanaDistributions; i++) {
+        // Create the token account.
+        const tokenOwner = Keypair.generate().publicKey;
+        const tokenAccount = (await getOrCreateAssociatedTokenAccount(connection, payer, mint, tokenOwner)).address;
+        refundAddresses.push(tokenOwner);
+        tokenAccounts.push(tokenAccount);
+
+        const [claimAccount] = PublicKey.findProgramAddressSync(
+          [Buffer.from("claim_account"), mint.toBuffer(), tokenOwner.toBuffer()],
+          program.programId
+        );
+
+        // Create instruction to initialize claim account.
+        initializeInstructions.push(
+          await program.methods.initializeClaimAccount().accounts({ mint, refundAddress: tokenOwner }).instruction()
+        );
+        claimAccounts.push(claimAccount);
+
+        refundAmounts.push(new BN(randomBigInt(2).toString()));
+
+        // Create instruction to claim refund to the token account.
+        const claimRelayerRefundAccounts = {
+          signer: owner,
+          initializer: owner,
+          state,
+          vault,
+          mint,
+          tokenAccount,
+          refundAddress: tokenOwner,
+          claimAccount,
+          tokenProgram: TOKEN_PROGRAM_ID,
+          program: program.programId,
+        };
+        claimInstructions.push(
+          await program.methods.claimRelayerRefundFor().accounts(claimRelayerRefundAccounts).instruction()
+        );
+      }
+
+      const { relayerRefundLeaves, merkleTree } = buildRelayerRefundMerkleTree({
+        totalEvmDistributions: evmDistributions,
+        totalSolanaDistributions: testConfig.solanaDistributions,
+        mixLeaves: false,
+        chainId: chainId.toNumber(),
+        mint,
+        svmRelayers: refundAddresses,
+        svmRefundAmounts: refundAmounts,
+      });
+
+      const root = merkleTree.getRoot();
+      const proof = merkleTree.getProof(relayerRefundLeaves[0]);
+      const leaf = relayerRefundLeaves[0] as RelayerRefundLeafSolana;
+
+      const stateAccountData = await program.account.state.fetch(state);
+      const rootBundleId = stateAccountData.rootBundleId;
+
+      const rootBundleIdBuffer = Buffer.alloc(4);
+      rootBundleIdBuffer.writeUInt32LE(rootBundleId);
+      const seeds = [Buffer.from("root_bundle"), seed.toArrayLike(Buffer, "le", 8), rootBundleIdBuffer];
+      const [rootBundle] = PublicKey.findProgramAddressSync(seeds, program.programId);
+
+      // Relay root bundle
+      const relayRootBundleAccounts = { state, rootBundle, signer: owner, payer: owner, program: program.programId };
+      await program.methods.relayRootBundle(Array.from(root), Array.from(root)).accounts(relayRootBundleAccounts).rpc();
+
+      // Verify valid leaf
+      const proofAsNumbers = proof.map((p) => Array.from(p));
+
+      const [instructionParams] = PublicKey.findProgramAddressSync(
+        [Buffer.from("instruction_params"), owner.toBuffer()],
+        program.programId
+      );
+
+      const executeAccounts = {
+        instructionParams,
+        state,
+        rootBundle: rootBundle,
+        signer: owner,
+        vault,
+        tokenProgram: TOKEN_PROGRAM_ID,
+        mint,
+        transferLiability,
+        program: program.programId,
+      };
+
+      const executeRemainingAccounts = claimAccounts.map((account) => ({
+        pubkey: account,
+        isWritable: true,
+        isSigner: false,
+      }));
+
+      // Build the instruction to execute relayer refund leaf and write its instruction args to the data account.
+      await loadExecuteRelayerRefundLeafParams(program, owner, stateAccountData.rootBundleId, leaf, proofAsNumbers);
+
+      const executeInstruction = await program.methods
+        .executeRelayerRefundLeafDeferred()
+        .accounts(executeAccounts)
+        .remainingAccounts(executeRemainingAccounts)
+        .instruction();
+
+      // Initialize, execute and claim depending on the chosen method.
+      const instructions = [...initializeInstructions, executeInstruction, ...claimInstructions];
+      if (!testConfig.separatePhases) {
+        // Pack all instructions in one transaction.
+        if (testConfig.useAddressLookup)
+          await sendTransactionWithLookupTable(
+            connection,
+            instructions,
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+        else
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(...instructions),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+      } else {
+        // Send claim account initialization, execution and claim in separate transactions.
+        if (testConfig.useAddressLookup) {
+          await sendTransactionWithLookupTable(
+            connection,
+            initializeInstructions,
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+          await sendTransactionWithLookupTable(
+            connection,
+            [executeInstruction],
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+          await sendTransactionWithLookupTable(
+            connection,
+            claimInstructions,
+            (anchor.AnchorProvider.env().wallet as anchor.Wallet).payer
+          );
+        } else {
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(...initializeInstructions),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(executeInstruction),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+          await web3.sendAndConfirmTransaction(
+            connection,
+            new web3.Transaction().add(...claimInstructions),
+            [(anchor.AnchorProvider.env().wallet as anchor.Wallet).payer],
+            {
+              commitment: "confirmed",
+            }
+          );
+        }
+      }
+
+      // Verify all refund token account balances.
+      const refundBalances = await Promise.all(
+        tokenAccounts.map(async (account) => {
+          return (await connection.getTokenAccountBalance(account)).value.amount;
+        })
+      );
+      refundBalances.forEach((balance, i) => {
+        assertSE(balance, refundAmounts[i].toString(), `Refund account ${i} balance should match refund amount`);
+      });
+    };
+
+    it("Execute Max multiple refunds with claims in one legacy transaction", async () => {
+      // Larger amount would hit transaction message size limit.
+      const solanaDistributions = 5;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: false, separatePhases: false });
+    });
+
+    it("Execute Max multiple refunds with claims in one versioned transaction", async () => {
+      // Larger amount would hit maximum instruction trace length limit.
+      const solanaDistributions = 12;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: true, separatePhases: false });
+    });
+
+    it("Execute Max multiple refunds with claims in separate phase legacy transactions", async () => {
+      // Larger amount would hit transaction message size limit.
+      const solanaDistributions = 7;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: false, separatePhases: true });
+    });
+
+    it("Execute Max multiple refunds with claims in separate phase versioned transactions", async () => {
+      // Larger amount would hit maximum instruction trace length limit.
+      const solanaDistributions = 21;
+      await executeMaxRefundClaims({ solanaDistributions, useAddressLookup: true, separatePhases: true });
+    });
+  });
 });