refactor. Add test for state transition from privacy preserving transaction

This commit is contained in:
Sergio Chouhy 2025-08-21 14:32:57 -03:00
parent a5dc01d85e
commit b6cabe8fb8
4 changed files with 176 additions and 116 deletions

View File

@ -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();
post_with_updated_nonce.nonce = *new_nonce;
let mut post_with_updated_values = post_states[i].clone();
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"),

View File

@ -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,
privacy_preserving_transaction::circuit::prove_privacy_preserving_execution_circuit,
program::Program, state::CommitmentSet,
Address, V01State,
merkle_tree::MerkleTree,
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 instruction_data = Program::serialize_instruction(balance_to_move).unwrap();
let private_account_keys = vec![(NullifierPublicKey::from(&[1; 32]), [2; 32], [3; 32])];
let private_account_nonces = vec![0xdeadbeef];
let private_account_auth = vec![];
let visibility_mask = vec![0, 2];
let commitment_set_digest = [99; 32];
let program = Program::authenticated_transfer_program();
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,
let (output, proof) = execute_and_prove(
&[sender, recipient],
&Program::serialize_instruction(balance_to_move).unwrap(),
&[0, 2],
&[0xdeadbeef],
&[(NullifierPublicKey::from(&[1; 32]), [2; 32], [3; 32])],
&[],
&Program::authenticated_transfer_program(),
&[99; 32],
)
.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(&private_key_2);
let Npk2 = NullifierPublicKey::from(&[99; 32]);
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(
&pre_states,
&instruction_data,
&private_account_keys,
&private_account_auth,
&visibility_mask,
&private_account_nonces,
commitment_set_digest,
&program,
let (output, proof) = execute_and_prove(
&[sender.clone(), recipient],
&Program::serialize_instruction(balance_to_move).unwrap(),
&[1, 2],
&[0xdeadbeef1, 0xdeadbeef2],
&[
(Npk1.clone(), [2; 32], [3; 32]),
(Npk2.clone(), [4; 32], [5; 32]),
],
&[(
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]);

View File

@ -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;
}

View File

@ -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,
public_transaction, signature::PrivateKey,
Address, PublicKey, PublicTransaction, V01State,
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
);
}
}