fix(transfer-switch): key the transfer switch on the source token account owner

tokens/token-2022/transfer-hook/transfer-switch/anchor/programs/transfer-switch/src/instructions/initialise_extra_account_metas_list.rs

@@ -4,9 +4,7 @@ use {
         system_program::{create_account, CreateAccount},
     },
     anchor_spl::token_interface::Mint,
-    spl_tlv_account_resolution::{
-        account::ExtraAccountMeta, seeds::Seed, state::ExtraAccountMetaList,
-    },
+    spl_tlv_account_resolution::{account::ExtraAccountMeta, seeds::Seed, state::ExtraAccountMetaList},
     spl_transfer_hook_interface::instruction::ExecuteInstruction,
 };
 
@@ -30,21 +28,20 @@ pub struct InitializeExtraAccountMetas<'info> {
 }
 
 impl<'info> InitializeExtraAccountMetas<'info> {
-    pub fn initialize_extra_account_metas_list(
-        &self,
-        bumps: InitializeExtraAccountMetasBumps,
-    ) -> Result<()> {
+    pub fn initialize_extra_account_metas_list(&self, bumps: InitializeExtraAccountMetasBumps) -> Result<()> {
         // .map_err() needed because spl-tlv-account-resolution uses solana-program-error 2.x
         // while anchor-lang 1.0 uses 3.x — structurally identical but different semver types
         let account_metas = vec![
-            // 5 - wallet (sender) config account
+            // 5 - sender (source token account owner) switch account
             ExtraAccountMeta::new_with_seeds(
                 &[
-                    Seed::AccountKey { index: 3 }, // sender index
+                    // owner field of the source token account; index 3 may be a delegate
+                    Seed::AccountData { account_index: 0, data_index: 32, length: 32 },
                 ],
                 false, // is_signer
                 false, // is_writable
-            ).map_err(|_| ProgramError::InvalidArgument)?,
+            )
+            .map_err(|_| ProgramError::InvalidArgument)?,
         ];
 
         // calculate account size
@@ -55,11 +52,7 @@ impl<'info> InitializeExtraAccountMetas<'info> {
         let lamports = Rent::get()?.minimum_balance(account_size as usize);
 
         let mint = self.token_mint.key();
-        let signer_seeds: &[&[&[u8]]] = &[&[
-            b"extra-account-metas",
-            mint.as_ref(),
-            &[bumps.extra_account_metas_list],
-        ]];
+        let signer_seeds: &[&[&[u8]]] = &[&[b"extra-account-metas", mint.as_ref(), &[bumps.extra_account_metas_list]]];
 
         create_account(
             CpiContext::new(
@@ -79,7 +72,8 @@ impl<'info> InitializeExtraAccountMetas<'info> {
         ExtraAccountMetaList::init::<ExecuteInstruction>(
             &mut self.extra_account_metas_list.try_borrow_mut_data()?,
             &account_metas,
-        ).map_err(|_| ProgramError::InvalidAccountData)?;
+        )
+        .map_err(|_| ProgramError::InvalidAccountData)?;
 
         Ok(())
     }

tokens/token-2022/transfer-hook/transfer-switch/anchor/programs/transfer-switch/src/instructions/transfer_hook.rs

@@ -3,21 +3,18 @@ use {
     anchor_lang::prelude::*,
     anchor_spl::{
         token_2022::spl_token_2022::{
-            extension::{
-                transfer_hook::TransferHookAccount, BaseStateWithExtensionsMut,
-                PodStateWithExtensionsMut,
-            },
+            extension::{transfer_hook::TransferHookAccount, BaseStateWithExtensionsMut, PodStateWithExtensionsMut},
             pod::PodAccount,
         },
-        token_interface::Mint,
+        token_interface::{Mint, TokenAccount},
     },
 };
 
 #[derive(Accounts)]
 pub struct TransferHook<'info> {
-    /// CHECK: Sender token account
+    /// Sender token account
     #[account()]
-    pub source_token_account: UncheckedAccount<'info>,
+    pub source_token_account: InterfaceAccount<'info, TokenAccount>,
 
     /// The mint of the token transferring
     #[account()]
@@ -27,7 +24,7 @@ pub struct TransferHook<'info> {
     #[account()]
     pub receiver_token_account: UncheckedAccount<'info>,
 
-    /// CHECK: the transfer sender
+    /// CHECK: the transfer authority (owner or delegate)
     #[account()]
     pub wallet: UncheckedAccount<'info>,
 
@@ -38,9 +35,9 @@ pub struct TransferHook<'info> {
     )]
     pub extra_account_metas_list: UncheckedAccount<'info>,
 
-    /// sender transfer switch
+    /// sender transfer switch, keyed by the source token account owner
     #[account(
-        seeds=[wallet.key().as_ref()],
+        seeds=[source_token_account.owner.as_ref()],
         bump,
     )]
     pub wallet_switch: Account<'info, TransferSwitch>,
