solana: add refund token to transfer authority PDA (#167)

Co-authored-by: A5 Pickle <[email protected]>

Author: a5-pickle

solana/programs/token-router/src/processor/market_order/prepare.rs

@@ -29,6 +29,7 @@ pub struct PrepareMarketOrder<'info> {
             TRANSFER_AUTHORITY_SEED_PREFIX,
             prepared_order.key().as_ref(),
             &args.hash().0,
+            refund_token.key().as_ref()
         ],
         bump,
         constraint = {
@@ -175,6 +176,14 @@ pub fn prepare_market_order(
         redeemer_message,
     } = args;
 
+    let token_program = &ctx.accounts.token_program;
+    let sender_token = &ctx.accounts.sender_token;
+    let custody_token = &ctx.accounts.prepared_custody_token;
+    let refund_token = &ctx.accounts.refund_token;
+
+    let prepared_order = &mut ctx.accounts.prepared_order;
+    let prepared_order_key = prepared_order.key();
+
     // Finally transfer amount to custody token account. We perform exclusive or because we do not
     // want to allow specifying more than one authority.
     let order_sender = match (
@@ -184,10 +193,10 @@ pub fn prepare_market_order(
         (Some(sender), None) => {
             token::transfer(
                 CpiContext::new(
-                    ctx.accounts.token_program.to_account_info(),
+                    token_program.to_account_info(),
                     token::Transfer {
-                        from: ctx.accounts.sender_token.to_account_info(),
-                        to: ctx.accounts.prepared_custody_token.to_account_info(),
+                        from: sender_token.to_account_info(),
+                        to: custody_token.to_account_info(),
                         authority: sender.to_account_info(),
                     },
                 ),
@@ -197,20 +206,19 @@ pub fn prepare_market_order(
             sender.key()
         }
         (None, Some(program_transfer_authority)) => {
-            let sender_token = &ctx.accounts.sender_token;
-
             token::transfer(
                 CpiContext::new_with_signer(
-                    ctx.accounts.token_program.to_account_info(),
+                    token_program.to_account_info(),
                     token::Transfer {
                         from: sender_token.to_account_info(),
-                        to: ctx.accounts.prepared_custody_token.to_account_info(),
+                        to: custody_token.to_account_info(),
                         authority: program_transfer_authority.to_account_info(),
                     },
                     &[&[
                         TRANSFER_AUTHORITY_SEED_PREFIX,
-                        ctx.accounts.prepared_order.key().as_ref(),
+                        prepared_order_key.as_ref(),
                         &hashed_args.0,
+                        refund_token.key().as_ref(),
                         &[ctx.bumps.program_transfer_authority.unwrap()],
                     ]],
                 ),
@@ -223,13 +231,13 @@ pub fn prepare_market_order(
     };
 
     // Set the values in prepared order account.
-    ctx.accounts.prepared_order.set_inner(PreparedOrder {
+    prepared_order.set_inner(PreparedOrder {
         info: PreparedOrderInfo {
             order_sender,
             prepared_by: ctx.accounts.payer.key(),
             order_type: OrderType::Market { min_amount_out },
-            src_token: ctx.accounts.sender_token.key(),
-            refund_token: ctx.accounts.refund_token.key(),
+            src_token: sender_token.key(),
+            refund_token: refund_token.key(),
             target_chain,
             redeemer,
             prepared_custody_token_bump: ctx.bumps.prepared_custody_token,

solana/ts/src/tokenRouter/index.ts

@@ -190,7 +190,11 @@ export class TokenRouterProgram {
         return this.program.account.preparedFill.fetch(addr);
     }
 
-    transferAuthorityAddress(preparedOrder: PublicKey, args: PrepareMarketOrderArgs): PublicKey {
+    transferAuthorityAddress(
+        preparedOrder: PublicKey,
+        args: PrepareMarketOrderArgs,
+        refundToken: PublicKey,
+    ): PublicKey {
         const { amountIn, minAmountOut, targetChain, redeemer, redeemerMessage } = args;
         const hasher = new Keccak(256);
         hasher.update(uint64ToBN(amountIn).toBuffer("be", 8));
@@ -208,7 +212,12 @@ export class TokenRouterProgram {
         hasher.update(redeemerMessage);
 
         return PublicKey.findProgramAddressSync(
-            [Buffer.from("transfer-authority"), preparedOrder.toBuffer(), hasher.digest()],
+            [
+                Buffer.from("transfer-authority"),
+                preparedOrder.toBuffer(),
+                hasher.digest(),
+                refundToken.toBuffer(),
+            ],
             this.ID,
         )[0];
     }
@@ -322,11 +331,12 @@ export class TokenRouterProgram {
         accounts: {
             preparedOrder: PublicKey;
             senderToken: PublicKey;
+            refundToken: PublicKey;
             senderTokenAuthority?: PublicKey;
         },
         args: PrepareMarketOrderArgs,
     ): Promise<{ transferAuthority: PublicKey; ix: TransactionInstruction }> {
-        const { preparedOrder, senderToken } = accounts;
+        const { preparedOrder, senderToken, refundToken } = accounts;
         const { amountIn } = args;
 
         let { senderTokenAuthority } = accounts;
@@ -338,7 +348,7 @@ export class TokenRouterProgram {
             return tokenAccount.owner;
         })();
 
-        const transferAuthority = this.transferAuthorityAddress(preparedOrder, args);
+        const transferAuthority = this.transferAuthorityAddress(preparedOrder, args, refundToken);
 
         return {
             transferAuthority,
@@ -380,7 +390,7 @@ export class TokenRouterProgram {
         if (programTransferAuthority === undefined) {
             if (useTransferAuthority) {
                 const approveResult = await this.approveTransferAuthorityIx(
-                    { preparedOrder, senderToken, senderTokenAuthority },
+                    { preparedOrder, senderToken, refundToken, senderTokenAuthority },
                     args,
                 );
                 programTransferAuthority = approveResult.transferAuthority;