fix: separate refund claim when there is no token account

programs/svm-spoke/src/error.rs

@@ -53,8 +53,6 @@ pub enum CustomError {
     InvalidRemoteSender,
     #[msg("Invalid Merkle proof!")]
     InvalidProof,
-    #[msg("Account not found!")]
-    AccountNotFound,
     #[msg("Fills are currently paused!")]
     FillsArePaused,
     #[msg("Invalid chain id!")]
@@ -75,4 +73,8 @@ pub enum CustomError {
     InvalidFillDeadline,
     #[msg("Overflow writing to parameters account!")]
     ParamsWriteOverflow,
+    #[msg("Invalid refund address!")]
+    InvalidRefund,
+    #[msg("Zero relayer refund claim!")]
+    ZeroRefundClaim,
 }

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

@@ -5,7 +5,7 @@ use crate::{
     constants::DISCRIMINATOR_SIZE,
     error::CustomError,
     event::ExecutedRelayerRefundRoot,
-    state::{ExecuteRelayerRefundLeafParams, RootBundle, State, TransferLiability},
+    state::{ExecuteRelayerRefundLeafParams, RefundAccount, RootBundle, State, TransferLiability},
     utils::{is_claimed, set_claimed, verify_merkle_proof},
 };
 
@@ -106,9 +106,12 @@ impl RelayerRefundLeaf {
     }
 }
 
-pub fn execute_relayer_refund_leaf<'info>(
-    ctx: Context<'_, '_, '_, 'info, ExecuteRelayerRefundLeaf<'info>>,
-) -> Result<()> {
+pub fn execute_relayer_refund_leaf<'c, 'info>(
+    ctx: Context<'_, '_, 'c, 'info, ExecuteRelayerRefundLeaf<'info>>,
+) -> Result<()>
+where
+    'c: 'info,
+{
     // Get pre-loaded instruction parameters.
     let instruction_params = &ctx.accounts.instruction_params;
     let root_bundle_id = instruction_params.root_bundle_id;
@@ -147,29 +150,41 @@ pub fn execute_relayer_refund_leaf<'info>(
     let signer_seeds = &[&seeds[..]];
 
     for (i, amount) in relayer_refund_leaf.refund_amounts.iter().enumerate() {
-        let refund_account = relayer_refund_leaf.refund_accounts[i];
         let amount = *amount as u64;
 
-        // TODO: we might be able to just use the refund_account and improve this block but it's not clear yet if that's possible.
-        let refund_account_info = ctx
-            .remaining_accounts
-            .iter()
-            .find(|account| account.key == &refund_account)
-            .cloned()
-            .ok_or(CustomError::AccountNotFound)?;
-
-        let transfer_accounts = TransferChecked {
-            from: ctx.accounts.vault.to_account_info(),
-            mint: ctx.accounts.mint.to_account_info(),
-            to: refund_account_info.to_account_info(),
-            authority: ctx.accounts.state.to_account_info(),
-        };
-        let cpi_context = CpiContext::new_with_signer(
-            ctx.accounts.token_program.to_account_info(),
-            transfer_accounts,
-            signer_seeds,
-        );
-        transfer_checked(cpi_context, amount, ctx.accounts.mint.decimals)?;
+        // Refund account holds either a regular token account or a claim account. This checks all required constraints.
+        let refund_account = RefundAccount::try_from_remaining_account(
+            ctx.remaining_accounts,
+            i,
+            &relayer_refund_leaf.refund_accounts[i],
+            &ctx.accounts.mint.key(),
+            &ctx.accounts.token_program.key(),
+        )?;
+
+        match refund_account {
+            // Valid token account was passed, transfer the refund atomically.
+            RefundAccount::TokenAccount(token_account) => {
+                let transfer_accounts = TransferChecked {
+                    from: ctx.accounts.vault.to_account_info(),
+                    mint: ctx.accounts.mint.to_account_info(),
+                    to: token_account.to_account_info(),
+                    authority: ctx.accounts.state.to_account_info(),
+                };
+                let cpi_context = CpiContext::new_with_signer(
+                    ctx.accounts.token_program.to_account_info(),
+                    transfer_accounts,
+                    signer_seeds,
+                );
+                transfer_checked(cpi_context, amount, ctx.accounts.mint.decimals)?;
+            }
+            // Valid claim account was passed, increment the claim account amount.
+            RefundAccount::ClaimAccount(mut claim_account) => {
+                claim_account.amount += amount;
+
+                // Persist the updated claim account (Anchor handles this only for static accounts).
+                claim_account.exit(ctx.program_id)?;
+            }
+        }
     }
 
     if relayer_refund_leaf.amount_to_return > 0 {

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

@@ -4,6 +4,7 @@ mod deposit;
 mod fill;
 mod handle_receive_message;
 mod instruction_params;
+mod refund_claims;
 mod slow_fill;
 mod testable;
 mod token_bridge;
@@ -14,6 +15,7 @@ pub use deposit::*;
 pub use fill::*;
 pub use handle_receive_message::*;
 pub use instruction_params::*;
+pub use refund_claims::*;
 pub use slow_fill::*;
 pub use testable::*;
 pub use token_bridge::*;

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

@@ -0,0 +1,96 @@
+use anchor_lang::prelude::*;
+use anchor_spl::token_interface::{
+    transfer_checked, Mint, TokenAccount, TokenInterface, TransferChecked,
+};
+
+use crate::{
+    constants::DISCRIMINATOR_SIZE,
+    error::CustomError,
+    state::{ClaimAccount, State},
+};
+
+#[derive(Accounts)]
+#[instruction(mint: Pubkey, token_account: Pubkey)]
+pub struct InitializeClaimAccount<'info> {
+    #[account(mut)]
+    pub signer: Signer<'info>,
+
+    #[account(
+        init,
+        payer = signer,
+        space = DISCRIMINATOR_SIZE + ClaimAccount::INIT_SPACE,
+        seeds = [b"claim_account", mint.as_ref(), token_account.as_ref()],
+        bump
+    )]
+    pub claim_account: Account<'info, ClaimAccount>,
+
+    pub system_program: Program<'info, System>,
+}
+
+#[derive(Accounts)]
+pub struct ClaimRelayerRefund<'info> {
+    #[account(mut)]
+    pub signer: Signer<'info>,
+
+    #[account(mut, seeds = [b"state", state.seed.to_le_bytes().as_ref()], bump)]
+    pub state: Account<'info, State>,
+
+    #[account(
+        mut,
+        associated_token::mint = mint,
+        associated_token::authority = state,
+        associated_token::token_program = token_program
+    )]
+    pub vault: InterfaceAccount<'info, TokenAccount>,
+
+    // Mint address has been checked when executing the relayer refund leaf and it is part of claim account derivation.
+    #[account(
+        mint::token_program = token_program,
+    )]
+    pub mint: InterfaceAccount<'info, Mint>,
+
+    // Token address has been checked when executing the relayer refund leaf and it is part of claim account derivation.
+    #[account(
+        mut,
+        token::mint = mint,
+        token::token_program = token_program
+    )]
+    pub token_account: InterfaceAccount<'info, TokenAccount>,
+
+    #[account(
+        mut,
+        seeds = [b"claim_account", mint.key().as_ref(), token_account.key().as_ref()],
+        bump
+    )]
+    pub claim_account: Account<'info, ClaimAccount>,
+
+    pub token_program: Interface<'info, TokenInterface>,
+}
+
+pub fn claim_relayer_refund(ctx: Context<ClaimRelayerRefund>) -> Result<()> {
+    // Ensure the claim account holds a non-zero amount.
+    let claim_amount = ctx.accounts.claim_account.amount;
+    require!(claim_amount > 0, CustomError::ZeroRefundClaim);
+
+    // Reset the claim amount.
+    ctx.accounts.claim_account.amount = 0;
+
+    // Derive the signer seeds for the state required for the transfer form vault.
+    let state_seed_bytes = ctx.accounts.state.seed.to_le_bytes();
+    let seeds = &[b"state", state_seed_bytes.as_ref(), &[ctx.bumps.state]];
+    let signer_seeds = &[&seeds[..]];
+
+    // Transfer the claim amount from the vault to the relayer token account.
+    let transfer_accounts = TransferChecked {
+        from: ctx.accounts.vault.to_account_info(),
+        mint: ctx.accounts.mint.to_account_info(),
+        to: ctx.accounts.token_account.to_account_info(),
+        authority: ctx.accounts.state.to_account_info(),
+    };
+    let cpi_context = 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)
+}

