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

Author: febo

Cargo.lock

@@ -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"

Cargo.toml

@@ -20,15 +20,21 @@ 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 }
 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/Cargo.toml

@@ -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() {