add test and refactor chain_caller program

nssa/core/src/program.rs

@@ -16,6 +16,12 @@ pub struct ProgramInput<T> {
 #[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 {

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<PdaSeed>) = (
+            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<PdaSeed>) = (
             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<PdaSeed>) = (
+            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);
+    }
 }

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<PdaSeed>);
 
 /// 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::<Instruction>();
 
-    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,
     );
 }