programs/svm-spoke/src/lib.rs

@@ -62,9 +62,12 @@ pub mod svm_spoke {
         instructions::relay_root_bundle(ctx, relayer_refund_root, slow_relay_root)
     }
 
-    pub fn execute_relayer_refund_leaf<'info>(
-        ctx: Context<'_, '_, '_, 'info, ExecuteRelayerRefundLeaf<'info>>,
-    ) -> Result<()> {
+    pub fn execute_relayer_refund_leaf<'c, 'info>(
+        ctx: Context<'_, '_, 'c, 'info, ExecuteRelayerRefundLeaf<'info>>,
+    ) -> Result<()>
+    where
+        'c: 'info,
+    {
         instructions::execute_relayer_refund_leaf(ctx)
     }
 
@@ -206,4 +209,16 @@ pub mod svm_spoke {
     pub fn close_instruction_params(ctx: Context<CloseInstructionParams>) -> Result<()> {
         instructions::close_instruction_params(ctx)
     }
+
+    pub fn initialize_claim_account(
+        _ctx: Context<InitializeClaimAccount>,
+        mint: Pubkey,
+        token_account: Pubkey,
+    ) -> Result<()> {
+        Ok(())
+    }
+
+    pub fn claim_relayer_refund(ctx: Context<ClaimRelayerRefund>) -> Result<()> {
+        instructions::claim_relayer_refund(ctx)
+    }
 }

programs/svm-spoke/src/state/mod.rs

@@ -1,12 +1,14 @@
 pub mod fill;
 pub mod instruction_params;
+pub mod refund_account;
 pub mod root_bundle;
 pub mod route;
 pub mod state;
 pub mod transfer_liability;
 
 pub use fill::*;
 pub use instruction_params::*;
+pub use refund_account::*;
 pub use root_bundle::*;
 pub use route::*;
 pub use state::*;

programs/svm-spoke/src/state/refund_account.rs

@@ -0,0 +1,131 @@
+use anchor_lang::prelude::*;
+use anchor_spl::token_interface::TokenAccount;
+
+use crate::error::CustomError;
+
+#[account]
+#[derive(InitSpace)]
+pub struct ClaimAccount {
+    pub amount: u64,
+}
+
+// When executing relayer refund leaf, refund accounts are passed as remaining accounts and can hold either a regular
+// token account or a claim account. This enum is used to differentiate between the two types.
+pub enum RefundAccount<'info> {
+    TokenAccount(InterfaceAccount<'info, TokenAccount>),
+    ClaimAccount(Account<'info, ClaimAccount>),
+}
+
+impl<'c, 'info> RefundAccount<'info>
+where
+    'c: 'info,
+{
+    // This function is used to parse a refund account from the remaining accounts list. It first tries to parse it as
+    // a token account and if that fails, it tries to parse it as a claim account.
+    pub fn try_from_remaining_account(
+        remaining_accounts: &'c [AccountInfo<'info>],
+        index: usize,
+        expected_token_account: &Pubkey,
+        expected_mint: &Pubkey,
+        token_program: &Pubkey,
+    ) -> Result<Self> {
+        let refund_account_info = remaining_accounts
+            .get(index)
+            .ok_or(ErrorCode::AccountNotEnoughKeys)?;
+
+        Self::try_token_account_from_account_info(
+            refund_account_info,
+            expected_token_account,
+            expected_mint,
+            token_program,
+        )
+        .map(Self::TokenAccount)
+        .or_else(|| {
+            Self::try_claim_account_from_account_info(
+                refund_account_info,
+                expected_mint,
+                expected_token_account,
+            )
+            .map(Self::ClaimAccount)
+        })
+        .ok_or_else(|| {
+            error::Error::from(CustomError::InvalidRefund)
+                .with_account_name(&format!("remaining_accounts[{}]", index))
+        })
+    }
+
+    // This implements the following Anchor account constraints when parsing remaining account as a token account:
+    // #[account(
+    //     mut,
+    //     address = expected_token_account @ CustomError::InvalidRefund,
+    //     token::mint = expected_mint,
+    //     token::token_program = token_program
+    // )]
+    // pub token_account: InterfaceAccount<'info, TokenAccount>,
+    // Note: All errors are ignored and Option is returned as we do not log them anyway due to memory constraints.
+    fn try_token_account_from_account_info(
+        account_info: &'info AccountInfo<'info>,
+        expected_token_account: &Pubkey,
+        expected_mint: &Pubkey,
+        token_program: &Pubkey,
+    ) -> Option<InterfaceAccount<'info, TokenAccount>> {
+        // Checks ownership on deserialization for the TokenAccount interface.
+        let token_account: InterfaceAccount<'info, TokenAccount> =
+            InterfaceAccount::try_from(account_info).ok()?;
+
+        // Checks if the token account is writable.
+        if !account_info.is_writable {
+            return None;
+        }
+
+        // Checks the token address matches.
+        if account_info.key != expected_token_account {
+            return None;
+        }
+
+        // Checks if the token account is associated with the expected mint.
+        if &token_account.mint != expected_mint {
+            return None;
+        }
+
+        // Checks ownership by specific token program.
+        if account_info.owner != token_program {
+            return None;
+        }
+
+        Some(token_account)
+    }
+
+    // This implements the following Anchor account constraints when parsing remaining account as a claim account:
+    // #[account(
+    //     mut,
+    //     seeds = [b"claim_account", mint.key().as_ref(), token_account.key().as_ref()],
+    //     bump
+    // )]
+    // pub claim_account: Account<'info, ClaimAccount>,
+    // Note: All errors are ignored and Option is returned as we do not log them anyway due to memory constraints.
+    fn try_claim_account_from_account_info(
+        account_info: &'info AccountInfo<'info>,
+        mint: &Pubkey,
+        token_account: &Pubkey,
+    ) -> Option<Account<'info, ClaimAccount>> {
+        // Checks ownership on deserialization for the ClaimAccount.
+        let claim_account: Account<'info, ClaimAccount> = Account::try_from(account_info).ok()?;
+
+        // Checks the PDA is derived from mint and token account keys.
+        let (pda_address, _bump) = Pubkey::find_program_address(
+            &[b"claim_account", mint.as_ref(), token_account.as_ref()],
+            &crate::ID,
+        );
+        if account_info.key() != pda_address {
+            return None;
+        }
+
+        // Checks if the claim account is writable.
+        if !account_info.is_writable {
+            return None;
+        }
+
+        Some(claim_account)
+    }
+}