p-token: Add explicitly ownership checks in `batch` (#80)

* Add transfer account owner checks

* Add batch transfer tests

* Update dependencies

* Fix lint

* Add more ownership checks

* More tests

* Typos

Cargo.lock

@@ -3072,11 +3072,15 @@ dependencies = [
 name = "pinocchio-token-program"
 version = "0.0.0"
 dependencies = [
+ "agave-feature-set",
  "assert_matches",
+ "mollusk-svm",
+ "mollusk-svm-fuzz-fixture",
  "num-traits",
  "pinocchio",
  "pinocchio-log",
  "pinocchio-token-interface",
+ "solana-account",
  "solana-instruction 2.3.0",
  "solana-keypair",
  "solana-program-error 2.2.2",
@@ -3084,6 +3088,8 @@ dependencies = [
  "solana-program-pack 2.2.1",
  "solana-program-test",
  "solana-pubkey 2.4.0",
+ "solana-rent",
+ "solana-sdk-ids 2.2.1",
  "solana-signature",
  "solana-signer",
  "solana-system-interface",

Cargo.toml

@@ -33,6 +33,8 @@ tag-message = "Publish {{crate_name}} v{{version}}"
 consolidate-commits = false
 
 [workspace.dependencies]
+mollusk-svm = "0.4.0"
+mollusk-svm-fuzz-fixture = "0.4.0"
 num-traits = "0.2"
 pinocchio = "0.9"
 solana-instruction = "2.3.0"

p-token/Cargo.toml

@@ -20,8 +20,12 @@ pinocchio-log = { version = "0.5", default-features = false }
 pinocchio-token-interface = { version = "^0", path = "../p-interface" }
 
 [dev-dependencies]
+agave-feature-set = "2.2.20"
 assert_matches = "1.5.0"
+mollusk-svm = { workspace = true }
+mollusk-svm-fuzz-fixture = { workspace = true }
 num-traits = { workspace = true }
+solana-account = "2.2.1"
 solana-instruction = { workspace = true }
 solana-keypair = "2.2.3"
 solana-program-error = { workspace = true }
@@ -29,6 +33,8 @@ solana-program-option = { workspace = true }
 solana-program-pack = { workspace = true }
 solana-program-test = "2.3.4"
 solana-pubkey = { workspace = true }
+solana-rent = "2.2.1"
+solana-sdk-ids = "2.2.1"
 solana-signature = "2.3.0"
 solana-signer = "2.2.1"
 solana-transaction = "2.2.3"

p-token/src/processor/batch.rs

@@ -1,5 +1,5 @@
 use {
-    crate::entrypoint::inner_process_instruction,
+    crate::{entrypoint::inner_process_instruction, processor::check_account_owner},
     pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
     pinocchio_token_interface::error::TokenError,
 };
@@ -46,6 +46,54 @@ pub fn process_batch(mut accounts: &[AccountInfo], mut instruction_data: &[u8])
             )
         };
 
+        // Few Instructions require specific account ownership checks when executed
+        // in a batch since ownership is only enforced by the runtime at the end of
+        // the batch processing.
+        //
+        // Instructions that do not appear in the list below do not require
+        // ownership checks since they either do not modify accounts or the ownership
+        // is already checked explicitly.
+        if let Some(&discriminator) = ix_data.first() {
+            match discriminator {
+                // 3 - Transfer
+                // 7 - MintTo
+                // 8 - Burn
+                // 14 - MintToChecked
+                // 15 - BurnChecked
+                3 | 7 | 8 | 14 | 15 => {
+                    let [a0, a1, ..] = ix_accounts else {
+                        return Err(ProgramError::NotEnoughAccountKeys);
+                    };
+                    check_account_owner(a0)?;
+                    check_account_owner(a1)?;
+                }
+                // 12 - TransferChecked
+                12 => {
+                    let [a0, _, a2, ..] = ix_accounts else {
+                        return Err(ProgramError::NotEnoughAccountKeys);
+                    };
+                    check_account_owner(a0)?;
+                    check_account_owner(a2)?;
+                }
+                // 4 - Approve
+                // 5 - Revoke
+                // 6 - SetAuthority
+                // 9 - CloseAccount
+                // 10 - FreezeAccount
+                // 11 - ThawAccount
+                // 13 - ApproveChecked
+                // 22 - InitializeImmutableOwner
+                // 38 - WithdrawExcessLamports
+                4..=13 | 22 | 38 => {
+                    let [a0, ..] = ix_accounts else {
+                        return Err(ProgramError::NotEnoughAccountKeys);
+                    };
+                    check_account_owner(a0)?;
+                }
+                _ => {}
+            }
+        }
+
         inner_process_instruction(ix_accounts, ix_data)?;
 
         if data_offset == instruction_data.len() {

p-token/tests/batch.rs

@@ -2,12 +2,24 @@ mod setup;
 
 use {
     crate::setup::TOKEN_PROGRAM_ID,
+    agave_feature_set::FeatureSet,
+    mollusk_svm::{result::Check, Mollusk},
+    pinocchio_token_interface::{
+        native_mint,
+        state::{
+            account::Account as TokenAccount, account_state::AccountState, load_mut_unchecked,
+            mint::Mint,
+        },
+    },
+    solana_account::Account,
     solana_instruction::{AccountMeta, Instruction},
     solana_keypair::Keypair,
     solana_program_error::ProgramError,
     solana_program_pack::Pack,
     solana_program_test::{tokio, ProgramTest},
     solana_pubkey::Pubkey,
+    solana_rent::Rent,
+    solana_sdk_ids::bpf_loader_upgradeable,
     solana_signer::Signer,
     solana_system_interface::instruction::create_account,
     solana_transaction::Transaction,
@@ -40,7 +52,7 @@ fn batch_instruction(instructions: Vec<Instruction>) -> Result<Instruction, Prog
 }
 
 #[tokio::test]
-async fn batch() {
+async fn batch_initialize_mint_transfer_close() {
     let context = ProgramTest::new("pinocchio_token_program", TOKEN_PROGRAM_ID, None)
         .start_with_context()
         .await;
@@ -218,4 +230,534 @@ async fn batch() {
     let owner_b_ta_a_account =
         spl_token::state::Account::unpack(&owner_b_ta_a_account.unwrap().data).unwrap();
     assert_eq!(owner_b_ta_a_account.amount, 1000000);
+
+    let closed_owner_a_ta_a = context
+        .banks_client
+        .get_account(owner_a_ta_a.pubkey())
+        .await
+        .unwrap();
+    assert!(closed_owner_a_ta_a.is_none());
+}
+
+fn create_mint(
+    mint_authority: &Pubkey,
+    supply: u64,
+    decimals: u8,
+    program_owner: &Pubkey,
+) -> Account {
+    let space = size_of::<Mint>();
+    let lamports = Rent::default().minimum_balance(space);
+
+    let mut data: Vec<u8> = vec![0u8; space];
+    let mint = unsafe { load_mut_unchecked::<Mint>(data.as_mut_slice()).unwrap() };
+    mint.set_initialized();
+    mint.set_supply(supply);
+    mint.set_mint_authority(mint_authority.as_array());
+    mint.decimals = decimals;
+
+    Account {
+        lamports,
+        data,
+        owner: *program_owner,
+        executable: false,
+        ..Default::default()
+    }
+}
+
+fn create_token_account(
+    mint: &Pubkey,
+    owner: &Pubkey,
+    is_native: bool,
+    amount: u64,
+    program_owner: &Pubkey,
+) -> Account {
+    let space = size_of::<TokenAccount>();
+    let mut lamports = Rent::default().minimum_balance(space);
+
+    let mut data: Vec<u8> = vec![0u8; space];
+    let token = unsafe { load_mut_unchecked::<TokenAccount>(data.as_mut_slice()).unwrap() };
+    token.set_account_state(AccountState::Initialized);
+    token.mint = *mint.as_array();
+    token.owner = *owner.as_array();
+    token.set_amount(amount);
+    token.set_native(is_native);
+    token.set_native_amount(amount);
+
+    if is_native {
+        lamports = lamports.saturating_add(amount);
+    }
+
+    Account {
+        lamports,
+        data,
+        owner: *program_owner,
+        executable: false,
+        ..Default::default()
+    }
+}
+
+/// Creates a Mollusk instance with the default feature set, excluding the
+/// `bpf_account_data_direct_mapping` feature.
+fn mollusk() -> Mollusk {
+    let feature_set = {
+        let mut fs = FeatureSet::all_enabled();
+        fs.active_mut()
+            .remove(&agave_feature_set::bpf_account_data_direct_mapping::id());
+        fs
+    };
+    let mut mollusk = Mollusk {
+        feature_set,
+        ..Default::default()
+    };
+    mollusk.add_program(
+        &TOKEN_PROGRAM_ID,
+        "pinocchio_token_program",
+        &bpf_loader_upgradeable::id(),
+    );
+    mollusk
+}
+
+#[tokio::test]
+async fn batch_transfer() {
+    let authority_key = Pubkey::new_unique();
+    let mint_key = Pubkey::new_unique();
+
+    // source account
+    //   - amount: 1_000_000_000
+    //   - mint: mint_key
+    //   - is_native: false
+    //   - program_id: TOKEN_PROGRAM_ID
+    let source_account_key = Pubkey::new_unique();
+    let source_account = create_token_account(
+        &mint_key,
+        &authority_key,
+        false,
+        1_000_000_000,
+        &TOKEN_PROGRAM_ID,
+    );
+
+    // destination account
+    //   - amount: 0
+    //   - mint: mint_key
+    //   - is_native: false
+    //   - program_id: TOKEN_PROGRAM_ID
+    let destination_account_key = Pubkey::new_unique();
+    let destination_account =
+        create_token_account(&mint_key, &authority_key, false, 0, &TOKEN_PROGRAM_ID);
+
+    let instruction = batch_instruction(vec![spl_token::instruction::transfer(
+        &TOKEN_PROGRAM_ID,
+        &source_account_key,
+        &destination_account_key,
+        &authority_key,
+        &[],
+        500_000_000,
+    )
+    .unwrap()])
+    .unwrap();
+
+    // Expected to succeed.
+
+    mollusk().process_and_validate_instruction_chain(
+        &[(&instruction, &[Check::success(), Check::all_rent_exempt()])],
+        &[
+            (source_account_key, source_account),
+            (destination_account_key, destination_account),
+            (
+                authority_key,
+                Account {
+                    lamports: Rent::default().minimum_balance(0),
+                    ..Default::default()
+                },
+            ),
+        ],
+    );
+}
+
+#[tokio::test]
+async fn batch_fail_transfer_with_invalid_program_owner() {
+    let invalid_program_id = Pubkey::new_from_array([2; 32]);
+    let native_mint = Pubkey::new_from_array(native_mint::ID);
+    let authority_key = Pubkey::new_unique();
+
+    // source account
+    //   - amount: 1_000_000_000
+    //   - mint: native_mint
+    //   - is_native: true
+    //   - program_id: invalid_program_id
+    let source_account_key = Pubkey::new_unique();
+    let source_account = create_token_account(
+        &native_mint,
+        &authority_key,
+        true,
+        1_000_000_000,
+        &invalid_program_id,
+    );
+
+    // destination account
+    //   - amount: 0
+    //   - mint: native_mint
+    //   - is_native: true
+    //   - program_id: TOKEN_PROGRAM_ID
+    let destination_account_key = Pubkey::new_unique();
+    let destination_account =
+        create_token_account(&native_mint, &authority_key, true, 0, &TOKEN_PROGRAM_ID);
+
+    let instruction = batch_instruction(vec![spl_token::instruction::transfer(
+        &TOKEN_PROGRAM_ID,
+        &source_account_key,
+        &destination_account_key,
+        &authority_key,
+        &[],
+        500_000_000,
+    )
+    .unwrap()])
+    .unwrap();
+
+    // Expected to fail since source account has an invalid program owner.
+
+    mollusk().process_and_validate_instruction_chain(
+        &[(
+            &instruction,
+            &[
+                Check::err(ProgramError::IncorrectProgramId),
+                Check::all_rent_exempt(),
+            ],
+        )],
+        &[
+            (source_account_key, source_account),
+            (destination_account_key, destination_account),
+            (
+                authority_key,
+                Account {
+                    lamports: Rent::default().minimum_balance(0),
+                    ..Default::default()
+                },
+            ),
+        ],
+    );
+}
+
+#[tokio::test]
+async fn batch_fail_transfer_checked_with_invalid_program_owner() {
+    let invalid_program_id = Pubkey::new_from_array([2; 32]);
+    let authority_key = Pubkey::new_unique();
+
+    let native_mint_key = Pubkey::new_from_array(native_mint::ID);
+    let native_mint = create_mint(&authority_key, 5_000_000_000, 9, &TOKEN_PROGRAM_ID);
+
+    // source account
+    //   - amount: 1_000_000_000
+    //   - mint: native_mint
+    //   - is_native: true
+    //   - program_id: invalid_program_id
+    let source_account_key = Pubkey::new_unique();
+    let source_account = create_token_account(
+        &native_mint_key,
+        &authority_key,
+        true,
+        1_000_000_000,
+        &invalid_program_id,
+    );
+
+    // destination account
+    //   - amount: 0
+    //   - mint: native_mint
+    //   - is_native: true
+    //   - program_id: TOKEN_PROGRAM_ID
+    let destination_account_key = Pubkey::new_unique();
+    let destination_account =
+        create_token_account(&native_mint_key, &authority_key, true, 0, &TOKEN_PROGRAM_ID);
+
+    let instruction = batch_instruction(vec![spl_token::instruction::transfer_checked(
+        &TOKEN_PROGRAM_ID,
+        &source_account_key,
+        &native_mint_key,
+        &destination_account_key,
+        &authority_key,
+        &[],
+        500_000_000,
+        9,
+    )
+    .unwrap()])
+    .unwrap();
+
+    // Expected to fail since source account has an invalid program owner.
+
+    mollusk().process_and_validate_instruction_chain(
+        &[(
+            &instruction,
+            &[
+                Check::err(ProgramError::IncorrectProgramId),
+                Check::all_rent_exempt(),
+            ],
+        )],
+        &[
+            (source_account_key, source_account),
+            (destination_account_key, destination_account),
+            (native_mint_key, native_mint),
+            (
+                authority_key,
+                Account {
+                    lamports: Rent::default().minimum_balance(0),
+                    ..Default::default()
+                },
+            ),
+        ],
+    );
+}
+
+#[tokio::test]
+async fn batch_fail_swap_tokens_with_invalid_program_owner() {
+    let native_mint = Pubkey::new_from_array(native_mint::ID);
+    let invalid_program_id = Pubkey::new_from_array([2; 32]);
+    let authority_key = Pubkey::new_unique();
+
+    // Account A
+    //   - amount: 1_000
+    //   - mint: native_mint
+    //   - is_native: false
+    //   - program_id: invalid_program_id
+    let account_a_key = Pubkey::new_unique();
+    let account_a = create_token_account(
+        &native_mint,
+        &authority_key,
+        false,
+        1_000,
+        &invalid_program_id,
+    );
+
+    // Account B
+    //   - amount: 0
+    //   - mint: native_mint
+    //   - is_native: true
+    //   - program_id: TOKEN_PROGRAM_ID
+    let account_b_key = Pubkey::new_unique();
+    let account_b = create_token_account(&native_mint, &authority_key, true, 0, &TOKEN_PROGRAM_ID);
+
+    // Account C
+    //   - amount: 0
+    //   - mint: native_mint
+    //   - is_native: true
+    //   - program_id: TOKEN_PROGRAM_ID
+    let account_c_key = Pubkey::new_unique();
+    let account_c =
+        create_token_account(&native_mint, &authority_key, true, 1_000, &TOKEN_PROGRAM_ID);
+
+    // Batch instruction to swap tokens
+    //   - transfer 300 from account A to account B
+    //   - transfer 300 from account C to account A
+    let instruction = batch_instruction(vec![
+        spl_token::instruction::sync_native(&TOKEN_PROGRAM_ID, &account_b_key).unwrap(),
+        spl_token::instruction::sync_native(&TOKEN_PROGRAM_ID, &account_c_key).unwrap(),
+        spl_token::instruction::transfer(
+            &TOKEN_PROGRAM_ID,
+            &account_a_key,
+            &account_b_key,
+            &authority_key,
+            &[],
+            300,
+        )
+        .unwrap(),
+        spl_token::instruction::transfer(
+            &TOKEN_PROGRAM_ID,
+            &account_c_key,
+            &account_a_key,
+            &authority_key,
+            &[],
+            300,
+        )
+        .unwrap(),
+    ])
+    .unwrap();
+
+    // Expected to fail since account A has an invalid program owner.
+
+    mollusk().process_and_validate_instruction_chain(
+        &[(
+            &instruction,
+            &[
+                Check::err(ProgramError::IncorrectProgramId),
+                Check::all_rent_exempt(),
+            ],
+        )],
+        &[
+            (account_a_key, account_a),
+            (account_b_key, account_b),
+            (account_c_key, account_c),
+            (
+                authority_key,
+                Account {
+                    lamports: Rent::default().minimum_balance(0),
+                    ..Default::default()
+                },
+            ),
+        ],
+    );
+}
+
+#[tokio::test]
+async fn batch_fail_mint_to_with_invalid_program_owner() {
+    let invalid_program_id = Pubkey::new_from_array([2; 32]);
+    let authority_key = Pubkey::new_unique();
+
+    let mint_key = Pubkey::new_unique();
+    let mint = create_mint(&authority_key, 0, 0, &TOKEN_PROGRAM_ID);
+
+    // account A (invalid)
+    //   - amount: 1_000_000_000
+    //   - mint: native_mint
+    //   - is_native: false
+    //   - program_id: invalid_program_id
+    let account_a_key = Pubkey::new_unique();
+    let account_a = create_token_account(&mint_key, &authority_key, false, 0, &invalid_program_id);
+
+    // account B
+    //   - amount: 0
+    //   - mint: native_mint
+    //   - is_native: false
+    //   - program_id: TOKEN_PROGRAM_ID
+    let account_b_key = Pubkey::new_unique();
+    let account_b = create_token_account(&mint_key, &authority_key, false, 0, &TOKEN_PROGRAM_ID);
+
+    let instruction = batch_instruction(vec![
+        spl_token::instruction::mint_to(
+            &TOKEN_PROGRAM_ID,
+            &mint_key,
+            &account_a_key,
+            &authority_key,
+            &[],
+            1_000_000_000,
+        )
+        .unwrap(),
+        spl_token::instruction::mint_to(
+            &TOKEN_PROGRAM_ID,
+            &mint_key,
+            &account_b_key,
+            &authority_key,
+            &[],
+            1_000_000_000,
+        )
+        .unwrap(),
+    ])
+    .unwrap();
+
+    // Expected to fail since source account has an invalid program owner.
+
+    mollusk().process_and_validate_instruction_chain(
+        &[(
+            &instruction,
+            &[
+                Check::err(ProgramError::IncorrectProgramId),
+                Check::all_rent_exempt(),
+            ],
+        )],
+        &[
+            (mint_key, mint),
+            (account_a_key, account_a),
+            (account_b_key, account_b),
+            (
+                authority_key,
+                Account {
+                    lamports: Rent::default().minimum_balance(0),
+                    ..Default::default()
+                },
+            ),
+        ],
+    );
+}
+
+#[tokio::test]
+async fn batch_fail_burn_with_invalid_program_owner() {
+    let invalid_program_id = Pubkey::new_from_array([2; 32]);
+    let authority_key = Pubkey::new_unique();
+
+    let mint_key = Pubkey::new_unique();
+    let mint = create_mint(&authority_key, 2_000_000_000, 0, &TOKEN_PROGRAM_ID);
+
+    // account A
+    //   - amount: 0
+    //   - mint: native_mint
+    //   - is_native: false
+    //   - program_id: TOKEN_PROGRAM_ID
+    let account_a_key = Pubkey::new_unique();
+    let account_a = create_token_account(&mint_key, &authority_key, false, 0, &TOKEN_PROGRAM_ID);
+
+    // account B (invalid)
+    //   - amount: 1_000_000_000
+    //   - mint: native_mint
+    //   - is_native: false
+    //   - program_id: invalid_program_id
+    let account_b_key = Pubkey::new_unique();
+    let account_b = create_token_account(
+        &mint_key,
+        &authority_key,
+        false,
+        1_000_000_000,
+        &invalid_program_id,
+    );
+
+    let instruction = batch_instruction(vec![
+        spl_token::instruction::mint_to(
+            &TOKEN_PROGRAM_ID,
+            &mint_key,
+            &account_a_key,
+            &authority_key,
+            &[],
+            1_000_000_000,
+        )
+        .unwrap(),
+        spl_token::instruction::mint_to(
+            &TOKEN_PROGRAM_ID,
+            &mint_key,
+            &account_b_key,
+            &authority_key,
+            &[],
+            1_000_000_000,
+        )
+        .unwrap(),
+        spl_token::instruction::burn(
+            &TOKEN_PROGRAM_ID,
+            &account_a_key,
+            &mint_key,
+            &authority_key,
+            &[],
+            1_000_000_000,
+        )
+        .unwrap(),
+        spl_token::instruction::burn(
+            &TOKEN_PROGRAM_ID,
+            &account_b_key,
+            &mint_key,
+            &authority_key,
+            &[],
+            1_000_000_000,
+        )
+        .unwrap(),
+    ])
+    .unwrap();
+
+    // Expected to fail since source account has an invalid program owner.
+
+    mollusk().process_and_validate_instruction_chain(
+        &[(
+            &instruction,
+            &[
+                Check::err(ProgramError::IncorrectProgramId),
+                Check::all_rent_exempt(),
+            ],
+        )],
+        &[
+            (mint_key, mint),
+            (account_a_key, account_a),
+            (account_b_key, account_b),
+            (
+                authority_key,
+                Account {
+                    lamports: Rent::default().minimum_balance(0),
+                    ..Default::default()
+                },
+            ),
+        ],
+    );
 }

program/Cargo.toml

@@ -35,8 +35,8 @@ thiserror = "2.0"
 
 [dev-dependencies]
 lazy_static = "1.5.0"
-mollusk-svm = "0.4.0"
-mollusk-svm-fuzz-fixture = "0.4.0"
+mollusk-svm = { workspace = true }
+mollusk-svm-fuzz-fixture = { workspace = true }
 proptest = "1.5"
 serial_test = "3.2.0"
 solana-account = "2.2.1"