refactor. Add test for state transition from privacy preserving transaction

nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -2,7 +2,7 @@ use risc0_zkvm::{guest::env, serde::to_vec};
 
 use nssa_core::{
     account::{Account, AccountWithMetadata, Commitment, Nullifier, NullifierPublicKey},
-    program::{validate_execution, ProgramOutput},
+    program::{validate_execution, ProgramOutput, DEFAULT_PROGRAM_ID},
     verify_membership_proof, EncryptedAccountData, EphemeralPublicKey, EphemeralSecretKey,
     IncomingViewingPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, Tag,
 };
@@ -101,16 +101,20 @@ fn main() {
                 }
 
                 // Update post-state with new nonce
-                let mut post_with_updated_nonce = post_states[i].clone();
+                let mut post_with_updated_values = post_states[i].clone();
-                post_with_updated_nonce.nonce = *new_nonce;
+                post_with_updated_values.nonce = *new_nonce;
+
+                if post_with_updated_values.program_owner == DEFAULT_PROGRAM_ID {
+                    post_with_updated_values.program_owner = program_id;
+                }
 
                 // Compute commitment and push
-                let commitment_post = Commitment::new(Npk, &post_with_updated_nonce);
+                let commitment_post = Commitment::new(Npk, &post_with_updated_values);
                 new_commitments.push(commitment_post);
 
                 // Encrypt and push post state
                 let encrypted_account =
-                    EncryptedAccountData::new(&post_with_updated_nonce, esk, Npk, Ipk);
+                    EncryptedAccountData::new(&post_with_updated_values, esk, Npk, Ipk);
                 encrypted_private_post_states.push(encrypted_account);
             }
             _ => panic!("Invalid visibility mask value"),

nssa/src/privacy_preserving_transaction/mod.rs

@@ -4,6 +4,7 @@ mod transaction;
 mod witness_set;
 
 pub use message::Message;
+pub use witness_set::WitnessSet;
 pub use transaction::PrivacyPreservingTransaction;
 
 pub mod circuit {
@@ -11,7 +12,7 @@ pub mod circuit {
         CommitmentSetDigest, EphemeralSecretKey, IncomingViewingPublicKey, MembershipProof,
         PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
         account::{Account, AccountWithMetadata, Nonce, NullifierPublicKey, NullifierSecretKey},
-        program::{InstructionData, ProgramOutput},
+        program::{InstructionData, ProgramId, ProgramOutput},
     };
     use risc0_zkvm::{ExecutorEnv, InnerReceipt, Receipt, default_prover};
 
@@ -30,6 +31,56 @@ pub mod circuit {
         }
     }
 
