diff --git a/docs/IMPROVEMENTS.md b/docs/IMPROVEMENTS.md index dfb7b9d..40d8f5b 100644 --- a/docs/IMPROVEMENTS.md +++ b/docs/IMPROVEMENTS.md @@ -13,3 +13,5 @@ The following enhancements could be considered for future iterations of the prog 5. **Receipt Seed Space Optimization** - The current `receipt_seed` uses a 32-byte `Address` type. Two alternatives could save space: - **Use `u8` counter**: Change to a simple counter (0-255), saving 31 bytes per receipt. Limits to 256 receipts per depositor/escrow/mint combination, which is acceptable for most use cases. - **Single receipt with `deposit_additional` instruction**: Allow users to add to an existing receipt rather than creating new ones. This would require handling complexities around `deposited_at` timestamps (e.g., weighted average, use latest, or track per-deposit). + +6. **Two-Step Admin Transfer** - The current `UpdateAdmin` instruction requires both the current and new admin to sign the same transaction. This is problematic when transferring to/from multisig wallets (e.g., Squads), since both parties must be present in one transaction. A 2-step pattern (`ProposeAdmin` → `AcceptAdmin`, with optional `CancelAdminTransfer` and a timeout) would allow async coordination between parties and is the standard pattern for admin handoffs in production programs. diff --git a/program/src/state/allowed_mint.rs b/program/src/state/allowed_mint.rs index e9de052..689a824 100644 --- a/program/src/state/allowed_mint.rs +++ b/program/src/state/allowed_mint.rs @@ -65,8 +65,14 @@ impl AllowedMint { 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)?; + let derived = Address::derive_address( + &[AllowedMintPda::PREFIX, escrow.as_ref(), mint.as_ref()], + Some(state.bump), + program_id, + ); + if account.address() != &derived { + return Err(ProgramError::InvalidSeeds); + } Ok(state) } } diff --git a/program/src/state/escrow.rs b/program/src/state/escrow.rs index a2a99af..18b45bf 100644 --- a/program/src/state/escrow.rs +++ b/program/src/state/escrow.rs @@ -79,6 +79,15 @@ impl PdaAccount for Escrow { fn bump(&self) -> u8 { self.bump } + + #[inline(always)] + fn validate_self(&self, account: &AccountView, program_id: &Address) -> Result<(), ProgramError> { + let derived = Address::derive_address(&[Self::PREFIX, self.escrow_seed.as_ref()], Some(self.bump), program_id); + if account.address() != &derived { + return Err(ProgramError::InvalidSeeds); + } + Ok(()) + } } impl Escrow { diff --git a/program/src/state/escrow_extensions.rs b/program/src/state/escrow_extensions.rs index f619c71..0b1f2b3 100644 --- a/program/src/state/escrow_extensions.rs +++ b/program/src/state/escrow_extensions.rs @@ -374,11 +374,8 @@ pub fn update_or_append_extension( /// /// Returns `Ok(())` if: /// - Extensions account is the correct PDA for this escrow -/// - If data exists, the stored bump matches the canonical bump +/// - If data exists, the stored bump matches the derived address pub fn validate_extensions_pda(escrow: &AccountView, extensions: &AccountView, program_id: &Address) -> ProgramResult { - let extensions_pda = ExtensionsPda::new(escrow.address()); - 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); @@ -386,9 +383,16 @@ pub fn validate_extensions_pda(escrow: &AccountView, extensions: &AccountView, p let data = extensions.try_borrow()?; let header = EscrowExtensionsHeader::from_bytes(&data)?; - if header.bump != expected_bump { + + let derived = + Address::derive_address(&[ExtensionsPda::PREFIX, escrow.address().as_ref()], Some(header.bump), program_id); + if extensions.address() != &derived { return Err(ProgramError::InvalidSeeds); } + } else { + // Extensions account not yet created — derive canonical bump via find_program_address + let extensions_pda = ExtensionsPda::new(escrow.address()); + extensions_pda.validate_pda_address(extensions, program_id)?; } Ok(()) diff --git a/program/src/state/receipt.rs b/program/src/state/receipt.rs index 82917a6..ba6ae27 100644 --- a/program/src/state/receipt.rs +++ b/program/src/state/receipt.rs @@ -124,6 +124,25 @@ impl PdaAccount for Receipt { fn bump(&self) -> u8 { self.bump } + + #[inline(always)] + fn validate_self(&self, account: &AccountView, program_id: &Address) -> Result<(), ProgramError> { + let derived = Address::derive_address( + &[ + Self::PREFIX, + self.escrow.as_ref(), + self.depositor.as_ref(), + self.mint.as_ref(), + self.receipt_seed.as_ref(), + ], + Some(self.bump), + program_id, + ); + if account.address() != &derived { + return Err(ProgramError::InvalidSeeds); + } + Ok(()) + } } impl Receipt { diff --git a/program/src/traits/pda.rs b/program/src/traits/pda.rs index 174f127..e382fb0 100644 --- a/program/src/traits/pda.rs +++ b/program/src/traits/pda.rs @@ -21,7 +21,10 @@ pub trait PdaSeeds { Address::find_program_address(&seeds, program_id) } - /// Validate that account matches derived PDA + /// Validate that account matches the canonical PDA for these seeds. + /// + /// This enforces the canonical bump (highest valid bump) returned by + /// `find_program_address`. #[inline(always)] fn validate_pda(&self, account: &AccountView, program_id: &Address, expected_bump: u8) -> Result<(), ProgramError> { let (derived, bump) = self.derive_address(program_id); @@ -34,7 +37,10 @@ pub trait PdaSeeds { Ok(()) } - /// Validate that account address matches derived PDA, returns canonical bump + /// Validate that account address matches derived PDA and return the canonical bump. + /// + /// Uses `find_program_address` — only call this when the bump is not already known. + /// When bump is available, prefer `validate_pda` or `validate_self`. #[inline(always)] fn validate_pda_address(&self, account: &AccountView, program_id: &Address) -> Result { let (derived, bump) = self.derive_address(program_id); diff --git a/tests/integration-tests/src/test_add_timelock.rs b/tests/integration-tests/src/test_add_timelock.rs index df558f9..1417915 100644 --- a/tests/integration-tests/src/test_add_timelock.rs +++ b/tests/integration-tests/src/test_add_timelock.rs @@ -2,9 +2,9 @@ use crate::{ fixtures::{AddTimelockFixture, CreateEscrowFixture, SetImmutableFixture}, utils::{ assert_escrow_error, assert_extensions_header, assert_instruction_error, assert_timelock_extension, - find_escrow_pda, find_extensions_pda, test_empty_data, test_missing_signer, test_not_writable, - test_truncated_data, test_wrong_account, test_wrong_current_program, test_wrong_system_program, EscrowError, - InstructionTestFixture, TestContext, RANDOM_PUBKEY, + find_escrow_pda, find_extensions_pda, find_noncanonical_program_address, test_empty_data, test_missing_signer, + test_not_writable, test_truncated_data, test_wrong_account, test_wrong_current_program, + test_wrong_system_program, EscrowError, InstructionTestFixture, TestContext, RANDOM_PUBKEY, }, }; use solana_sdk::{instruction::InstructionError, signature::Signer}; @@ -53,6 +53,24 @@ fn test_add_timelock_invalid_extensions_bump() { assert_instruction_error(error, InstructionError::InvalidSeeds); } +#[test] +fn test_add_timelock_noncanonical_extensions_pda_rejected() { + let mut ctx = TestContext::new(); + let test_ix = AddTimelockFixture::build_valid(&mut ctx); + let escrow = test_ix.instruction.accounts[2].pubkey; + let program_id = test_ix.instruction.program_id; + + let (noncanonical_extensions, noncanonical_bump) = + find_noncanonical_program_address(&[b"extensions", escrow.as_ref()], &program_id) + .expect("expected at least one noncanonical bump"); + + let error = test_ix + .with_account_at(3, noncanonical_extensions) + .with_data_byte_at(1, noncanonical_bump) + .send_expect_error(&mut ctx); + assert_instruction_error(error, InstructionError::InvalidSeeds); +} + #[test] fn test_add_timelock_empty_data() { let mut ctx = TestContext::new(); diff --git a/tests/integration-tests/src/test_allow_mint.rs b/tests/integration-tests/src/test_allow_mint.rs index 91b3bc3..de21288 100644 --- a/tests/integration-tests/src/test_allow_mint.rs +++ b/tests/integration-tests/src/test_allow_mint.rs @@ -2,8 +2,9 @@ use crate::{ fixtures::{AllowMintFixture, AllowMintSetup}, utils::{ assert_account_exists, assert_allowed_mint_account, assert_escrow_error, assert_instruction_error, - find_allowed_mint_pda, test_missing_signer, test_not_writable, test_wrong_current_program, - test_wrong_system_program, EscrowError, InstructionTestFixture, TestContext, RANDOM_PUBKEY, + find_allowed_mint_pda, find_noncanonical_program_address, test_missing_signer, test_not_writable, + test_wrong_current_program, test_wrong_system_program, EscrowError, InstructionTestFixture, TestContext, + RANDOM_PUBKEY, }, }; use escrow_program_client::instructions::{AllowMintBuilder, SetImmutableBuilder}; @@ -64,6 +65,26 @@ fn test_allow_mint_invalid_bump() { assert_instruction_error(error, InstructionError::InvalidSeeds); } +#[test] +fn test_allow_mint_noncanonical_pda_rejected() { + let mut ctx = TestContext::new(); + let setup = AllowMintSetup::new(&mut ctx); + let valid_ix = setup.build_instruction(&ctx); + let program_id = valid_ix.instruction.program_id; + + let (noncanonical_allowed_mint, noncanonical_bump) = find_noncanonical_program_address( + &[b"allowed_mint", setup.escrow_pda.as_ref(), setup.mint_pubkey.as_ref()], + &program_id, + ) + .expect("expected at least one noncanonical bump"); + + let error = valid_ix + .with_account_at(5, noncanonical_allowed_mint) + .with_data_byte_at(1, noncanonical_bump) + .send_expect_error(&mut ctx); + assert_instruction_error(error, InstructionError::InvalidSeeds); +} + #[test] fn test_allow_mint_wrong_admin() { let mut ctx = TestContext::new(); diff --git a/tests/integration-tests/src/test_create_escrow.rs b/tests/integration-tests/src/test_create_escrow.rs index 5d312ce..f9f9a35 100644 --- a/tests/integration-tests/src/test_create_escrow.rs +++ b/tests/integration-tests/src/test_create_escrow.rs @@ -1,8 +1,8 @@ use crate::{ fixtures::CreateEscrowFixture, utils::{ - assert_escrow_account, assert_escrow_mutability, assert_instruction_error, test_empty_data, - test_missing_signer, test_not_writable, test_wrong_account, test_wrong_current_program, + assert_escrow_account, assert_escrow_mutability, assert_instruction_error, find_noncanonical_program_address, + test_empty_data, test_missing_signer, test_not_writable, test_wrong_account, test_wrong_current_program, test_wrong_system_program, InstructionTestFixture, TestContext, }, }; @@ -70,6 +70,25 @@ fn test_create_escrow_invalid_bump() { assert_instruction_error(error, InstructionError::InvalidSeeds); } +#[test] +fn test_create_escrow_noncanonical_bump_rejected() { + let mut ctx = TestContext::new(); + let valid_ix = CreateEscrowFixture::build_valid(&mut ctx); + + let escrow_seed = valid_ix.instruction.accounts[2].pubkey; + let program_id = valid_ix.instruction.program_id; + let (noncanonical_escrow, noncanonical_bump) = + find_noncanonical_program_address(&[b"escrow", escrow_seed.as_ref()], &program_id) + .expect("expected at least one noncanonical bump"); + + let error = valid_ix + .with_account_at(3, noncanonical_escrow) + .with_data_byte_at(1, noncanonical_bump) + .send_expect_error(&mut ctx); + + assert_instruction_error(error, InstructionError::InvalidSeeds); +} + #[test] fn test_create_escrow_empty_data() { let mut ctx = TestContext::new(); diff --git a/tests/integration-tests/src/test_deposit.rs b/tests/integration-tests/src/test_deposit.rs index f8b9571..ba585b5 100644 --- a/tests/integration-tests/src/test_deposit.rs +++ b/tests/integration-tests/src/test_deposit.rs @@ -4,10 +4,10 @@ use crate::{ 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, - test_wrong_system_program, test_wrong_token_program, EscrowError, TestContext, TestInstruction, - TEST_HOOK_ALLOW_ID, TEST_HOOK_DENY_ERROR, TEST_HOOK_DENY_ID, + assert_custom_error, assert_escrow_error, assert_instruction_error, find_noncanonical_program_address, + find_receipt_pda, test_empty_data, 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::DepositBuilder; @@ -84,6 +84,29 @@ fn test_deposit_invalid_bump() { assert_instruction_error(error, InstructionError::InvalidSeeds); } +#[test] +fn test_deposit_noncanonical_receipt_pda_rejected() { + let mut ctx = TestContext::new(); + let setup = DepositSetup::new(&mut ctx); + let valid_ix = setup.build_instruction(&ctx); + let program_id = valid_ix.instruction.program_id; + + let depositor = setup.depositor.pubkey(); + let mint = setup.mint.pubkey(); + let receipt_seed = setup.receipt_seed.pubkey(); + let (noncanonical_receipt, noncanonical_bump) = find_noncanonical_program_address( + &[b"receipt", setup.escrow_pda.as_ref(), depositor.as_ref(), mint.as_ref(), receipt_seed.as_ref()], + &program_id, + ) + .expect("expected at least one noncanonical bump"); + + let error = valid_ix + .with_account_at(5, noncanonical_receipt) + .with_data_byte_at(1, noncanonical_bump) + .send_expect_error(&mut ctx); + assert_instruction_error(error, InstructionError::InvalidSeeds); +} + #[test] fn test_deposit_wrong_token_program() { let mut ctx = TestContext::new(); diff --git a/tests/integration-tests/src/utils/pda_utils.rs b/tests/integration-tests/src/utils/pda_utils.rs index 4882fbe..550ffea 100644 --- a/tests/integration-tests/src/utils/pda_utils.rs +++ b/tests/integration-tests/src/utils/pda_utils.rs @@ -20,3 +20,23 @@ pub fn find_receipt_pda(escrow: &Pubkey, depositor: &Pubkey, mint: &Pubkey, rece pub fn find_allowed_mint_pda(escrow: &Pubkey, mint: &Pubkey) -> (Pubkey, u8) { AllowedMint::find_pda(escrow, mint) } + +pub fn find_noncanonical_program_address(seeds: &[&[u8]], program_id: &Pubkey) -> Option<(Pubkey, u8)> { + let (_, canonical_bump) = Pubkey::find_program_address(seeds, program_id); + + for bump in (0u8..=u8::MAX).rev() { + if bump == canonical_bump { + continue; + } + + let bump_seed = [bump]; + let mut seeds_with_bump = seeds.to_vec(); + seeds_with_bump.push(&bump_seed); + + if let Ok(address) = Pubkey::create_program_address(&seeds_with_bump, program_id) { + return Some((address, bump)); + } + } + + None +}