diff --git a/idl/escrow_program.json b/idl/escrow_program.json index 32f7a65..4b7a7d8 100644 --- a/idl/escrow_program.json +++ b/idl/escrow_program.json @@ -61,7 +61,7 @@ { "defaultValue": { "kind": "numberValueNode", - "number": 3 + "number": 4 }, "kind": "structFieldTypeNode", "name": "discriminator", @@ -103,7 +103,7 @@ { "defaultValue": { "kind": "numberValueNode", - "number": 0 + "number": 1 }, "kind": "structFieldTypeNode", "name": "discriminator", @@ -159,7 +159,7 @@ { "defaultValue": { "kind": "numberValueNode", - "number": 1 + "number": 2 }, "kind": "structFieldTypeNode", "name": "discriminator", @@ -210,7 +210,7 @@ { "defaultValue": { "kind": "numberValueNode", - "number": 2 + "number": 3 }, "kind": "structFieldTypeNode", "name": "discriminator", diff --git a/program/src/instructions/block_mint/processor.rs b/program/src/instructions/block_mint/processor.rs index b6f6bba..e1b12ff 100644 --- a/program/src/instructions/block_mint/processor.rs +++ b/program/src/instructions/block_mint/processor.rs @@ -3,8 +3,8 @@ use pinocchio::{account::AccountView, Address, ProgramResult}; use crate::{ events::BlockMintEvent, instructions::BlockMint, - state::{AllowedMint, AllowedMintPda, Escrow}, - traits::{EventSerialize, PdaSeeds}, + state::{AllowedMint, Escrow}, + traits::EventSerialize, utils::{close_pda_account, emit_event}, }; @@ -19,13 +19,15 @@ pub fn process_block_mint(program_id: &Address, accounts: &[AccountView], instru let escrow = Escrow::from_account(&escrow_data, ix.accounts.escrow, program_id)?; escrow.validate_admin(ix.accounts.admin.address())?; - // Verify allowed_mint exists and is valid + // Verify allowed_mint account exists and self-validates against escrow + mint PDA derivation let allowed_mint_data = ix.accounts.allowed_mint.try_borrow()?; - let allowed_mint = AllowedMint::from_account(&allowed_mint_data)?; - - // Validate that allowed_mint PDA matches the escrow + mint combination - let pda_seeds = AllowedMintPda::new(ix.accounts.escrow.address(), ix.accounts.mint.address()); - pda_seeds.validate_pda(ix.accounts.allowed_mint, program_id, allowed_mint.bump)?; + let _allowed_mint = AllowedMint::from_account( + &allowed_mint_data, + ix.accounts.allowed_mint, + program_id, + ix.accounts.escrow.address(), + ix.accounts.mint.address(), + )?; drop(allowed_mint_data); // Close the AllowedMint account and return lamports to rent_recipient diff --git a/program/src/instructions/deposit/accounts.rs b/program/src/instructions/deposit/accounts.rs index 55e5e48..8a75f05 100644 --- a/program/src/instructions/deposit/accounts.rs +++ b/program/src/instructions/deposit/accounts.rs @@ -4,8 +4,8 @@ use crate::{ traits::InstructionAccounts, utils::{ validate_associated_token_account, verify_current_program, verify_current_program_account, - verify_event_authority, verify_readonly, verify_signer, verify_system_program, verify_token_program, - verify_writable, + verify_event_authority, verify_owned_by, verify_readonly, verify_signer, verify_system_program, + verify_token_program, verify_writable, }, }; @@ -77,6 +77,7 @@ impl<'a> TryFrom<&'a [AccountView]> for DepositAccounts<'a> { // 4. Validate program IDs verify_token_program(token_program)?; + verify_owned_by(mint, token_program.address())?; verify_system_program(system_program)?; verify_current_program(escrow_program)?; verify_event_authority(event_authority)?; diff --git a/program/src/instructions/deposit/processor.rs b/program/src/instructions/deposit/processor.rs index 850f3b3..2d692e9 100644 --- a/program/src/instructions/deposit/processor.rs +++ b/program/src/instructions/deposit/processor.rs @@ -13,11 +13,11 @@ use crate::{ events::DepositEvent, instructions::Deposit, state::{ - get_extensions_from_account, validate_extensions_pda, AllowedMint, AllowedMintPda, Escrow, ExtensionType, - HookData, HookPoint, Receipt, + get_extensions_from_account, validate_extensions_pda, AllowedMint, Escrow, ExtensionType, HookData, HookPoint, + Receipt, }, traits::{AccountSerialize, AccountSize, EventSerialize, ExtensionData, PdaSeeds}, - utils::{create_pda_account, emit_event, get_mint_decimals}, + utils::{create_pda_account, emit_event, get_mint_decimals, validate_mint_extensions}, }; /// Processes the Deposit instruction. @@ -30,15 +30,16 @@ pub fn process_deposit(program_id: &Address, accounts: &[AccountView], instructi let escrow_data = ix.accounts.escrow.try_borrow()?; let _escrow = Escrow::from_account(&escrow_data, ix.accounts.escrow, program_id)?; - // Verify allowed_mint PDA exists and matches expected derivation + // Verify allowed_mint account exists and self-validates against escrow + mint PDA derivation let allowed_mint_data = ix.accounts.allowed_mint.try_borrow()?; - let allowed_mint = AllowedMint::from_account(&allowed_mint_data).map_err(|_| EscrowProgramError::MintNotAllowed)?; - - // Validate that the allowed_mint PDA is derived from the correct escrow + mint - let pda_seeds = AllowedMintPda::new(ix.accounts.escrow.address(), ix.accounts.mint.address()); - pda_seeds - .validate_pda(ix.accounts.allowed_mint, program_id, allowed_mint.bump) - .map_err(|_| EscrowProgramError::MintNotAllowed)?; + let _allowed_mint = AllowedMint::from_account( + &allowed_mint_data, + ix.accounts.allowed_mint, + program_id, + ix.accounts.escrow.address(), + ix.accounts.mint.address(), + ) + .map_err(|_| EscrowProgramError::MintNotAllowed)?; // Get current timestamp from Clock sysvar let clock = Clock::get()?; @@ -74,6 +75,10 @@ pub fn process_deposit(program_id: &Address, accounts: &[AccountView], instructi // Validate extensions PDA validate_extensions_pda(ix.accounts.escrow, ix.accounts.extensions, program_id)?; + // Re-check mint extensions against the current escrow blocklist. + // This prevents stale AllowedMint entries from bypassing new blocklist rules. + validate_mint_extensions(ix.accounts.mint, ix.accounts.extensions)?; + // Get hook extension if present let exts = get_extensions_from_account(ix.accounts.extensions, &[ExtensionType::Hook])?; let hook_data = exts[0].as_ref().map(|b| HookData::from_bytes(b)).transpose()?; @@ -83,7 +88,7 @@ pub fn process_deposit(program_id: &Address, accounts: &[AccountView], instructi hook.invoke( HookPoint::PreDeposit, ix.accounts.remaining_accounts, - &[ix.accounts.escrow, ix.accounts.depositor, ix.accounts.mint, ix.accounts.receipt], + &[ix.accounts.escrow, ix.accounts.mint, ix.accounts.receipt], )?; } @@ -106,7 +111,7 @@ pub fn process_deposit(program_id: &Address, accounts: &[AccountView], instructi hook.invoke( HookPoint::PostDeposit, ix.accounts.remaining_accounts, - &[ix.accounts.escrow, ix.accounts.depositor, ix.accounts.mint, ix.accounts.receipt], + &[ix.accounts.escrow, ix.accounts.mint, ix.accounts.receipt], )?; } diff --git a/program/src/instructions/extensions/add_timelock/processor.rs b/program/src/instructions/extensions/add_timelock/processor.rs index e6fabb1..3ebd0ee 100644 --- a/program/src/instructions/extensions/add_timelock/processor.rs +++ b/program/src/instructions/extensions/add_timelock/processor.rs @@ -4,9 +4,9 @@ use pinocchio::{account::AccountView, cpi::Seed, error::ProgramError, Address, P use crate::{ events::TimelockAddedEvent, instructions::AddTimelock, - state::{append_extension, Escrow, ExtensionType, ExtensionsPda, TimelockData}, - traits::{EventSerialize, PdaSeeds}, - utils::{emit_event, TlvWriter}, + state::{update_or_append_extension, Escrow, ExtensionType, ExtensionsPda, TimelockData}, + traits::{EventSerialize, ExtensionData, PdaSeeds}, + utils::emit_event, }; /// Processes the AddTimelock instruction. @@ -24,23 +24,22 @@ pub fn process_add_timelock(program_id: &Address, accounts: &[AccountView], inst let extensions_pda = ExtensionsPda::new(ix.accounts.escrow.address()); extensions_pda.validate_pda(ix.accounts.extensions, program_id, ix.data.extensions_bump)?; - // Build TLV data + // Build extension data let timelock = TimelockData::new(ix.data.lock_duration); - let mut tlv_writer = TlvWriter::new(); - tlv_writer.write_timelock(&timelock); + let timelock_bytes = timelock.to_bytes(); - // Get seeds and append extension + // Get seeds and append/update extension let extensions_bump_seed = [ix.data.extensions_bump]; let extensions_seeds: Vec = extensions_pda.seeds_with_bump(&extensions_bump_seed); let extensions_seeds_array: [Seed; 3] = extensions_seeds.try_into().map_err(|_| ProgramError::InvalidArgument)?; - append_extension( + update_or_append_extension( ix.accounts.payer, ix.accounts.extensions, program_id, ix.data.extensions_bump, ExtensionType::Timelock, - &tlv_writer.into_bytes(), + &timelock_bytes, extensions_seeds_array, )?; diff --git a/program/src/instructions/extensions/set_arbiter/processor.rs b/program/src/instructions/extensions/set_arbiter/processor.rs index bedfbcd..bc60322 100644 --- a/program/src/instructions/extensions/set_arbiter/processor.rs +++ b/program/src/instructions/extensions/set_arbiter/processor.rs @@ -4,15 +4,14 @@ use pinocchio::{account::AccountView, cpi::Seed, error::ProgramError, Address, P use crate::{ events::ArbiterSetEvent, instructions::SetArbiter, - state::{append_extension, ArbiterData, Escrow, ExtensionType, ExtensionsPda}, - traits::{EventSerialize, PdaSeeds}, - utils::{emit_event, TlvWriter}, + state::{update_or_append_extension, ArbiterData, Escrow, ExtensionType, ExtensionsPda}, + traits::{EventSerialize, ExtensionData, PdaSeeds}, + utils::emit_event, }; /// Processes the SetArbiter instruction. /// /// Sets the arbiter on an escrow. Creates extensions PDA if it doesn't exist. -/// The arbiter is immutable — this instruction will fail if an arbiter is already set. pub fn process_set_arbiter(program_id: &Address, accounts: &[AccountView], instruction_data: &[u8]) -> ProgramResult { let ix = SetArbiter::try_from((instruction_data, accounts))?; @@ -25,23 +24,22 @@ pub fn process_set_arbiter(program_id: &Address, accounts: &[AccountView], instr let extensions_pda = ExtensionsPda::new(ix.accounts.escrow.address()); extensions_pda.validate_pda(ix.accounts.extensions, program_id, ix.data.extensions_bump)?; - // Build TLV data + // Build extension data let arbiter = ArbiterData::new(*ix.accounts.arbiter.address()); - let mut tlv_writer = TlvWriter::new(); - tlv_writer.write_arbiter(&arbiter); + let arbiter_bytes = arbiter.to_bytes(); - // Get seeds and append extension (fails if arbiter already exists, enforcing immutability) + // Get seeds and append/update extension let extensions_bump_seed = [ix.data.extensions_bump]; let extensions_seeds: Vec = extensions_pda.seeds_with_bump(&extensions_bump_seed); let extensions_seeds_array: [Seed; 3] = extensions_seeds.try_into().map_err(|_| ProgramError::InvalidArgument)?; - append_extension( + update_or_append_extension( ix.accounts.payer, ix.accounts.extensions, program_id, ix.data.extensions_bump, ExtensionType::Arbiter, - &tlv_writer.into_bytes(), + &arbiter_bytes, extensions_seeds_array, )?; diff --git a/program/src/instructions/extensions/set_hook/processor.rs b/program/src/instructions/extensions/set_hook/processor.rs index 6c9e027..8529194 100644 --- a/program/src/instructions/extensions/set_hook/processor.rs +++ b/program/src/instructions/extensions/set_hook/processor.rs @@ -4,9 +4,9 @@ use pinocchio::{account::AccountView, cpi::Seed, error::ProgramError, Address, P use crate::{ events::HookSetEvent, instructions::SetHook, - state::{append_extension, Escrow, ExtensionType, ExtensionsPda, HookData}, - traits::{EventSerialize, PdaSeeds}, - utils::{emit_event, TlvWriter}, + state::{update_or_append_extension, Escrow, ExtensionType, ExtensionsPda, HookData}, + traits::{EventSerialize, ExtensionData, PdaSeeds}, + utils::emit_event, }; /// Processes the SetHook instruction. @@ -24,23 +24,22 @@ pub fn process_set_hook(program_id: &Address, accounts: &[AccountView], instruct let extensions_pda = ExtensionsPda::new(ix.accounts.escrow.address()); extensions_pda.validate_pda(ix.accounts.extensions, program_id, ix.data.extensions_bump)?; - // Build TLV data + // Build extension data let hook = HookData::new(ix.data.hook_program); - let mut tlv_writer = TlvWriter::new(); - tlv_writer.write_hook(&hook); + let hook_bytes = hook.to_bytes(); - // Get seeds and append extension + // Get seeds and append/update extension let extensions_bump_seed = [ix.data.extensions_bump]; let extensions_seeds: Vec = extensions_pda.seeds_with_bump(&extensions_bump_seed); let extensions_seeds_array: [Seed; 3] = extensions_seeds.try_into().map_err(|_| ProgramError::InvalidArgument)?; - append_extension( + update_or_append_extension( ix.accounts.payer, ix.accounts.extensions, program_id, ix.data.extensions_bump, ExtensionType::Hook, - &tlv_writer.into_bytes(), + &hook_bytes, extensions_seeds_array, )?; diff --git a/program/src/instructions/withdraw/accounts.rs b/program/src/instructions/withdraw/accounts.rs index 99c8604..2ab6f29 100644 --- a/program/src/instructions/withdraw/accounts.rs +++ b/program/src/instructions/withdraw/accounts.rs @@ -4,8 +4,8 @@ use crate::{ traits::InstructionAccounts, utils::{ validate_associated_token_account, verify_current_program, verify_current_program_account, - verify_event_authority, verify_readonly, verify_signer, verify_system_program, verify_token_program, - verify_writable, + verify_event_authority, verify_owned_by, verify_readonly, verify_signer, verify_system_program, + verify_token_program, verify_writable, }, }; @@ -74,6 +74,7 @@ impl<'a> TryFrom<&'a [AccountView]> for WithdrawAccounts<'a> { // 4. Validate program IDs verify_token_program(token_program)?; + verify_owned_by(mint, token_program.address())?; verify_system_program(system_program)?; verify_current_program(escrow_program)?; verify_event_authority(event_authority)?; diff --git a/program/src/instructions/withdraw/processor.rs b/program/src/instructions/withdraw/processor.rs index 404e513..5747df4 100644 --- a/program/src/instructions/withdraw/processor.rs +++ b/program/src/instructions/withdraw/processor.rs @@ -1,4 +1,4 @@ -use pinocchio::{account::AccountView, Address, ProgramResult}; +use pinocchio::{account::AccountView, error::ProgramError, Address, ProgramResult}; use pinocchio_token_2022::instructions::TransferChecked; use crate::{ @@ -25,7 +25,7 @@ pub fn process_withdraw(program_id: &Address, accounts: &[AccountView], instruct } // Read and validate receipt - let (amount, receipt_seed, mint, deposited_at) = { + let (amount, receipt_seed, receipt_mint, deposited_at) = { let receipt_data = ix.accounts.receipt.try_borrow()?; let receipt = Receipt::from_account(&receipt_data, ix.accounts.receipt, program_id)?; @@ -35,6 +35,11 @@ pub fn process_withdraw(program_id: &Address, accounts: &[AccountView], instruct (receipt.amount, receipt.receipt_seed, receipt.mint, receipt.deposited_at) }; + // Ensure the mint account matches the receipt's mint to prevent cross-mint withdrawals. + if receipt_mint != *ix.accounts.mint.address() { + return Err(ProgramError::InvalidAccountData); + } + // Validate extensions PDA validate_extensions_pda(ix.accounts.escrow, ix.accounts.extensions, program_id)?; @@ -68,7 +73,7 @@ pub fn process_withdraw(program_id: &Address, accounts: &[AccountView], instruct hook.invoke( HookPoint::PreWithdraw, remaining_accounts, - &[ix.accounts.escrow, ix.accounts.withdrawer, ix.accounts.mint, ix.accounts.receipt], + &[ix.accounts.escrow, ix.accounts.mint, ix.accounts.receipt], )?; } @@ -92,23 +97,23 @@ pub fn process_withdraw(program_id: &Address, accounts: &[AccountView], instruct })?; } - // Close receipt account and return lamports to rent_recipient - close_pda_account(ix.accounts.receipt, ix.accounts.rent_recipient)?; - - // Invoke post-withdraw hook if configured (receipt is closed, don't pass it) + // Invoke post-withdraw hook if configured (receipt is still open, pass it for context) if let Some(ref hook) = hook_data { hook.invoke( HookPoint::PostWithdraw, remaining_accounts, - &[ix.accounts.escrow, ix.accounts.withdrawer, ix.accounts.mint], + &[ix.accounts.escrow, ix.accounts.mint, ix.accounts.receipt], )?; } + // Close receipt account and return lamports to rent_recipient + close_pda_account(ix.accounts.receipt, ix.accounts.rent_recipient)?; + // Emit event let event = WithdrawEvent::new( *ix.accounts.escrow.address(), *ix.accounts.withdrawer.address(), - mint, + receipt_mint, receipt_seed, amount, ); diff --git a/program/src/state/allowed_mint.rs b/program/src/state/allowed_mint.rs index e938b31..e9de052 100644 --- a/program/src/state/allowed_mint.rs +++ b/program/src/state/allowed_mint.rs @@ -1,7 +1,7 @@ use alloc::vec; use alloc::vec::Vec; use codama::CodamaAccount; -use pinocchio::{cpi::Seed, error::ProgramError, Address}; +use pinocchio::{account::AccountView, cpi::Seed, error::ProgramError, Address}; use crate::assert_no_padding; use crate::traits::{ @@ -17,7 +17,7 @@ use crate::traits::{ /// # PDA Seeds /// `[b"allowed_mint", escrow.as_ref(), mint.as_ref()]` #[derive(Clone, Debug, PartialEq, CodamaAccount)] -#[codama(field("discriminator", number(u8), default_value = 3))] +#[codama(field("discriminator", number(u8), default_value = 4))] #[codama(discriminator(field = "discriminator"))] #[codama(seed(type = string(utf8), value = "allowed_mint"))] #[codama(seed(name = "escrow", type = public_key))] @@ -57,8 +57,17 @@ impl AllowedMint { } #[inline(always)] - pub fn from_account(data: &[u8]) -> Result<&Self, ProgramError> { - Self::from_bytes(data) + pub fn from_account<'a>( + data: &'a [u8], + account: &AccountView, + program_id: &Address, + escrow: &Address, + mint: &Address, + ) -> Result<&'a Self, ProgramError> { + let state = Self::from_bytes(data)?; + let pda = AllowedMintPda::new(escrow, mint); + pda.validate_pda(account, program_id, state.bump)?; + Ok(state) } } diff --git a/program/src/state/escrow.rs b/program/src/state/escrow.rs index 1fe5605..3899e48 100644 --- a/program/src/state/escrow.rs +++ b/program/src/state/escrow.rs @@ -20,7 +20,7 @@ use crate::traits::{ /// # PDA Seeds /// `[b"escrow", escrow_seed.as_ref()]` #[derive(Clone, Debug, PartialEq, CodamaAccount)] -#[codama(field("discriminator", number(u8), default_value = 0))] +#[codama(field("discriminator", number(u8), default_value = 1))] #[codama(discriminator(field = "discriminator"))] #[codama(seed(type = string(utf8), value = "escrow"))] #[codama(seed(name = "escrowSeed", type = public_key))] diff --git a/program/src/state/escrow_extensions.rs b/program/src/state/escrow_extensions.rs index 49479f8..4ea6fe7 100644 --- a/program/src/state/escrow_extensions.rs +++ b/program/src/state/escrow_extensions.rs @@ -46,7 +46,7 @@ pub const TLV_HEADER_SIZE: usize = 4; /// [discriminator: 1][version: 1][header: 2][TLV extensions: variable] /// ``` #[derive(Clone, Debug, PartialEq, CodamaAccount)] -#[codama(field("discriminator", number(u8), default_value = 1))] +#[codama(field("discriminator", number(u8), default_value = 2))] #[codama(discriminator(field = "discriminator"))] #[codama(pda = "extensions")] #[codama(seed(type = string(utf8), value = "extensions"))] @@ -86,6 +86,9 @@ impl EscrowExtensionsHeader { require_len!(data, Self::LEN); validate_discriminator!(data, Self::DISCRIMINATOR); + if data[1] != Self::VERSION { + return Err(ProgramError::InvalidAccountData); + } // Skip discriminator (byte 0) and version (byte 1) Ok(Self { bump: data[2], extension_count: data[3] }) @@ -309,6 +312,10 @@ pub fn validate_extensions_pda(escrow: &AccountView, extensions: &AccountView, p let expected_bump = extensions_pda.validate_pda_address(extensions, program_id)?; if extensions.data_len() > 0 { + if !extensions.owned_by(program_id) { + return Err(ProgramError::InvalidAccountOwner); + } + let data = extensions.try_borrow()?; let header = EscrowExtensionsHeader::from_bytes(&data)?; if header.bump != expected_bump { @@ -405,6 +412,16 @@ mod tests { assert_eq!(parsed.extension_count, header.extension_count); } + #[test] + fn test_header_from_bytes_wrong_version() { + let header = EscrowExtensionsHeader::new(100, 2); + let mut bytes = header.to_bytes(); + bytes[1] = EscrowExtensionsHeader::VERSION.wrapping_add(1); + + let result = EscrowExtensionsHeader::from_bytes(&bytes); + assert_eq!(result, Err(ProgramError::InvalidAccountData)); + } + #[test] fn test_calculate_extensions_account_size() { let no_extensions = calculate_extensions_account_size(false); diff --git a/program/src/state/extensions/hook.rs b/program/src/state/extensions/hook.rs index f7d0c0d..c3118a6 100644 --- a/program/src/state/extensions/hook.rs +++ b/program/src/state/extensions/hook.rs @@ -51,7 +51,7 @@ impl HookData { /// # Arguments /// * `hook_point` - The hook point discriminator /// * `remaining_accounts` - Remaining accounts slice: [hook_program, extra_accounts...] - /// * `core_accounts` - Core accounts to pass to hook (escrow, actor, mint, receipt, vault) + /// * `core_accounts` - Core accounts to pass to hook (escrow, mint, receipt) /// /// # Returns /// * `Ok(())` if hook succeeds @@ -65,6 +65,10 @@ impl HookData { self.validate(remaining_accounts)?; let extra_accounts = remaining_accounts.get(1..).unwrap_or(&[]); + if extra_accounts.iter().any(|acc| acc.is_signer()) { + return Err(EscrowProgramError::HookRejected.into()); + } + let all_accounts: Vec<&AccountView> = core_accounts.iter().copied().chain(extra_accounts.iter()).collect(); // Build instruction accounts - ALL accounts are read-only diff --git a/program/src/state/receipt.rs b/program/src/state/receipt.rs index 945bffd..82917a6 100644 --- a/program/src/state/receipt.rs +++ b/program/src/state/receipt.rs @@ -15,7 +15,7 @@ use crate::{assert_no_padding, require_account_len, validate_discriminator}; /// # PDA Seeds /// `[b"receipt", escrow.as_ref(), depositor.as_ref(), mint.as_ref(), receipt_seed.as_ref()]` #[derive(Clone, Debug, PartialEq, CodamaAccount)] -#[codama(field("discriminator", number(u8), default_value = 2))] +#[codama(field("discriminator", number(u8), default_value = 3))] #[codama(discriminator(field = "discriminator"))] #[codama(seed(type = string(utf8), value = "receipt"))] #[codama(seed(name = "escrow", type = public_key))] @@ -56,6 +56,9 @@ impl AccountParse for Receipt { fn parse_from_bytes(data: &[u8]) -> Result { require_account_len!(data, Self::LEN); validate_discriminator!(data, Self::DISCRIMINATOR); + if data[1] != Self::VERSION { + return Err(ProgramError::InvalidAccountData); + } // Skip discriminator (byte 0) and version (byte 1) let data = &data[2..]; @@ -267,4 +270,14 @@ mod tests { let result = Receipt::parse_from_bytes(&bytes); assert!(result.is_err()); } + + #[test] + fn test_receipt_parse_from_bytes_wrong_version() { + let receipt = create_test_receipt(); + let mut bytes = receipt.to_bytes(); + bytes[1] = Receipt::VERSION.wrapping_add(1); + + let result = Receipt::parse_from_bytes(&bytes); + assert_eq!(result, Err(ProgramError::InvalidAccountData)); + } } diff --git a/program/src/traits/account.rs b/program/src/traits/account.rs index f486938..d80caa0 100644 --- a/program/src/traits/account.rs +++ b/program/src/traits/account.rs @@ -29,6 +29,9 @@ pub trait AccountDeserialize: AccountSize { fn from_bytes(data: &[u8]) -> Result<&Self, ProgramError> { require_len!(data, Self::LEN); validate_discriminator!(data, Self::DISCRIMINATOR); + if data[1] != Self::VERSION { + return Err(ProgramError::InvalidAccountData); + } // Skip discriminator (byte 0) and version (byte 1) unsafe { Self::from_bytes_unchecked(&data[2..]) } @@ -52,6 +55,9 @@ pub trait AccountDeserialize: AccountSize { fn from_bytes_mut(data: &mut [u8]) -> Result<&mut Self, ProgramError> { require_len!(data, Self::LEN); validate_discriminator!(data, Self::DISCRIMINATOR); + if data[1] != Self::VERSION { + return Err(ProgramError::InvalidAccountData); + } // Skip discriminator (byte 0) and version (byte 1) unsafe { Self::from_bytes_mut_unchecked(&mut data[2..]) } @@ -74,10 +80,10 @@ pub trait AccountDeserialize: AccountSize { /// Account discriminator values for this program #[repr(u8)] pub enum EscrowAccountDiscriminators { - EscrowDiscriminator = 0, - EscrowExtensionsDiscriminator = 1, - ReceiptDiscriminator = 2, - AllowedMintDiscriminator = 3, + EscrowDiscriminator = 1, + EscrowExtensionsDiscriminator = 2, + ReceiptDiscriminator = 3, + AllowedMintDiscriminator = 4, } /// Manual account deserialization (non-zero-copy) @@ -178,6 +184,30 @@ mod tests { assert_eq!(deserialized.admin, escrow.admin); } + #[test] + fn test_from_bytes_wrong_version() { + let escrow_seed = Address::new_from_array([1u8; 32]); + let admin = Address::new_from_array([2u8; 32]); + let escrow = Escrow::new(100, escrow_seed, admin); + let mut bytes = escrow.to_bytes(); + bytes[1] = Escrow::VERSION.wrapping_add(1); + + let result = Escrow::from_bytes(&bytes); + assert_eq!(result, Err(ProgramError::InvalidAccountData)); + } + + #[test] + fn test_from_bytes_mut_wrong_version() { + let escrow_seed = Address::new_from_array([1u8; 32]); + let admin = Address::new_from_array([2u8; 32]); + let escrow = Escrow::new(100, escrow_seed, admin); + let mut bytes = escrow.to_bytes(); + bytes[1] = Escrow::VERSION.wrapping_add(1); + + let result = Escrow::from_bytes_mut(&mut bytes); + assert_eq!(result, Err(ProgramError::InvalidAccountData)); + } + #[test] fn test_write_to_slice_exact_size() { let escrow_seed = Address::new_from_array([1u8; 32]); @@ -203,4 +233,12 @@ mod tests { assert_eq!(bytes[0], Escrow::DISCRIMINATOR); assert_eq!(bytes[1], Escrow::VERSION); } + + #[test] + fn test_account_discriminators_are_non_zero_and_stable() { + assert_eq!(EscrowAccountDiscriminators::EscrowDiscriminator as u8, 1); + assert_eq!(EscrowAccountDiscriminators::EscrowExtensionsDiscriminator as u8, 2); + assert_eq!(EscrowAccountDiscriminators::ReceiptDiscriminator as u8, 3); + assert_eq!(EscrowAccountDiscriminators::AllowedMintDiscriminator as u8, 4); + } } diff --git a/program/src/utils/pda_utils.rs b/program/src/utils/pda_utils.rs index 610fb37..9ba15c8 100644 --- a/program/src/utils/pda_utils.rs +++ b/program/src/utils/pda_utils.rs @@ -18,7 +18,11 @@ pub fn close_pda_account(pda_account: &AccountView, recipient: &AccountView) -> /// Create a PDA account for the given seeds. /// -/// Will return an error if the account already exists (has lamports). +/// Strict create-once semantics: +/// - If account has data: returns `AccountAlreadyInitialized` +/// - If account has lamports but no data: completes initialization via +/// transfer (if needed) + allocate + assign +/// - If account is fully absent (0 lamports): uses `CreateAccount` pub fn create_pda_account( payer: &AccountView, space: usize, @@ -33,7 +37,18 @@ pub fn create_pda_account( let signers = [Signer::from(&pda_signer_seeds)]; if pda_account.lamports() > 0 { - Err(ProgramError::AccountAlreadyInitialized) + if pda_account.data_len() > 0 { + return Err(ProgramError::AccountAlreadyInitialized); + } + + // PDA was prefunded but not initialized yet. + let additional_lamports = required_lamports.saturating_sub(pda_account.lamports()); + if additional_lamports > 0 { + Transfer { from: payer, to: pda_account, lamports: additional_lamports }.invoke()?; + } + + Allocate { account: pda_account, space: space as u64 }.invoke_signed(&signers)?; + Assign { account: pda_account, owner }.invoke_signed(&signers) } else { CreateAccount { from: payer, to: pda_account, lamports: required_lamports, space: space as u64, owner } .invoke_signed(&signers) diff --git a/program/src/utils/token2022_utils.rs b/program/src/utils/token2022_utils.rs index 627091f..235e663 100644 --- a/program/src/utils/token2022_utils.rs +++ b/program/src/utils/token2022_utils.rs @@ -1,6 +1,6 @@ //! Token2022 extension validation utilities. -use alloc::vec::Vec; +use alloc::vec; use pinocchio::{account::AccountView, error::ProgramError}; use pinocchio_token_2022::ID as TOKEN_2022_PROGRAM_ID; use spl_token_2022::{ @@ -19,6 +19,8 @@ use crate::{errors::EscrowProgramError, utils::TlvReader}; /// - `PermanentDelegate`: Authority can transfer/burn tokens from ANY account /// - `NonTransferable`: Tokens cannot be transferred /// - `Pausable`: Authority can pause all transfers +/// - `TransferFeeConfig`: Transfer fees break escrow accounting invariants +/// - `MintCloseAuthority`: Mint can be closed and recreated with unsafe configuration /// /// # Escrow-Specific Blocklist /// If `extensions` account is provided and contains a `BlockTokenExtensions` extension, @@ -46,12 +48,14 @@ pub fn validate_mint_extensions(mint: &AccountView, extensions: &AccountView) -> let extension_types = mint_state.get_extension_types()?; // Build combined blocklist: global + escrow-specific (as u16 values) - let mut blocked_types_u16 = Vec::new(); - // Add global blocklist (convert ExtensionType enum to u16) - blocked_types_u16.push(ExtensionType::PermanentDelegate as u16); - blocked_types_u16.push(ExtensionType::NonTransferable as u16); - blocked_types_u16.push(ExtensionType::Pausable as u16); + let mut blocked_types_u16 = vec![ + ExtensionType::PermanentDelegate as u16, + ExtensionType::NonTransferable as u16, + ExtensionType::Pausable as u16, + ExtensionType::TransferFeeConfig as u16, + ExtensionType::MintCloseAuthority as u16, + ]; // Add escrow-specific blocklist if extensions account has data if extensions.data_len() > 0 { @@ -78,6 +82,12 @@ pub fn validate_mint_extensions(mint: &AccountView, extensions: &AccountView) -> ExtensionType::Pausable => { return Err(EscrowProgramError::PausableNotAllowed.into()); } + ExtensionType::TransferFeeConfig => { + return Err(EscrowProgramError::MintNotAllowed.into()); + } + ExtensionType::MintCloseAuthority => { + return Err(EscrowProgramError::MintNotAllowed.into()); + } _ => { // Escrow-specific blocked extension return Err(EscrowProgramError::MintNotAllowed.into()); diff --git a/tests/integration-tests/src/fixtures/deposit.rs b/tests/integration-tests/src/fixtures/deposit.rs index 615180b..6a71485 100644 --- a/tests/integration-tests/src/fixtures/deposit.rs +++ b/tests/integration-tests/src/fixtures/deposit.rs @@ -5,6 +5,7 @@ use solana_sdk::{ pubkey::Pubkey, signature::{Keypair, Signer}, }; +use spl_token_2022::extension::ExtensionType; use spl_token_2022::ID as TOKEN_2022_PROGRAM_ID; use spl_token_interface::ID as TOKEN_PROGRAM_ID; @@ -85,11 +86,12 @@ pub struct DepositSetupBuilder<'a> { ctx: &'a mut TestContext, token_program: Pubkey, hook_program: Option, + mint_extension: Option, } impl<'a> DepositSetupBuilder<'a> { fn new(ctx: &'a mut TestContext) -> Self { - Self { ctx, token_program: TOKEN_PROGRAM_ID, hook_program: None } + Self { ctx, token_program: TOKEN_PROGRAM_ID, hook_program: None, mint_extension: None } } pub fn token_2022(mut self) -> Self { @@ -107,6 +109,12 @@ impl<'a> DepositSetupBuilder<'a> { self } + pub fn mint_extension(mut self, extension: ExtensionType) -> Self { + self.mint_extension = Some(extension); + self.token_program = TOKEN_2022_PROGRAM_ID; + self + } + pub fn build(self) -> DepositSetup { let admin = self.ctx.create_funded_keypair(); let escrow_seed = Keypair::new(); @@ -140,12 +148,19 @@ impl<'a> DepositSetupBuilder<'a> { let token_program = self.token_program; let (vault, depositor_token_account); - if token_program == TOKEN_2022_PROGRAM_ID { - self.ctx.create_token_2022_mint(&mint, &self.ctx.payer.pubkey(), 6); - vault = self.ctx.create_token_2022_account(&escrow_pda, &mint.pubkey()); - } else { - self.ctx.create_mint(&mint, &self.ctx.payer.pubkey(), 6); - vault = self.ctx.create_token_account(&escrow_pda, &mint.pubkey()); + match self.mint_extension { + Some(extension) => { + self.ctx.create_token_2022_mint_with_extension(&mint, &self.ctx.payer.pubkey(), 6, extension); + vault = self.ctx.create_token_2022_account(&escrow_pda, &mint.pubkey()); + } + None if token_program == TOKEN_2022_PROGRAM_ID => { + self.ctx.create_token_2022_mint(&mint, &self.ctx.payer.pubkey(), 6); + vault = self.ctx.create_token_2022_account(&escrow_pda, &mint.pubkey()); + } + None => { + self.ctx.create_mint(&mint, &self.ctx.payer.pubkey(), 6); + vault = self.ctx.create_token_account(&escrow_pda, &mint.pubkey()); + } } let (allowed_mint_pda, allowed_mint_bump) = find_allowed_mint_pda(&escrow_pda, &mint.pubkey()); diff --git a/tests/integration-tests/src/test_add_timelock.rs b/tests/integration-tests/src/test_add_timelock.rs index ccfdefd..d61ec95 100644 --- a/tests/integration-tests/src/test_add_timelock.rs +++ b/tests/integration-tests/src/test_add_timelock.rs @@ -95,7 +95,7 @@ fn test_add_timelock_escrow_not_owned_by_program() { } #[test] -fn test_add_timelock_duplicate_extension() { +fn test_add_timelock_updates_existing_extension() { let mut ctx = TestContext::new(); let escrow_ix = CreateEscrowFixture::build_valid(&mut ctx); @@ -104,13 +104,19 @@ fn test_add_timelock_duplicate_extension() { escrow_ix.send_expect_success(&mut ctx); let (escrow_pda, _) = find_escrow_pda(&escrow_seed); + let (extensions_pda, extensions_bump) = find_extensions_pda(&escrow_pda); let first_ix = AddTimelockFixture::build_with_escrow(&mut ctx, escrow_pda, admin.insecure_clone(), 3600); first_ix.send_expect_success(&mut ctx); + assert_extensions_header(&ctx, &extensions_pda, extensions_bump, 1); + assert_timelock_extension(&ctx, &extensions_pda, 3600); let second_ix = AddTimelockFixture::build_with_escrow(&mut ctx, escrow_pda, admin, 7200); - let error = second_ix.send_expect_error(&mut ctx); - assert_instruction_error(error, InstructionError::AccountAlreadyInitialized); + second_ix.send_expect_success(&mut ctx); + + // Updating timelock should replace value in place without increasing count. + assert_extensions_header(&ctx, &extensions_pda, extensions_bump, 1); + assert_timelock_extension(&ctx, &extensions_pda, 7200); } // ============================================================================ diff --git a/tests/integration-tests/src/test_allow_mint.rs b/tests/integration-tests/src/test_allow_mint.rs index a3634ef..d1c1719 100644 --- a/tests/integration-tests/src/test_allow_mint.rs +++ b/tests/integration-tests/src/test_allow_mint.rs @@ -7,7 +7,7 @@ use crate::{ }, }; use escrow_program_client::instructions::AllowMintBuilder; -use solana_sdk::{instruction::InstructionError, signature::Signer}; +use solana_sdk::{account::Account, instruction::InstructionError, pubkey::Pubkey, signature::Signer}; use spl_associated_token_account::get_associated_token_address; use spl_token_2022::extension::ExtensionType; @@ -147,6 +147,26 @@ fn test_allow_mint_success() { assert_allowed_mint_account(&ctx, &setup.allowed_mint_pda, setup.allowed_mint_bump); } +#[test] +fn test_allow_mint_prefunded_allowed_mint_pda_succeeds() { + let mut ctx = TestContext::new(); + let setup = AllowMintSetup::new(&mut ctx); + + // Simulate griefing by pre-funding the AllowedMint PDA before initialization. + ctx.svm + .set_account( + setup.allowed_mint_pda, + Account { lamports: 1, data: vec![], owner: Pubkey::default(), executable: false, rent_epoch: 0 }, + ) + .unwrap(); + + let test_ix = setup.build_instruction(&ctx); + test_ix.send_expect_success(&mut ctx); + + assert_account_exists(&ctx, &setup.allowed_mint_pda); + assert_allowed_mint_account(&ctx, &setup.allowed_mint_pda, setup.allowed_mint_bump); +} + #[test] fn test_allow_mint_multiple_mints() { let mut ctx = TestContext::new(); @@ -234,6 +254,28 @@ fn test_allow_mint_rejects_pausable() { assert_escrow_error(error, EscrowError::PausableNotAllowed); } +#[test] +fn test_allow_mint_rejects_transfer_fee_config() { + let mut ctx = TestContext::new(); + let setup = AllowMintSetup::new_with_extension(&mut ctx, ExtensionType::TransferFeeConfig); + + let test_ix = setup.build_instruction(&ctx); + let error = test_ix.send_expect_error(&mut ctx); + + assert_escrow_error(error, EscrowError::MintNotAllowed); +} + +#[test] +fn test_allow_mint_rejects_mint_close_authority() { + let mut ctx = TestContext::new(); + let setup = AllowMintSetup::new_with_extension(&mut ctx, ExtensionType::MintCloseAuthority); + + let test_ix = setup.build_instruction(&ctx); + let error = test_ix.send_expect_error(&mut ctx); + + assert_escrow_error(error, EscrowError::MintNotAllowed); +} + // ============================================================================ // Escrow-Specific Blocked Extension Tests // ============================================================================ @@ -242,8 +284,9 @@ fn test_allow_mint_rejects_pausable() { fn test_allow_mint_rejects_escrow_blocked_extension() { let mut ctx = TestContext::new(); - // Create escrow with TransferFeeConfig blocked, then try to allow a mint with that extension - let setup = AllowMintSetup::new_with_escrow_blocked_extension(&mut ctx, ExtensionType::TransferFeeConfig); + // Create escrow with MetadataPointer blocked, then try to allow a mint with that extension. + // MetadataPointer is not globally blocked, so rejection should come from escrow-specific blocklist. + let setup = AllowMintSetup::new_with_escrow_blocked_extension(&mut ctx, ExtensionType::MetadataPointer); let test_ix = setup.build_instruction(&ctx); let error = test_ix.send_expect_error(&mut ctx); @@ -255,12 +298,8 @@ fn test_allow_mint_rejects_escrow_blocked_extension() { fn test_allow_mint_accepts_mint_without_escrow_blocked_extension() { let mut ctx = TestContext::new(); - // Block TransferFeeConfig, but create a mint with MintCloseAuthority (different extension) - let setup = AllowMintSetup::new_with_different_extension_blocked( - &mut ctx, - ExtensionType::TransferFeeConfig, - ExtensionType::MintCloseAuthority, - ); + // Block MetadataPointer, but create a Token-2022 mint without that extension. + let setup = AllowMintSetup::builder(&mut ctx).block_extension(ExtensionType::MetadataPointer).token_2022().build(); let test_ix = setup.build_instruction(&ctx); test_ix.send_expect_success(&mut ctx); diff --git a/tests/integration-tests/src/test_create_escrow.rs b/tests/integration-tests/src/test_create_escrow.rs index 772171e..5a1afd9 100644 --- a/tests/integration-tests/src/test_create_escrow.rs +++ b/tests/integration-tests/src/test_create_escrow.rs @@ -6,7 +6,7 @@ use crate::{ }, }; use escrow_program_client::instructions::CreatesEscrowBuilder; -use solana_sdk::{instruction::InstructionError, signature::Signer}; +use solana_sdk::{account::Account, instruction::InstructionError, pubkey::Pubkey, signature::Signer}; // ============================================================================ // Error Tests - Using Generic Test Helpers @@ -94,6 +94,29 @@ fn test_create_escrow_success() { assert_escrow_account(&ctx, &escrow_pda, &admin_pubkey, bump, &escrow_seed_pubkey); } +#[test] +fn test_create_escrow_prefunded_pda_succeeds() { + let mut ctx = TestContext::new(); + let test_ix = CreateEscrowFixture::build_valid(&mut ctx); + + let admin_pubkey = test_ix.signers[0].pubkey(); + let escrow_seed_pubkey = test_ix.signers[1].pubkey(); + let escrow_pda = test_ix.instruction.accounts[3].pubkey; + let bump = test_ix.instruction.data[1]; + + // Simulate griefing by pre-funding the PDA before initialization. + ctx.svm + .set_account( + escrow_pda, + Account { lamports: 1, data: vec![], owner: Pubkey::default(), executable: false, rent_epoch: 0 }, + ) + .unwrap(); + + test_ix.send_expect_success(&mut ctx); + + assert_escrow_account(&ctx, &escrow_pda, &admin_pubkey, bump, &escrow_seed_pubkey); +} + // ============================================================================ // Re-initialization Protection Tests // ============================================================================ diff --git a/tests/integration-tests/src/test_deposit.rs b/tests/integration-tests/src/test_deposit.rs index e6c2ff5..236268b 100644 --- a/tests/integration-tests/src/test_deposit.rs +++ b/tests/integration-tests/src/test_deposit.rs @@ -1,5 +1,5 @@ use crate::{ - fixtures::{DepositFixture, DepositSetup, DEFAULT_DEPOSIT_AMOUNT}, + fixtures::{AddBlockTokenExtensionsFixture, DepositFixture, DepositSetup, DEFAULT_DEPOSIT_AMOUNT}, utils::{ assert_custom_error, assert_escrow_error, assert_instruction_error, find_receipt_pda, test_empty_data, test_missing_signer, test_not_writable, test_wrong_account, test_wrong_current_program, test_wrong_owner, @@ -7,8 +7,16 @@ use crate::{ TEST_HOOK_ALLOW_ID, TEST_HOOK_DENY_ERROR, TEST_HOOK_DENY_ID, }, }; +use escrow_program_client::instructions::AddTimelockBuilder; use escrow_program_client::instructions::DepositBuilder; -use solana_sdk::{instruction::InstructionError, pubkey::Pubkey, signature::Signer}; +use solana_sdk::{ + account::Account, + instruction::{AccountMeta, InstructionError}, + pubkey::Pubkey, + signature::Signer, +}; +use spl_token_2022::extension::ExtensionType; +use spl_token_2022::ID as TOKEN_2022_PROGRAM_ID; // ============================================================================ // Error Tests - Using Generic Test Helpers @@ -80,6 +88,20 @@ fn test_deposit_wrong_token_program() { test_wrong_token_program::(&mut ctx, 9); } +#[test] +fn test_deposit_mint_owner_mismatch_token_program() { + let mut ctx = TestContext::new(); + let setup = DepositSetup::new(&mut ctx); + + let mut mint_account = ctx.get_account(&setup.mint.pubkey()).expect("Mint account should exist"); + mint_account.owner = TOKEN_2022_PROGRAM_ID; + ctx.svm.set_account(setup.mint.pubkey(), mint_account).unwrap(); + + let test_ix = setup.build_instruction(&ctx); + let error = test_ix.send_expect_error(&mut ctx); + assert_instruction_error(error, InstructionError::InvalidAccountOwner); +} + #[test] fn test_deposit_wrong_escrow_owner() { let mut ctx = TestContext::new(); @@ -92,6 +114,69 @@ fn test_deposit_wrong_allowed_mint_owner() { test_wrong_owner::(&mut ctx, 3); } +#[test] +fn test_deposit_initialized_extensions_wrong_owner() { + let mut ctx = TestContext::new(); + let setup = DepositSetup::new(&mut ctx); + + let (extensions_pda, extensions_bump) = crate::utils::find_extensions_pda(&setup.escrow_pda); + let add_timelock_ix = AddTimelockBuilder::new() + .payer(ctx.payer.pubkey()) + .admin(setup.admin.pubkey()) + .escrow(setup.escrow_pda) + .extensions(extensions_pda) + .extensions_bump(extensions_bump) + .lock_duration(1) + .instruction(); + ctx.send_transaction(add_timelock_ix, &[&setup.admin]).unwrap(); + + let mut extensions_account = ctx.get_account(&setup.extensions_pda).expect("Extensions account should exist"); + extensions_account.owner = Pubkey::new_unique(); + ctx.svm.set_account(setup.extensions_pda, extensions_account).unwrap(); + + let test_ix = setup.build_instruction(&ctx); + let error = test_ix.send_expect_error(&mut ctx); + assert_instruction_error(error, InstructionError::InvalidAccountOwner); +} + +#[test] +fn test_deposit_rejects_newly_blocked_mint_extension() { + let mut ctx = TestContext::new(); + let setup = DepositSetup::builder(&mut ctx).mint_extension(ExtensionType::MetadataPointer).build(); + + let block_extension_ix = AddBlockTokenExtensionsFixture::build_with_escrow( + &mut ctx, + setup.escrow_pda, + setup.admin.insecure_clone(), + ExtensionType::MetadataPointer as u16, + ); + block_extension_ix.send_expect_success(&mut ctx); + + let test_ix = setup.build_instruction(&ctx); + let error = test_ix.send_expect_error(&mut ctx); + assert_escrow_error(error, EscrowError::MintNotAllowed); +} + +#[test] +fn test_deposit_prefunded_receipt_pda_succeeds() { + let mut ctx = TestContext::new(); + let setup = DepositSetup::new(&mut ctx); + + // Simulate griefing by pre-funding the receipt PDA before initialization. + ctx.svm + .set_account( + setup.receipt_pda, + Account { lamports: 1, data: vec![], owner: Pubkey::default(), executable: false, rent_epoch: 0 }, + ) + .unwrap(); + + let test_ix = setup.build_instruction(&ctx); + test_ix.send_expect_success(&mut ctx); + + let receipt_account = ctx.get_account(&setup.receipt_pda).expect("Deposit receipt should exist"); + assert!(!receipt_account.data.is_empty()); +} + #[test] fn test_deposit_zero_amount() { let mut ctx = TestContext::new(); @@ -337,6 +422,20 @@ fn test_deposit_with_hook_rejected() { assert_custom_error(error, TEST_HOOK_DENY_ERROR); } +#[test] +fn test_deposit_with_hook_extra_signer_rejected() { + let mut ctx = TestContext::new(); + let setup = DepositSetup::new_with_hook(&mut ctx, TEST_HOOK_ALLOW_ID); + let extra_signer = ctx.create_funded_keypair(); + + let mut test_ix = setup.build_instruction(&ctx); + test_ix.instruction.accounts.push(AccountMeta::new_readonly(extra_signer.pubkey(), true)); + test_ix.signers.push(extra_signer.insecure_clone()); + + let error = test_ix.send_expect_error(&mut ctx); + assert_escrow_error(error, EscrowError::HookRejected); +} + // ============================================================================ // Additional Tests // ============================================================================ diff --git a/tests/integration-tests/src/test_set_arbiter.rs b/tests/integration-tests/src/test_set_arbiter.rs index 382be3a..aad3cf2 100644 --- a/tests/integration-tests/src/test_set_arbiter.rs +++ b/tests/integration-tests/src/test_set_arbiter.rs @@ -106,7 +106,7 @@ fn test_set_arbiter_escrow_not_owned_by_program() { } #[test] -fn test_set_arbiter_duplicate_extension() { +fn test_set_arbiter_updates_existing_extension() { let mut ctx = TestContext::new(); let escrow_ix = CreateEscrowFixture::build_valid(&mut ctx); @@ -115,15 +115,23 @@ fn test_set_arbiter_duplicate_extension() { escrow_ix.send_expect_success(&mut ctx); let (escrow_pda, _) = find_escrow_pda(&escrow_seed); - let arbiter = Keypair::new(); + let (extensions_pda, extensions_bump) = find_extensions_pda(&escrow_pda); + let first_arbiter = Keypair::new(); + let first_arbiter_pubkey = first_arbiter.pubkey(); - let first_ix = SetArbiterFixture::build_with_escrow(&mut ctx, escrow_pda, admin.insecure_clone(), arbiter); + let first_ix = SetArbiterFixture::build_with_escrow(&mut ctx, escrow_pda, admin.insecure_clone(), first_arbiter); first_ix.send_expect_success(&mut ctx); + assert_extensions_header(&ctx, &extensions_pda, extensions_bump, 1); + assert_arbiter_extension(&ctx, &extensions_pda, &first_arbiter_pubkey); - // Second attempt should fail — arbiter is immutable - let second_ix = SetArbiterFixture::build_with_escrow(&mut ctx, escrow_pda, admin, Keypair::new()); - let error = second_ix.send_expect_error(&mut ctx); - assert_instruction_error(error, InstructionError::AccountAlreadyInitialized); + let second_arbiter = Keypair::new(); + let second_arbiter_pubkey = second_arbiter.pubkey(); + let second_ix = SetArbiterFixture::build_with_escrow(&mut ctx, escrow_pda, admin, second_arbiter); + second_ix.send_expect_success(&mut ctx); + + // Updating arbiter should replace value in place without increasing count. + assert_extensions_header(&ctx, &extensions_pda, extensions_bump, 1); + assert_arbiter_extension(&ctx, &extensions_pda, &second_arbiter_pubkey); } // ============================================================================ diff --git a/tests/integration-tests/src/test_set_hook.rs b/tests/integration-tests/src/test_set_hook.rs index 34889c1..98cc556 100644 --- a/tests/integration-tests/src/test_set_hook.rs +++ b/tests/integration-tests/src/test_set_hook.rs @@ -98,7 +98,7 @@ fn test_set_hook_escrow_not_owned_by_program() { } #[test] -fn test_set_hook_duplicate_extension() { +fn test_set_hook_updates_existing_extension() { let mut ctx = TestContext::new(); let escrow_ix = CreateEscrowFixture::build_valid(&mut ctx); @@ -107,14 +107,21 @@ fn test_set_hook_duplicate_extension() { escrow_ix.send_expect_success(&mut ctx); let (escrow_pda, _) = find_escrow_pda(&escrow_seed); - let hook_program = Pubkey::new_unique(); + let (extensions_pda, extensions_bump) = find_extensions_pda(&escrow_pda); + let first_hook_program = Pubkey::new_unique(); - let first_ix = SetHookFixture::build_with_escrow(&mut ctx, escrow_pda, admin.insecure_clone(), hook_program); + let first_ix = SetHookFixture::build_with_escrow(&mut ctx, escrow_pda, admin.insecure_clone(), first_hook_program); first_ix.send_expect_success(&mut ctx); + assert_extensions_header(&ctx, &extensions_pda, extensions_bump, 1); + assert_hook_extension(&ctx, &extensions_pda, &first_hook_program); - let second_ix = SetHookFixture::build_with_escrow(&mut ctx, escrow_pda, admin, Pubkey::new_unique()); - let error = second_ix.send_expect_error(&mut ctx); - assert_instruction_error(error, InstructionError::AccountAlreadyInitialized); + let second_hook_program = Pubkey::new_unique(); + let second_ix = SetHookFixture::build_with_escrow(&mut ctx, escrow_pda, admin, second_hook_program); + second_ix.send_expect_success(&mut ctx); + + // Updating hook should replace value in place without increasing count. + assert_extensions_header(&ctx, &extensions_pda, extensions_bump, 1); + assert_hook_extension(&ctx, &extensions_pda, &second_hook_program); } // ============================================================================ diff --git a/tests/integration-tests/src/test_withdraw.rs b/tests/integration-tests/src/test_withdraw.rs index 6e6a55e..fb734d4 100644 --- a/tests/integration-tests/src/test_withdraw.rs +++ b/tests/integration-tests/src/test_withdraw.rs @@ -1,19 +1,22 @@ use crate::{ fixtures::{AllowMintSetup, WithdrawFixture, WithdrawSetup, DEFAULT_DEPOSIT_AMOUNT}, utils::{ - assert_custom_error, assert_escrow_error, assert_instruction_error, test_missing_signer, test_not_writable, - test_wrong_account, test_wrong_current_program, test_wrong_owner, test_wrong_system_program, + assert_custom_error, assert_escrow_error, assert_instruction_error, find_allowed_mint_pda, test_missing_signer, + test_not_writable, test_wrong_account, test_wrong_current_program, test_wrong_owner, test_wrong_system_program, test_wrong_token_program, EscrowError, TestContext, TestInstruction, TEST_HOOK_ALLOW_ID, TEST_HOOK_DENY_ERROR, TEST_HOOK_DENY_ID, }, }; +use escrow_program_client::instructions::AddTimelockBuilder; +use escrow_program_client::instructions::AllowMintBuilder; use escrow_program_client::instructions::WithdrawBuilder; use solana_sdk::{ account::Account, instruction::{AccountMeta, InstructionError}, pubkey::Pubkey, - signature::Signer, + signature::{Keypair, Signer}, }; +use spl_token_2022::ID as TOKEN_2022_PROGRAM_ID; // ============================================================================ // Error Tests - Using Generic Test Helpers @@ -67,6 +70,20 @@ fn test_withdraw_wrong_token_program() { test_wrong_token_program::(&mut ctx, 8); } +#[test] +fn test_withdraw_mint_owner_mismatch_token_program() { + let mut ctx = TestContext::new(); + let setup = WithdrawSetup::new(&mut ctx); + + let mut mint_account = ctx.get_account(&setup.mint.pubkey()).expect("Mint account should exist"); + mint_account.owner = TOKEN_2022_PROGRAM_ID; + ctx.svm.set_account(setup.mint.pubkey(), mint_account).unwrap(); + + let test_ix = setup.build_instruction(&ctx); + let error = test_ix.send_expect_error(&mut ctx); + assert_instruction_error(error, InstructionError::InvalidAccountOwner); +} + #[test] fn test_withdraw_wrong_escrow_owner() { let mut ctx = TestContext::new(); @@ -79,6 +96,31 @@ fn test_withdraw_wrong_receipt_owner() { test_wrong_owner::(&mut ctx, 4); } +#[test] +fn test_withdraw_initialized_extensions_wrong_owner() { + let mut ctx = TestContext::new(); + let setup = WithdrawSetup::new(&mut ctx); + + let (extensions_pda, extensions_bump) = crate::utils::find_extensions_pda(&setup.escrow_pda); + let add_timelock_ix = AddTimelockBuilder::new() + .payer(ctx.payer.pubkey()) + .admin(setup.admin.pubkey()) + .escrow(setup.escrow_pda) + .extensions(extensions_pda) + .extensions_bump(extensions_bump) + .lock_duration(1) + .instruction(); + ctx.send_transaction(add_timelock_ix, &[&setup.admin]).unwrap(); + + let mut extensions_account = ctx.get_account(&setup.extensions_pda).expect("Extensions account should exist"); + extensions_account.owner = Pubkey::new_unique(); + ctx.svm.set_account(setup.extensions_pda, extensions_account).unwrap(); + + let test_ix = setup.build_instruction(&ctx); + let error = test_ix.send_expect_error(&mut ctx); + assert_instruction_error(error, InstructionError::InvalidAccountOwner); +} + #[test] fn test_withdraw_wrong_extensions_account() { let mut ctx = TestContext::new(); @@ -407,6 +449,21 @@ fn test_withdraw_with_hook_rejected() { assert_eq!(final_vault_balance, initial_vault_balance, "Vault balance should be unchanged"); } +#[test] +fn test_withdraw_with_hook_extra_signer_rejected() { + let mut ctx = TestContext::new(); + let setup = WithdrawSetup::new_with_hook(&mut ctx, TEST_HOOK_ALLOW_ID); + let extra_signer = ctx.create_funded_keypair(); + + let mut test_ix = setup.build_instruction(&ctx); + test_ix.instruction.accounts.push(AccountMeta::new_readonly(extra_signer.pubkey(), true)); + test_ix.signers.push(extra_signer.insecure_clone()); + + let error = test_ix.send_expect_error(&mut ctx); + assert_escrow_error(error, EscrowError::HookRejected); + assert!(ctx.get_account(&setup.receipt_pda).is_some(), "Receipt should remain after rejected hook"); +} + // ============================================================================ // Cross-Escrow Protection Tests // ============================================================================ @@ -440,6 +497,49 @@ fn test_withdraw_receipt_for_different_escrow_fails() { assert_escrow_error(error, EscrowError::InvalidReceiptEscrow); } +#[test] +fn test_withdraw_receipt_mint_mismatch_fails() { + let mut ctx = TestContext::new(); + let setup = WithdrawSetup::new(&mut ctx); + + let second_mint = Keypair::new(); + ctx.create_mint(&second_mint, &ctx.payer.pubkey(), 6); + let second_vault = + ctx.create_token_account_with_balance(&setup.escrow_pda, &second_mint.pubkey(), DEFAULT_DEPOSIT_AMOUNT); + let second_withdrawer_token_account = ctx.create_token_account(&setup.depositor.pubkey(), &second_mint.pubkey()); + + let (second_allowed_mint, second_allowed_mint_bump) = + find_allowed_mint_pda(&setup.escrow_pda, &second_mint.pubkey()); + let allow_second_mint_ix = AllowMintBuilder::new() + .payer(ctx.payer.pubkey()) + .admin(setup.admin.pubkey()) + .escrow(setup.escrow_pda) + .escrow_extensions(setup.extensions_pda) + .mint(second_mint.pubkey()) + .allowed_mint(second_allowed_mint) + .vault(second_vault) + .token_program(setup.token_program) + .bump(second_allowed_mint_bump) + .instruction(); + ctx.send_transaction(allow_second_mint_ix, &[&setup.admin]).unwrap(); + + let instruction = WithdrawBuilder::new() + .rent_recipient(ctx.payer.pubkey()) + .withdrawer(setup.depositor.pubkey()) + .escrow(setup.escrow_pda) + .extensions(setup.extensions_pda) + .receipt(setup.receipt_pda) + .vault(second_vault) + .withdrawer_token_account(second_withdrawer_token_account) + .mint(second_mint.pubkey()) + .token_program(setup.token_program) + .instruction(); + + let error = ctx.send_transaction_expect_error(instruction, &[&setup.depositor]); + assert_instruction_error(error, InstructionError::InvalidAccountData); + assert!(ctx.get_account(&setup.receipt_pda).is_some(), "Receipt should remain after rejected withdraw"); +} + // ============================================================================ // Closed Account Protection Tests // ============================================================================ diff --git a/tests/integration-tests/src/utils/extensions_utils.rs b/tests/integration-tests/src/utils/extensions_utils.rs index ca221dc..9b6d956 100644 --- a/tests/integration-tests/src/utils/extensions_utils.rs +++ b/tests/integration-tests/src/utils/extensions_utils.rs @@ -5,7 +5,7 @@ pub const EXTENSION_TYPE_HOOK: u16 = 1; pub const EXTENSION_TYPE_BLOCK_TOKEN_EXTENSIONS: u16 = 2; pub const EXTENSION_TYPE_ARBITER: u16 = 3; -pub const ESCROW_EXTENSIONS_DISCRIMINATOR: u8 = 1; +pub const ESCROW_EXTENSIONS_DISCRIMINATOR: u8 = 2; pub const ESCROW_EXTENSIONS_HEADER_LEN: usize = 4; // discriminator + bump + version + extension_count pub const TIMELOCK_DATA_LEN: usize = 8; diff --git a/tests/integration-tests/src/utils/token_utils.rs b/tests/integration-tests/src/utils/token_utils.rs index de1b4ac..8a467d7 100644 --- a/tests/integration-tests/src/utils/token_utils.rs +++ b/tests/integration-tests/src/utils/token_utils.rs @@ -8,9 +8,9 @@ use solana_sdk::{ use spl_associated_token_account::{get_associated_token_address, get_associated_token_address_with_program_id}; use spl_token_2022::{ extension::{ - mint_close_authority::MintCloseAuthority, non_transferable::NonTransferable, pausable::PausableConfig, - permanent_delegate::PermanentDelegate, transfer_fee::TransferFeeConfig, BaseStateWithExtensionsMut, - ExtensionType, StateWithExtensionsMut, + metadata_pointer::MetadataPointer, mint_close_authority::MintCloseAuthority, non_transferable::NonTransferable, + pausable::PausableConfig, permanent_delegate::PermanentDelegate, transfer_fee::TransferFeeConfig, + BaseStateWithExtensionsMut, ExtensionType, StateWithExtensionsMut, }, state::Mint as Token2022Mint, ID as TOKEN_2022_PROGRAM_ID, @@ -204,6 +204,9 @@ impl TestContext { let ext = state.init_extension::(true).unwrap(); ext.close_authority = COption::Some(*mint_authority).try_into().unwrap(); } + ExtensionType::MetadataPointer => { + state.init_extension::(true).unwrap(); + } _ => panic!("Unsupported extension type for test helper"), } diff --git a/tests/test-hook-program/src/lib.rs b/tests/test-hook-program/src/lib.rs index 4768e54..3e4b72f 100644 --- a/tests/test-hook-program/src/lib.rs +++ b/tests/test-hook-program/src/lib.rs @@ -17,9 +17,23 @@ pinocchio::nostd_panic_handler!(); #[cfg(feature = "allow")] pub fn process_instruction( _program_id: &Address, - _accounts: &[AccountView], - _instruction_data: &[u8], + accounts: &[AccountView], + instruction_data: &[u8], ) -> ProgramResult { + use pinocchio::error::ProgramError; + + // Validate core context shape so integration tests catch missing account context. + // hook_point: 0=PreDeposit, 1=PostDeposit, 2=PreWithdraw, 3=PostWithdraw + let hook_point = *instruction_data.first().ok_or(ProgramError::InvalidInstructionData)?; + match hook_point { + 0..=3 => { + if accounts.len() < 3 { + return Err(ProgramError::Custom(42)); + } + } + _ => return Err(ProgramError::InvalidInstructionData), + } + Ok(()) }