perf(program): reduce PDA validation compute units (#39)

* perf(program): reduce PDA validation compute units

Replace find_program_address in validate_pda with create_program_address
since the bump is already provided in instruction data. Override
validate_self in Escrow, Receipt, AllowedMint, and EscrowExtensions to
use derive_address with the stored bump, avoiding the off-curve check
that is unnecessary for accounts validated at creation.

Also documents 2-step admin transfer as a future improvement.

* fix(program): enforce canonical PDA validation on init

Restore canonical bump enforcement for instruction-provided PDA bumps while preserving a fast path for trusted stored bumps.

Add integration regressions that verify noncanonical-but-valid PDA addresses are rejected for escrow, allowed mint, receipt, and extensions initialization flows.

* chore(program): remove stale PDA CU comment

* refactor(program): remove dead `validate_pda_with_bump` trait method

No concrete type called this — all override `validate_self` with
`Address::derive_address` directly. Default `validate_self` now
falls back to `validate_pda` (canonical bump enforcement).

---------

Co-authored-by: Jo D <dev-jodee@users.noreply.github.com>

Author: dev-jodee

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.

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)
     }
 }

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 {

program/src/state/escrow_extensions.rs

@@ -374,21 +374,25 @@ pub fn update_or_append_extension<const N: usize>(
 ///
 /// 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);
         }
 
         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(())

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 {

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<u8, ProgramError> {
         let (derived, bump) = self.derive_address(program_id);

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();

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();

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();

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();