From e61a971790cf5ba9cde4e7de9a34b401f2460aff Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Thu, 27 Nov 2025 13:49:56 -0300 Subject: [PATCH] add test and refactor chain_caller program --- nssa/core/src/program.rs | 6 ++ nssa/src/state.rs | 60 +++++++++++++++++-- .../guest/src/bin/chain_caller.rs | 50 +++++++++------- 3 files changed, 88 insertions(+), 28 deletions(-) diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs index ad9bbab..dfd52e4 100644 --- a/nssa/core/src/program.rs +++ b/nssa/core/src/program.rs @@ -16,6 +16,12 @@ pub struct ProgramInput { #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct PdaSeed([u8; 32]); +impl PdaSeed { + pub fn new(value: [u8; 32]) -> Self { + Self(value) + } +} + #[cfg(feature = "host")] impl From<(&ProgramId, &PdaSeed)> for AccountId { fn from(value: (&ProgramId, &PdaSeed)) -> Self { diff --git a/nssa/src/state.rs b/nssa/src/state.rs index cef7791..79541a3 100644 --- a/nssa/src/state.rs +++ b/nssa/src/state.rs @@ -250,7 +250,7 @@ pub mod tests { Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, account::{Account, AccountId, AccountWithMetadata, Nonce}, encryption::{EphemeralPublicKey, IncomingViewingPublicKey, Scalar}, - program::ProgramId, + program::{PdaSeed, ProgramId}, }; use crate::{ @@ -2092,14 +2092,18 @@ pub mod tests { let key = PrivateKey::try_new([1; 32]).unwrap(); let from = AccountId::from(&PublicKey::new_from_private_key(&key)); let to = AccountId::new([2; 32]); - let initial_balance = 100; + let initial_balance = 1000; let initial_data = [(from, initial_balance), (to, 0)]; let mut state = V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs(); let from_key = key; - let amount: u128 = 0; - let instruction: (u128, ProgramId, u32) = - (amount, Program::authenticated_transfer_program().id(), 2); + let amount: u128 = 37; + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + Program::authenticated_transfer_program().id(), + 2, + None, + ); let expected_to_post = Account { program_owner: Program::authenticated_transfer_program().id(), @@ -2139,10 +2143,11 @@ pub mod tests { V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs(); let from_key = key; let amount: u128 = 0; - let instruction: (u128, ProgramId, u32) = ( + let instruction: (u128, ProgramId, u32, Option) = ( amount, Program::authenticated_transfer_program().id(), MAX_NUMBER_CHAINED_CALLS as u32 + 1, + None, ); let message = public_transaction::Message::try_new( @@ -2162,4 +2167,47 @@ pub mod tests { Err(NssaError::MaxChainedCallsDepthExceeded) )); } + #[test] + fn test_execution_that_requires_authentication_of_a_program_derived_account_id_succeeds() { + let chain_caller = Program::chain_caller(); + let pda_seed = PdaSeed::new([37; 32]); + let from = AccountId::from((&chain_caller.id(), &pda_seed)); + let to = AccountId::new([2; 32]); + let initial_balance = 1000; + let initial_data = [(from, initial_balance), (to, 0)]; + let mut state = + V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs(); + let amount: u128 = 58; + let instruction: (u128, ProgramId, u32, Option) = ( + amount, + Program::authenticated_transfer_program().id(), + 1, + Some(pda_seed), + ); + + let expected_to_post = Account { + program_owner: Program::authenticated_transfer_program().id(), + balance: amount, // The `chain_caller` chains the program twice + ..Account::default() + }; + + let message = public_transaction::Message::try_new( + chain_caller.id(), + vec![to, from], // The chain_caller program permutes the account order in the chain + // call + vec![], + instruction, + ) + .unwrap(); + let witness_set = public_transaction::WitnessSet::for_message(&message, &[]); + let tx = PublicTransaction::new(message, witness_set); + + state.transition_from_public_transaction(&tx).unwrap(); + + let from_post = state.get_account_by_id(&from); + let to_post = state.get_account_by_id(&to); + // The `chain_caller` program calls the program twice + assert_eq!(from_post.balance, initial_balance - amount); + assert_eq!(to_post, expected_to_post); + } } diff --git a/nssa/test_program_methods/guest/src/bin/chain_caller.rs b/nssa/test_program_methods/guest/src/bin/chain_caller.rs index 23b0244..1885da2 100644 --- a/nssa/test_program_methods/guest/src/bin/chain_caller.rs +++ b/nssa/test_program_methods/guest/src/bin/chain_caller.rs @@ -1,45 +1,51 @@ use nssa_core::program::{ - ChainedCall, ProgramId, ProgramInput, read_nssa_inputs, write_nssa_outputs_with_chained_call, + ChainedCall, PdaSeed, ProgramId, ProgramInput, read_nssa_inputs, + write_nssa_outputs_with_chained_call, }; use risc0_zkvm::serde::to_vec; -type Instruction = (u128, ProgramId, u32); +type Instruction = (u128, ProgramId, u32, Option); /// A program that calls another program `num_chain_calls` times. /// It permutes the order of the input accounts on the subsequent call +/// The `ProgramId` in the instruction must be the program_id of the authenticated transfers program fn main() { let ProgramInput { pre_states, - instruction: (balance, program_id, num_chain_calls), + instruction: (balance, auth_transfer_id, num_chain_calls, pda_seed), } = read_nssa_inputs::(); - let [sender_pre, receiver_pre] = match pre_states.try_into() { + let [recipient_pre, sender_pre] = match pre_states.try_into() { Ok(array) => array, Err(_) => return, }; let instruction_data = to_vec(&balance).unwrap(); - let mut chained_call = vec![ - ChainedCall { - program_id, - instruction_data: instruction_data.clone(), - pre_states: vec![receiver_pre.clone(), sender_pre.clone()], // <- Account order permutation here - pda_seeds: vec![] - }; - num_chain_calls as usize - 1 - ]; + let mut running_recipient_pre = recipient_pre.clone(); + let mut running_sender_pre = sender_pre.clone(); - chained_call.push(ChainedCall { - program_id, - instruction_data, - pre_states: vec![receiver_pre.clone(), sender_pre.clone()], // <- Account order permutation here - pda_seeds: vec![], - }); + if pda_seed.is_some() { + running_sender_pre.is_authorized = true; + } + + let mut chained_calls = Vec::new(); + for _i in 0..num_chain_calls { + let new_chained_call = ChainedCall { + program_id: auth_transfer_id, + instruction_data: instruction_data.clone(), + pre_states: vec![running_sender_pre.clone(), running_recipient_pre.clone()], // <- Account order permutation here + pda_seeds: pda_seed.iter().cloned().collect(), + }; + chained_calls.push(new_chained_call); + + running_sender_pre.account.balance -= balance; + running_recipient_pre.account.balance += balance; + } write_nssa_outputs_with_chained_call( - vec![sender_pre.clone(), receiver_pre.clone()], - vec![sender_pre.account, receiver_pre.account], - chained_call, + vec![recipient_pre.clone(), sender_pre.clone()], + vec![recipient_pre.account, sender_pre.account], + chained_calls, ); }