fix: check public account authorization in privacy preserving circuit

nssa/core/src/program.rs

@@ -30,6 +30,20 @@ impl PdaSeed {
     }
 }
 
+pub fn compute_authorized_pdas(
+    caller_program_id: Option<ProgramId>,
+    pda_seeds: &[PdaSeed],
+) -> HashSet<AccountId> {
+    caller_program_id
+        .map(|caller_program_id| {
+            pda_seeds
+                .iter()
+                .map(|pda_seed| AccountId::from((&caller_program_id, pda_seed)))
+                .collect()
+        })
+        .unwrap_or_default()
+}
+
 impl From<(&ProgramId, &PdaSeed)> for AccountId {
     fn from(value: (&ProgramId, &PdaSeed)) -> Self {
         use risc0_zkvm::sha::{Impl, Sha256};

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -43,7 +43,7 @@ impl From<Program> for ProgramWithDependencies {
 }
 
 /// Generates a proof of the execution of a NSSA program inside the privacy preserving execution
-/// circuit
+/// circuit.
 #[expect(clippy::too_many_arguments, reason = "TODO: fix later")]
 pub fn execute_and_prove(
     pre_states: Vec<AccountWithMetadata>,
@@ -87,8 +87,6 @@ pub fn execute_and_prove(
             .decode()
             .map_err(|e| NssaError::ProgramOutputDeserializationError(e.to_string()))?;
 
-        // TODO: Why private execution doesn't care about public account authorization?
-
         // TODO: remove clone
         program_outputs.push(program_output.clone());
 

nssa/src/program.rs

@@ -235,6 +235,17 @@ mod tests {
             }
         }
 
+        pub fn malicious_authorization_changer() -> Self {
+            use test_program_methods::{
+                MALICIOUS_AUTHORIZATION_CHANGER_ELF, MALICIOUS_AUTHORIZATION_CHANGER_ID,
+            };
+
+            Program {
+                id: MALICIOUS_AUTHORIZATION_CHANGER_ID,
+                elf: MALICIOUS_AUTHORIZATION_CHANGER_ELF.to_vec(),
+            }
+        }
+
         pub fn modified_transfer_program() -> Self {
             use test_program_methods::MODIFIED_TRANSFER_ELF;
             // This unwrap won't panic since the `MODIFIED_TRANSFER_ELF` comes from risc0 build of

nssa/src/public_transaction/transaction.rs

@@ -4,7 +4,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use log::debug;
 use nssa_core::{
     account::{Account, AccountId, AccountWithMetadata},
-    program::{ChainedCall, DEFAULT_PROGRAM_ID, PdaSeed, ProgramId, validate_execution},
+    program::{ChainedCall, DEFAULT_PROGRAM_ID, validate_execution},
 };
 use sha2::{Digest, digest::FixedOutput};
 
@@ -135,8 +135,10 @@ impl PublicTransaction {
                 chained_call.program_id, program_output
             );
 
-            let authorized_pdas =
-                Self::compute_authorized_pdas(&caller_program_id, &chained_call.pda_seeds);
+            let authorized_pdas = nssa_core::program::compute_authorized_pdas(
+                caller_program_id,
+                &chained_call.pda_seeds,
+            );
 
             for pre in &program_output.pre_states {
                 let account_id = pre.account_id;
@@ -200,20 +202,6 @@ impl PublicTransaction {
 
         Ok(state_diff)
     }
-
-    fn compute_authorized_pdas(
-        caller_program_id: &Option<ProgramId>,
-        pda_seeds: &[PdaSeed],
-    ) -> HashSet<AccountId> {
-        if let Some(caller_program_id) = caller_program_id {
-            pda_seeds
-                .iter()
-                .map(|pda_seed| AccountId::from((caller_program_id, pda_seed)))
-                .collect()
-        } else {
-            HashSet::new()
-        }
-    }
 }
 
 #[cfg(test)]

nssa/src/state.rs

@@ -4369,4 +4369,60 @@ pub mod tests {
 
         assert!(matches!(res, Err(NssaError::CircuitProvingError(_))));
     }
+
+    #[test]
+    fn test_malicious_authorization_changer_should_fail_in_privacy_preserving_circuit() {
+        // Arrange
+        let malicious_program = Program::malicious_authorization_changer();
+        let auth_transfers = Program::authenticated_transfer_program();
+        let sender_keys = test_public_account_keys_1();
+        let recipient_keys = test_private_account_keys_1();
+
+        let sender_account = AccountWithMetadata::new(
+            Account {
+                program_owner: auth_transfers.id(),
+                balance: 100,
+                ..Default::default()
+            },
+            false,
+            sender_keys.account_id(),
+        );
+        let recipient_account =
+            AccountWithMetadata::new(Account::default(), true, &recipient_keys.npk());
+
+        let recipient_commitment =
+            Commitment::new(&recipient_keys.npk(), &recipient_account.account);
+        let state = V02State::new_with_genesis_accounts(
+            &[(sender_account.account_id, sender_account.account.balance)],
+            std::slice::from_ref(&recipient_commitment),
+        )
+        .with_test_programs();
+
+        let balance_to_transfer = 10u128;
+        let instruction = (balance_to_transfer, auth_transfers.id());
+
+        let recipient_esk = [3; 32];
+        let recipient = SharedSecretKey::new(&recipient_esk, &recipient_keys.ivk());
+
+        let mut dependencies = HashMap::new();
+        dependencies.insert(auth_transfers.id(), auth_transfers);
+        let program_with_deps = ProgramWithDependencies::new(malicious_program, dependencies);
+
+        let recipient_new_nonce = 0xdeadbeef1;
+
+        // Act - execute the malicious program - this should fail during proving
+        let result = execute_and_prove(
+            vec![sender_account, recipient_account],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![0, 1],
+            vec![recipient_new_nonce],
+            vec![(recipient_keys.npk(), recipient)],
+            vec![recipient_keys.nsk],
+            vec![state.get_proof_for_commitment(&recipient_commitment)],
+            &program_with_deps,
+        );
+
+        // Assert - should fail because the malicious program tries to manipulate is_authorized
+        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
+    }
 }

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::{HashMap, VecDeque},
+    collections::{HashMap, HashSet, VecDeque, hash_map::Entry},
     convert::Infallible,
 };
 
@@ -10,8 +10,8 @@ use nssa_core::{
     account::{Account, AccountId, AccountWithMetadata, Nonce},
     compute_digest_for_path,
     program::{
-        ChainedCall, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, ProgramId, ProgramOutput,
-        validate_execution,
+        AccountPostState, ChainedCall, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, ProgramId,
+        ProgramOutput, validate_execution,
     },
 };
 use risc0_zkvm::{guest::env, serde::to_vec};
@@ -51,7 +51,7 @@ impl ExecutionState {
     /// Validate program outputs and derive the overall execution state.
     pub fn derive_from_outputs(program_id: ProgramId, program_outputs: Vec<ProgramOutput>) -> Self {
         let Some(first_output) = program_outputs.first() else {
-            panic!("Program outputs is empty")
+            panic!("No program outputs provided");
         };
 
         let initial_call = ChainedCall {
@@ -60,7 +60,7 @@ impl ExecutionState {
             pre_states: first_output.pre_states.clone(),
             pda_seeds: Vec::new(),
         };
-        let mut chained_calls = VecDeque::from_iter([initial_call]);
+        let mut chained_calls = VecDeque::from_iter([(initial_call, None)]);
 
         let mut execution_state = ExecutionState {
             pre_states: Vec::new(),
@@ -69,7 +69,7 @@ impl ExecutionState {
         let mut last_program_id = program_id;
         let mut program_outputs_iter = program_outputs.into_iter();
         let mut chain_calls_counter = 0;
-        while let Some(chained_call) = chained_calls.pop_front() {
+        while let Some((chained_call, caller_program_id)) = chained_calls.pop_front() {
             assert!(
                 chain_calls_counter <= MAX_NUMBER_CHAINED_CALLS,
                 "Max chained calls depth is exceeded"
@@ -93,8 +93,6 @@ impl ExecutionState {
                 |_: Infallible| unreachable!("Infallible error is never constructed"),
             );
 
-            // TODO: Why private execution doesn't care about public account authorization?
-
             // Check that the program is well behaved.
             // See the # Programs section for the definition of the `validate_execution` method.
             let execution_valid = validate_execution(
@@ -105,10 +103,19 @@ impl ExecutionState {
             assert!(execution_valid, "Bad behaved program");
 
             for next_call in program_output.chained_calls.iter().rev() {
-                chained_calls.push_front(next_call.clone());
+                chained_calls.push_front((next_call.clone(), Some(chained_call.program_id)));
             }
 
-            execution_state.populate_from_output(chained_call.program_id, program_output);
+            let authorized_pdas = nssa_core::program::compute_authorized_pdas(
+                caller_program_id,
+                &chained_call.pda_seeds,
+            );
+            execution_state.validate_and_sync_states(
+                chained_call.program_id,
+                authorized_pdas,
+                program_output.pre_states,
+                program_output.post_states,
+            );
             last_program_id = chained_call.program_id;
             chain_calls_counter += 1;
         }
@@ -118,7 +125,7 @@ impl ExecutionState {
             "Inner call without a chained call found",
         );
 
-        // Claim accounts
+        // Claim accounts which were not explicitly claimed during execution
         for account in execution_state.post_states.values_mut() {
             if account.program_owner == DEFAULT_PROGRAM_ID {
                 account.program_owner = last_program_id;
@@ -128,17 +135,48 @@ impl ExecutionState {
         execution_state
     }
 
-    fn populate_from_output(&mut self, program_id: ProgramId, program_output: ProgramOutput) {
-        for (pre, mut post) in program_output
-            .pre_states
-            .into_iter()
-            .zip(program_output.post_states)
-        {
+    /// Validate program pre and post states and populate the execution state.
+    fn validate_and_sync_states(
+        &mut self,
+        program_id: ProgramId,
+        authorized_pdas: HashSet<AccountId>,
+        pre_states: Vec<AccountWithMetadata>,
+        post_states: Vec<AccountPostState>,
+    ) {
+        for (pre, mut post) in pre_states.into_iter().zip(post_states) {
             let pre_account_id = pre.account_id;
-            if let Some(account_pre) = self.post_states.get(&pre_account_id) {
-                assert_eq!(account_pre, &pre.account, "Inconsistent pre state");
-            } else {
-                self.pre_states.push(pre);
+            let post_states_entry = self.post_states.entry(pre.account_id);
+            match &post_states_entry {
+                Entry::Occupied(occupied) => {
+                    // Ensure that new pre state is the same as known post state
+                    assert_eq!(
+                        occupied.get(),
+                        &pre.account,
+                        "Inconsistent pre state for account {pre_account_id:?}",
+                    );
+
+                    let previous_is_authorized = self
+                        .pre_states
+                        .iter()
+                        .find(|acc| acc.account_id == pre_account_id)
+                        .map(|acc| acc.is_authorized)
+                        .unwrap_or_else(|| {
+                            panic!(
+                                "Pre state must exist in execution state for account {pre_account_id:?}",
+                            )
+                        });
+
+                    let is_authorized =
+                        previous_is_authorized || authorized_pdas.contains(&pre_account_id);
+
+                    assert_eq!(
+                        pre.is_authorized, is_authorized,
+                        "Inconsistent authorization for account {pre_account_id:?}",
+                    );
+                }
+                Entry::Vacant(_) => {
+                    self.pre_states.push(pre);
+                }
             }
 
             if post.requires_claim() {
@@ -146,11 +184,11 @@ impl ExecutionState {
                 if post.account().program_owner == DEFAULT_PROGRAM_ID {
                     post.account_mut().program_owner = program_id;
                 } else {
-                    panic!("Cannot claim an initialized account")
+                    panic!("Cannot claim an initialized account {pre_account_id:?}");
                 }
             }
 
-            self.post_states.insert(pre_account_id, post.into_account());
+            post_states_entry.insert_entry(post.into_account());
         }
     }
 

test_program_methods/guest/src/bin/malicious_authorization_changer.rs

@@ -0,0 +1,53 @@
+use nssa_core::{
+    account::AccountWithMetadata,
+    program::{
+        AccountPostState, ChainedCall, ProgramId, ProgramInput, read_nssa_inputs,
+        write_nssa_outputs_with_chained_call,
+    },
+};
+use risc0_zkvm::serde::to_vec;
+
+type Instruction = (u128, ProgramId);
+
+/// A malicious test program that attempts to change authorization status.
+/// It accepts two accounts and executes a native token transfer program via chain call,
+/// but sets the `is_authorized` field of the first account to true.
+fn main() {
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: (balance, transfer_program_id),
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
+
+    let [sender, receiver] = match pre_states.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    // Maliciously set is_authorized to true for the first account
+    let authorised_sender = AccountWithMetadata {
+        is_authorized: true,
+        ..sender.clone()
+    };
+
+    let instruction_data = to_vec(&balance).unwrap();
+
+    let chained_call = ChainedCall {
+        program_id: transfer_program_id,
+        instruction_data,
+        pre_states: vec![authorised_sender.clone(), receiver.clone()],
+        pda_seeds: vec![],
+    };
+
+    write_nssa_outputs_with_chained_call(
+        instruction_words,
+        vec![sender.clone(), receiver.clone()],
+        vec![
+            AccountPostState::new(sender.account),
+            AccountPostState::new(receiver.account),
+        ],
+        vec![chained_call],
+    );
+}