+    pub fn execute_and_prove(
+        pre_states: &[AccountWithMetadata],
+        instruction_data: &InstructionData,
+        visibility_mask: &[u8],
+        private_account_nonces: &[u128],
+        private_account_keys: &[(
+            NullifierPublicKey,
+            IncomingViewingPublicKey,
+            EphemeralSecretKey,
+        )],
+        private_account_auth: &[(NullifierSecretKey, MembershipProof)],
+        program: &Program,
+        commitment_set_digest: &CommitmentSetDigest,
+    ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> {
+        let inner_receipt = execute_and_prove_program(program, pre_states, instruction_data)?;
+
+        let program_output: ProgramOutput = inner_receipt
+            .journal
+            .decode()
+            .map_err(|e| NssaError::ProgramOutputDeserializationError(e.to_string()))?;
+
+        let circuit_input = PrivacyPreservingCircuitInput {
+            program_output,
+            visibility_mask: visibility_mask.to_vec(),
+            private_account_nonces: private_account_nonces.to_vec(),
+            private_account_keys: private_account_keys.to_vec(),
+            private_account_auth: private_account_auth.to_vec(),
+            program_id: program.id(),
+            commitment_set_digest: *commitment_set_digest,
+        };
+
+        // Prove circuit.
+        let mut env_builder = ExecutorEnv::builder();
+        env_builder.add_assumption(inner_receipt);
+        env_builder.write(&circuit_input).unwrap();
+        let env = env_builder.build().unwrap();
+        let prover = default_prover();
+        let prove_info = prover.prove(env, PRIVACY_PRESERVING_CIRCUIT_ELF).unwrap();
+
+        let proof = Proof(borsh::to_vec(&prove_info.receipt.inner)?);
+
+        let circuit_output: PrivacyPreservingCircuitOutput = prove_info
+            .receipt
+            .journal
+            .decode()
+            .map_err(|e| NssaError::CircuitOutputDeserializationError(e.to_string()))?;
+
+        Ok((circuit_output, proof))
+    }
+
     fn execute_and_prove_program(
         program: &Program,
         pre_states: &[AccountWithMetadata],
@@ -47,56 +98,6 @@ pub mod circuit {
             .map_err(|e| NssaError::ProgramProveFailed(e.to_string()))?
             .receipt)
     }
-
-    pub fn prove_privacy_preserving_execution_circuit(
-        pre_states: &[AccountWithMetadata],
-        instruction_data: &InstructionData,
-        private_account_keys: &[(
-            NullifierPublicKey,
-            IncomingViewingPublicKey,
-            EphemeralSecretKey,
-        )],
-        private_account_auth: &[(NullifierSecretKey, MembershipProof)],
-        visibility_mask: &[u8],
-        private_account_nonces: &[u128],
-        commitment_set_digest: CommitmentSetDigest,
-        program: &Program,
-    ) -> Result<(Proof, PrivacyPreservingCircuitOutput), NssaError> {
-        let inner_receipt = execute_and_prove_program(program, pre_states, instruction_data)?;
-
-        let program_output: ProgramOutput = inner_receipt
-            .journal
-            .decode()
-            .map_err(|e| NssaError::ProgramOutputDeserializationError(e.to_string()))?;
-
-        let circuit_input = PrivacyPreservingCircuitInput {
-            program_output,
-            visibility_mask: visibility_mask.to_vec(),
-            private_account_nonces: private_account_nonces.to_vec(),
-            private_account_keys: private_account_keys.to_vec(),
-            private_account_auth: private_account_auth.to_vec(),
-            program_id: program.id(),
-            commitment_set_digest,
-        };
-
-        // Prove circuit.
-        let mut env_builder = ExecutorEnv::builder();
-        env_builder.add_assumption(inner_receipt);
-        env_builder.write(&circuit_input).unwrap();
-        let env = env_builder.build().unwrap();
-        let prover = default_prover();
-        let prove_info = prover.prove(env, PRIVACY_PRESERVING_CIRCUIT_ELF).unwrap();
-
-        let proof = Proof(borsh::to_vec(&prove_info.receipt.inner)?);
-
-        let circuit_output: PrivacyPreservingCircuitOutput = prove_info
-            .receipt
-            .journal
-            .decode()
-            .map_err(|e| NssaError::CircuitOutputDeserializationError(e.to_string()))?;
-
-        Ok((proof, circuit_output))
-    }
 }
 
 #[cfg(test)]