@@ -61,8 +58,8 @@ impl<'info> TransferHook<'info> {
         // while anchor-lang 1.0 uses 3.x — structurally identical but different semver types
         let mut account = PodStateWithExtensionsMut::<PodAccount>::unpack(*account_data_ref)
             .map_err(|_| ProgramError::InvalidAccountData)?;
-        let account_extension = account.get_extension_mut::<TransferHookAccount>()
-            .map_err(|_| ProgramError::InvalidAccountData)?;
+        let account_extension =
+            account.get_extension_mut::<TransferHookAccount>().map_err(|_| ProgramError::InvalidAccountData)?;
 
         if !bool::from(account_extension.transferring) {
             return err!(TransferError::IsNotCurrentlyTransferring);

tokens/token-2022/transfer-hook/transfer-switch/anchor/tests/litesvm.test.ts

@@ -1,6 +1,7 @@
 import * as anchor from '@anchor-lang/core';
 import {
     ASSOCIATED_TOKEN_PROGRAM_ID,
+    createApproveCheckedInstruction,
     createAssociatedTokenAccountInstruction,
     createInitializeMintInstruction,
     createInitializeTransferHookInstruction,
@@ -15,19 +16,20 @@ import { Keypair, PublicKey, SystemProgram, Transaction, type TransactionInstruc
 import { getTokenDecoder } from '@solana-program/token-2022';
 import { LiteSVMProvider } from 'anchor-litesvm';
 import { assert } from 'chai';
-import { LiteSVM } from 'litesvm';
+import { FailedTransactionMetadata, LiteSVM } from 'litesvm';
 import IDL from '../target/idl/transfer_switch.json';
 import type { TransferSwitch } from '../target/types/transfer_switch';
 
 const PROGRAM_ID = new PublicKey(IDL.address);
 
 const expectRevert = async (promise: Promise<any>) => {
+    let reverted = false;
     try {
         await promise;
-        throw new Error('Expected a revert');
     } catch {
-        return;
+        reverted = true;
     }
+    assert(reverted, 'Expected a revert');
 };
 
 describe('Transfer switch', () => {
@@ -259,4 +261,123 @@ describe('Transfer switch', () => {
 
         assert(recipientBalance === bigIntAmount, 'transfer was not successful');
     });
+    // A delegate approved by the sender can move the sender's tokens, but the
+    // switch that gates the transfer must still be the SENDER's (source token
+    // account owner), not the delegate's.
+    describe('delegate transfers', () => {
+        const delegate = Keypair.generate();
+        const tokenBalance = (tokenAccount: PublicKey) =>
+            getTokenDecoder().decode(client.getAccount(tokenAccount).data).amount;
+
+        it('turn transfers off for sender, on for delegate', async () => {
+            // same bytes as the earlier switch(false) for sender
+            client.expireBlockhash();
+            await program.methods
+                .switch(false)
+                .accountsPartial({ wallet: sender.publicKey, admin: payer.publicKey })
+                .signers([payer])
+                .rpc();
+            await program.methods
+                .switch(true)
+                .accountsPartial({ wallet: delegate.publicKey, admin: payer.publicKey })
+                .signers([payer])
+                .rpc();
+
+            const senderSwitch = await program.account.transferSwitch.fetch(
+                walletTransferSwitchAddress(sender.publicKey),
+            );
+            const delegateSwitch = await program.account.transferSwitch.fetch(
+                walletTransferSwitchAddress(delegate.publicKey),
+            );
+            assert(!senderSwitch.on, 'sender switch not set to false');
+            assert(delegateSwitch.on, 'delegate switch not set to true');
+        });
+
+        it('sender approves the delegate', async () => {
+            const amount = BigInt(10 * 10 ** decimals);
+            const transaction = new Transaction().add(
+                createApproveCheckedInstruction(
+                    senderTokenAccount,
+                    mint.publicKey,
+                    delegate.publicKey,
+                    sender.publicKey,
+                    amount,
+                    decimals,
+                    [],
+                    TOKEN_2022_PROGRAM_ID,
+                ),
+            );
+            await provider.sendAndConfirm(transaction, [sender]);
+
+            const senderAccount = getTokenDecoder().decode(client.getAccount(senderTokenAccount).data);
+            assert(senderAccount.delegate.__option === 'Some', 'delegate not set');
+            assert(senderAccount.delegate.value === delegate.publicKey.toBase58(), 'delegate does not match');
+            assert(senderAccount.delegatedAmount === amount, 'delegated amount does not match');
+        });
+
+        it('Delegate transfer while sender switch is off, should fail!', async () => {
+            const amount = BigInt(1 * 10 ** decimals);
+            const [_recipient, recipientTokenAccount, recipientTokenAccountCreateIx] = newUser();
+            await provider.sendAndConfirm(new Transaction().add(recipientTokenAccountCreateIx));
+            const senderBalanceBefore = tokenBalance(senderTokenAccount);
+
+            const transferInstruction = await createTransferCheckedWithTransferHookInstruction(
+                connection,
+                senderTokenAccount,
+                mint.publicKey,
+                recipientTokenAccount,
+                delegate.publicKey, // delegate signs as the transfer authority
+                amount,
+                decimals,
+                [],
+                'confirmed',
+                TOKEN_2022_PROGRAM_ID,
+            );
+
+            const transaction = new Transaction().add(transferInstruction);
+            transaction.feePayer = payer.publicKey;
+            transaction.recentBlockhash = client.latestBlockhash();
+            transaction.sign(payer, delegate);
+
+            const result = client.sendTransaction(transaction);
+            assert(result instanceof FailedTransactionMetadata, 'delegate transfer succeeded with sender switch off');
+            assert(result.toString().includes('SwitchNotOn'), `unexpected error: ${result.toString()}`);
+
+            assert(tokenBalance(recipientTokenAccount) === BigInt(0), 'recipient received tokens');
+            assert(tokenBalance(senderTokenAccount) === senderBalanceBefore, 'sender balance changed');
+        });
+
+        it('turn on for sender, delegate transfer succeeds', async () => {
+            // same bytes as the earlier switch(true) for sender
+            client.expireBlockhash();
+            await program.methods
+                .switch(true)
+                .accountsPartial({ wallet: sender.publicKey, admin: payer.publicKey })
+                .signers([payer])
+                .rpc();
+
+            const amount = BigInt(1 * 10 ** decimals);
+            const [_recipient, recipientTokenAccount, recipientTokenAccountCreateIx] = newUser();
+            await provider.sendAndConfirm(new Transaction().add(recipientTokenAccountCreateIx));
+            const senderBalanceBefore = tokenBalance(senderTokenAccount);
+
+            const transferInstruction = await createTransferCheckedWithTransferHookInstruction(
+                connection,
+                senderTokenAccount,
+                mint.publicKey,
+                recipientTokenAccount,
+                delegate.publicKey,
+                amount,
+                decimals,
+                [],
+                'confirmed',
+                TOKEN_2022_PROGRAM_ID,
+            );
+
+            await provider.sendAndConfirm(new Transaction().add(transferInstruction), [delegate]);
+
+            assert(tokenBalance(recipientTokenAccount) === amount, 'transfer was not successful');
+            assert(tokenBalance(senderTokenAccount) === senderBalanceBefore - amount, 'sender balance not debited');
+        });
+    });
 });