fix account passing mechanism

nssa/core/src/program.rs

@@ -17,7 +17,7 @@ pub struct ProgramInput<T> {
 pub struct ChainedCall {
     pub program_id: ProgramId,
     pub instruction_data: InstructionData,
-    pub account_indices: Vec<usize>,
+    pub pre_states: Vec<AccountWithMetadata>,
 }
 
 #[derive(Serialize, Deserialize, Clone)]

nssa/src/public_transaction/transaction.rs

@@ -115,11 +115,22 @@ impl PublicTransaction {
 
             let mut program_output = program.execute(&input_pre_states, &instruction_data)?;
 
-            // This check is equivalent to checking that the program output pre_states coinicide
-            // with the values in the public state or with any modifications to those values
-            // during the chain of calls.
-            if input_pre_states != program_output.pre_states {
-                return Err(NssaError::InvalidProgramBehavior);
+            for pre in program_output.pre_states.iter() {
+                let account_id = pre.account_id;
+                // Check that the program output pre_states coinicide with the values in the public
+                // state or with any modifications to those values during the chain of calls.
+                let expected_pre = state_diff
+                    .get(&account_id)
+                    .cloned()
+                    .unwrap_or_else(|| state.get_account_by_address(&account_id));
+                if pre.account != expected_pre {
+                    return Err(NssaError::InvalidProgramBehavior);
+                }
+
+                // Check that authorization flags are consistent with the provided ones
+                if pre.is_authorized && !signer_addresses.contains(&account_id) {
+                    return Err(NssaError::InvalidProgramBehavior);
+                }
             }
 
             // Verify execution corresponds to a well-behaved program.
@@ -153,31 +164,7 @@ impl PublicTransaction {
             if let Some(next_chained_call) = chained_calls.pop() {
                 program_id = next_chained_call.program_id;
                 instruction_data = next_chained_call.instruction_data;
-
-                // Build post states with metadata for next call
-                let mut post_states_with_metadata = Vec::new();
-                for (pre, post) in program_output
-                    .pre_states
-                    .iter()
-                    .zip(program_output.post_states)
-                {
-                    let mut post_with_metadata = pre.clone();
-                    post_with_metadata.account = post.clone();
-                    post_states_with_metadata.push(post_with_metadata);
-                }
-
-                input_pre_states = next_chained_call
-                    .account_indices
-                    .iter()
-                    .map(|&i| {
-                        post_states_with_metadata
-                            .get(i)
-                            .ok_or_else(|| {
-                                NssaError::InvalidInput("Invalid account indices".into())
-                            })
-                            .cloned()
-                    })
-                    .collect::<Result<Vec<_>, NssaError>>()?;
+                input_pre_states = next_chained_call.pre_states;
             } else {
                 break;
             };

nssa/src/state.rs

@@ -2085,27 +2085,26 @@ pub mod tests {
     fn test_chained_call_succeeds() {
         let program = Program::chain_caller();
         let key = PrivateKey::try_new([1; 32]).unwrap();
-        let address = Address::from(&PublicKey::new_from_private_key(&key));
+        let from_address = Address::from(&PublicKey::new_from_private_key(&key));
+        let to_address = Address::new([2; 32]);
         let initial_balance = 100;
-        let initial_data = [(address, initial_balance)];
+        let initial_data = [(from_address, initial_balance), (to_address, 0)];
         let mut state =
             V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
-        let from = address;
         let from_key = key;
-        let to = Address::new([2; 32]);
-        let amount: u128 = 37;
+        let amount: u128 = 0;
         let instruction: (u128, ProgramId, u32) =
             (amount, Program::authenticated_transfer_program().id(), 2);
 
         let expected_to_post = Account {
-            program_owner: Program::chain_caller().id(),
+            program_owner: Program::authenticated_transfer_program().id(),
             balance: amount * 2, // The `chain_caller` chains the program twice
             ..Account::default()
         };
 
         let message = public_transaction::Message::try_new(
             program.id(),
-            vec![to, from], //The chain_caller program permutes the account order in the chain call
+            vec![to_address, from_address], //The chain_caller program permutes the account order in the chain call
             vec![0],
             instruction,
         )
@@ -2115,8 +2114,8 @@ pub mod tests {
 
         state.transition_from_public_transaction(&tx).unwrap();
 
-        let from_post = state.get_account_by_address(&from);
-        let to_post = state.get_account_by_address(&to);
+        let from_post = state.get_account_by_address(&from_address);
+        let to_post = state.get_account_by_address(&to_address);
         // The `chain_caller` program calls the program twice
         assert_eq!(from_post.balance, initial_balance - 2 * amount);
         assert_eq!(to_post, expected_to_post);
@@ -2126,14 +2125,13 @@ pub mod tests {
     fn test_execution_fails_if_chained_calls_exceeds_depth() {
         let program = Program::chain_caller();
         let key = PrivateKey::try_new([1; 32]).unwrap();
-        let address = Address::from(&PublicKey::new_from_private_key(&key));
+        let from_address = Address::from(&PublicKey::new_from_private_key(&key));
+        let to_address = Address::new([2; 32]);
         let initial_balance = 100;
-        let initial_data = [(address, initial_balance)];
+        let initial_data = [(from_address, initial_balance), (to_address, 0)];
         let mut state =
             V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
-        let from = address;
         let from_key = key;
-        let to = Address::new([2; 32]);
         let amount: u128 = 0;
         let instruction: (u128, ProgramId, u32) = (
             amount,
@@ -2143,7 +2141,7 @@ pub mod tests {
 
         let message = public_transaction::Message::try_new(
             program.id(),
-            vec![to, from], //The chain_caller program permutes the account order in the chain call
+            vec![to_address, from_address], //The chain_caller program permutes the account order in the chain call
             vec![0],
             instruction,
         )

nssa/test_program_methods/guest/src/bin/chain_caller.rs

@@ -24,7 +24,7 @@ fn main() {
         ChainedCall {
             program_id,
             instruction_data: instruction_data.clone(),
-            account_indices: vec![0, 1],
+            pre_states: vec![receiver_pre.clone(), sender_pre.clone()], // <- Account order permutation here
         };
         num_chain_calls as usize - 1
     ];
@@ -32,7 +32,7 @@ fn main() {
     chained_call.push(ChainedCall {
         program_id,
         instruction_data,
-        account_indices: vec![1, 0], // <- Account order permutation here
+        pre_states: vec![receiver_pre.clone(), sender_pre.clone()], // <- Account order permutation here
     });
 
     write_nssa_outputs_with_chained_call(