@@ -111,9 +112,11 @@ mod tests {
     use risc0_zkvm::{InnerReceipt, Journal, Receipt};
 
     use crate::{
-        Address, V01State, merkle_tree::MerkleTree,
+        Address, V01State,
-        privacy_preserving_transaction::circuit::prove_privacy_preserving_execution_circuit,
+        merkle_tree::MerkleTree,
-        program::Program, state::CommitmentSet,
+        privacy_preserving_transaction::circuit::{Proof, execute_and_prove},
+        program::Program,
+        state::CommitmentSet,
     };
 
     use rand::{Rng, RngCore, rngs::OsRng};
@@ -143,24 +146,15 @@ mod tests {
         };
 
         let expected_sender_pre = sender.clone();
-        let pre_states = vec![sender, recipient];
+        let (output, proof) = execute_and_prove(
-        let instruction_data = Program::serialize_instruction(balance_to_move).unwrap();
+            &[sender, recipient],
-        let private_account_keys = vec![(NullifierPublicKey::from(&[1; 32]), [2; 32], [3; 32])];
+            &Program::serialize_instruction(balance_to_move).unwrap(),
-        let private_account_nonces = vec![0xdeadbeef];
+            &[0, 2],
-
+            &[0xdeadbeef],
-        let private_account_auth = vec![];
+            &[(NullifierPublicKey::from(&[1; 32]), [2; 32], [3; 32])],
-        let visibility_mask = vec![0, 2];
+            &[],
-        let commitment_set_digest = [99; 32];
+            &Program::authenticated_transfer_program(),
-        let program = Program::authenticated_transfer_program();
+            &[99; 32],
-        let (proof, output) = prove_privacy_preserving_execution_circuit(
-            &pre_states,
-            &instruction_data,
-            &private_account_keys,
-            &private_account_auth,
-            &visibility_mask,
-            &private_account_nonces,
-            commitment_set_digest,
-            &program,
         )
         .unwrap();
 
@@ -172,7 +166,7 @@ mod tests {
         assert_eq!(sender_post, expected_sender_post);
         assert_eq!(output.new_commitments.len(), 1);
         assert_eq!(output.new_nullifiers.len(), 0);
-        assert_eq!(output.commitment_set_digest, commitment_set_digest);
+        assert_eq!(output.commitment_set_digest, [99; 32]);
         assert_eq!(output.encrypted_private_post_states.len(), 1);
         // TODO: replace with real assertion when encryption is implemented
         assert_eq!(output.encrypted_private_post_states[0].to_bytes(), vec![0]);
@@ -195,34 +189,19 @@ mod tests {
             account: Account::default(),
             is_authorized: false,
         };
-        let private_key_2 = [2; 32];
+        let Npk2 = NullifierPublicKey::from(&[99; 32]);
-        let Npk2 = NullifierPublicKey::from(&private_key_2);
         let balance_to_move: u128 = 37;
-        let private_account_nonces = vec![0xdeadbeef1, 0xdeadbeef2];
         let commitment_set =
             CommitmentSet(MerkleTree::new(vec![commitment_sender.to_byte_array()]));
-        let pre_states = vec![sender.clone(), recipient];
-        let instruction_data = Program::serialize_instruction(balance_to_move).unwrap();
-        let private_account_keys = vec![
-            (Npk1.clone(), [2; 32], [3; 32]),
-            (Npk2.clone(), [4; 32], [5; 32]),
-        ];
-        let private_account_auth = vec![(
-            private_key_1,
-            commitment_set.get_proof_for(&commitment_sender).unwrap(),
-        )];
-        let visibility_mask = vec![1, 2];
-        let program = Program::authenticated_transfer_program();
-        let mut commitment_set_digest = commitment_set.digest();
 
         let expected_private_account_1 = Account {
             balance: 100 - balance_to_move,
-            nonce: private_account_nonces[0],
+            nonce: 0xdeadbeef1,
             ..Default::default()
         };
         let expected_private_account_2 = Account {
             balance: balance_to_move,
-            nonce: private_account_nonces[1],
+            nonce: 0xdeadbeef2,
             ..Default::default()
         };
         let expected_new_commitments = vec![
@@ -231,15 +210,21 @@ mod tests {
         ];
         let expected_new_nullifiers = vec![Nullifier::new(&commitment_sender, &private_key_1)];
 
-        let (proof, output) = prove_privacy_preserving_execution_circuit(
+        let (output, proof) = execute_and_prove(
-            &pre_states,
+            &[sender.clone(), recipient],
-            &instruction_data,
+            &Program::serialize_instruction(balance_to_move).unwrap(),
-            &private_account_keys,
+            &[1, 2],
-            &private_account_auth,
+            &[0xdeadbeef1, 0xdeadbeef2],
-            &visibility_mask,
+            &[
-            &private_account_nonces,
+                (Npk1.clone(), [2; 32], [3; 32]),
-            commitment_set_digest,
+                (Npk2.clone(), [4; 32], [5; 32]),
-            &program,
+            ],
+            &[(
+                private_key_1,
+                commitment_set.get_proof_for(&commitment_sender).unwrap(),
+            )],
+            &Program::authenticated_transfer_program(),
+            &commitment_set.digest(),
         )
         .unwrap();
 
@@ -248,7 +233,7 @@ mod tests {
         assert!(output.public_post_states.is_empty());
         assert_eq!(output.new_commitments, expected_new_commitments);
         assert_eq!(output.new_nullifiers, expected_new_nullifiers);
-        assert_eq!(output.commitment_set_digest, commitment_set_digest);
+        assert_eq!(output.commitment_set_digest, commitment_set.digest());
         // TODO: replace with real assertion when encryption is implemented
         assert_eq!(output.encrypted_private_post_states.len(), 2);
         assert_eq!(output.encrypted_private_post_states[0].to_bytes(), vec![0]);

nssa/src/privacy_preserving_transaction/transaction.rs

@@ -165,8 +165,3 @@ fn n_unique<T: Eq + Hash>(data: &[T]) -> usize {
     let set: HashSet<&T> = data.iter().collect();
     set.len()
 }
-
-#[cfg(test)]
-mod tests {
-    use crate::privacy_preserving_transaction::message::tests::message_for_tests;
-}

nssa/src/state.rs

@@ -4,7 +4,9 @@ use crate::{
     public_transaction::PublicTransaction,
 };
 use nssa_core::{
-    account::{Account, Commitment, Nullifier}, program::{ProgramId, DEFAULT_PROGRAM_ID}, CommitmentSetDigest, MembershipProof
+    CommitmentSetDigest, MembershipProof,
+    account::{Account, Commitment, Nullifier},
+    program::{DEFAULT_PROGRAM_ID, ProgramId},
 };
 use std::collections::{HashMap, HashSet};
 
@@ -22,7 +24,8 @@ impl CommitmentSet {
     }
 
     pub fn get_proof_for(&self, commitment: &Commitment) -> Option<MembershipProof> {
-        self.0.get_authentication_path_for(&commitment.to_byte_array())
+        self.0
+            .get_authentication_path_for(&commitment.to_byte_array())
     }
 
     fn contains(&self, commitment: &Commitment) -> bool {
@@ -123,6 +126,7 @@ impl V01State {
             let current_account = self.get_account_by_address_mut(address);
             current_account.nonce += 1;
         }
+
         Ok(())
     }
 
@@ -180,10 +184,16 @@ mod tests {
     use std::collections::HashMap;
 
     use crate::{
-        Address, PublicKey, PublicTransaction, V01State, error::NssaError, program::Program,
+        Address, PublicKey, PublicTransaction, V01State,
-        public_transaction, signature::PrivateKey,
+        error::NssaError,
+        privacy_preserving_transaction::{
+            Message, PrivacyPreservingTransaction, WitnessSet, circuit::execute_and_prove,
+        },
+        program::Program,
+        public_transaction,
+        signature::PrivateKey,
     };
-    use nssa_core::account::Account;
+    use nssa_core::account::{Account, AccountWithMetadata, NullifierPublicKey};
 
     fn transfer_transaction(
         from: Address,
@@ -670,4 +680,70 @@ mod tests {
 
         assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
     }
+
+    #[test]
+    fn test_transition_from_privacy_preserving_transaction() {
+        let sender = AccountWithMetadata {
+            account: Account {
+                balance: 200,
+                program_owner: Program::authenticated_transfer_program().id(),
+                ..Account::default()
+            },
+            is_authorized: true,
+        };
+        let sender_signing_key = PrivateKey::try_new([1; 32]).unwrap();
+        let sender_address =
+            Address::from_public_key(&PublicKey::new_from_private_key(&sender_signing_key));
+
+        let recipient = AccountWithMetadata {
+            account: Account {
+                ..Account::default()
+            },
+            is_authorized: false,
+        };
+
+        let recipient_npk = NullifierPublicKey::from(&[1; 32]);
+        let recipient_ivk = [2; 32];
+        let esk = [3; 32];
+
+        let balance_to_move: u128 = 37;
+
+        let mut state = V01State::new_with_genesis_accounts(&[(*sender_address.value(), 200)]);
+
+        let (output, proof) = execute_and_prove(
+            &[sender, recipient],
+            &Program::serialize_instruction(balance_to_move).unwrap(),
+            &[0, 2],
+            &[0xdeadbeef],
+            &[(recipient_npk, recipient_ivk, esk)],
+            &[],
+            &Program::authenticated_transfer_program(),
+            &state.commitment_set_digest(),
+        )
+        .unwrap();
+
+        let message = Message::new(
+            vec![sender_address.clone()],
+            vec![0],
+            output.public_post_states,
+            output.encrypted_private_post_states,
+            output.new_commitments.clone(),
+            output.new_nullifiers,
+        );
+
+        let witness_set = WitnessSet::for_message(&message, proof, &[&sender_signing_key]);
+        let tx = PrivacyPreservingTransaction::new(message, witness_set);
+
+        assert!(!state.private_state.0.contains(&output.new_commitments[0]));
+
+        state
+            .transition_from_privacy_preserving_transaction(&tx)
+            .unwrap();
+
+        assert!(state.private_state.0.contains(&output.new_commitments[0]));
+        assert_eq!(
+            state.get_account_by_address(&sender_address).balance,
+            200 - 37
+        );
+    }
 }