diff --git a/nssa/Cargo.toml b/nssa/Cargo.toml index f479999..942834e 100644 --- a/nssa/Cargo.toml +++ b/nssa/Cargo.toml @@ -13,6 +13,7 @@ sha2 = "0.10.9" secp256k1 = "0.31.1" rand = "0.8" borsh = "1.5.7" +bytemuck = "1.13" [dev-dependencies] test-program-methods = { path = "test_program_methods" } diff --git a/nssa/core/src/lib.rs b/nssa/core/src/lib.rs index a616237..56de12c 100644 --- a/nssa/core/src/lib.rs +++ b/nssa/core/src/lib.rs @@ -1,4 +1,7 @@ -use risc0_zkvm::serde::to_vec; +use risc0_zkvm::{ + serde::to_vec, + sha::{Impl, Sha256}, +}; use serde::{Deserialize, Serialize}; #[cfg(feature = "host")] @@ -21,15 +24,35 @@ pub mod program; #[cfg(feature = "host")] pub mod error; -pub type CommitmentSetDigest = [u32; 8]; -pub type MembershipProof = Vec<[u8; 32]>; +pub type CommitmentSetDigest = [u8; 32]; +pub type MembershipProof = (usize, Vec<[u8; 32]>); pub fn verify_membership_proof( commitment: &Commitment, proof: &MembershipProof, digest: &CommitmentSetDigest, ) -> bool { - // TODO: implement - true + let value_bytes = commitment.to_byte_array(); + let mut result: [u8; 32] = Impl::hash_bytes(&value_bytes) + .as_bytes() + .try_into() + .unwrap(); + let mut level_index = proof.0; + for node in &proof.1 { + let is_left_child = level_index & 1 == 0; + if is_left_child { + let mut bytes = [0u8; 64]; + bytes[..32].copy_from_slice(&result); + bytes[32..].copy_from_slice(node); + result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap(); + } else { + let mut bytes = [0u8; 64]; + bytes[..32].copy_from_slice(node); + bytes[32..].copy_from_slice(&result); + result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap(); + } + level_index >>= 1; + } + &result == digest } pub type IncomingViewingPublicKey = [u8; 32]; @@ -158,7 +181,7 @@ mod tests { &Commitment::new(&NullifierPublicKey::from(&[2; 32]), &Account::default()), &[1; 32], )], - commitment_set_digest: [0, 1, 0, 1, 0, 1, 0, 1], + commitment_set_digest: [0xab; 32], }; let bytes = output.to_bytes(); let output_from_slice: PrivacyPreservingCircuitOutput = from_slice(&bytes).unwrap(); diff --git a/nssa/src/merkle_tree.rs b/nssa/src/merkle_tree.rs index 697cd7c..9126150 100644 --- a/nssa/src/merkle_tree.rs +++ b/nssa/src/merkle_tree.rs @@ -93,7 +93,7 @@ impl MerkleTree { let base_length = self.capacity; let mut layer_node = hash_value(&value); let mut layer_index = new_index + base_length - 1; - self.node_map.insert(layer_index, layer_node); + self.set_node(layer_index, layer_node); let mut layer = 0; let mut top_layer = self.depth(); @@ -159,6 +159,10 @@ impl MerkleTree { } Some((*value_index, result)) } + + pub(crate) fn contains(&self, value: &[u8; 32]) -> bool { + self.index_map.contains_key(value) + } } // Reference implementation @@ -518,7 +522,6 @@ mod tests { assert!(tree.get_authentication_path_for(&value).is_none()); } - #[test] fn test_authentication_path_5() { let values = vec![[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; diff --git a/nssa/src/privacy_preserving_transaction/mod.rs b/nssa/src/privacy_preserving_transaction/mod.rs index 17fe677..f4e2002 100644 --- a/nssa/src/privacy_preserving_transaction/mod.rs +++ b/nssa/src/privacy_preserving_transaction/mod.rs @@ -113,14 +113,16 @@ pub mod circuit { mod tests { use nssa_core::{ EncryptedAccountData, - account::{Account, AccountWithMetadata, NullifierPublicKey, NullifierSecretKey}, + account::{ + Account, AccountWithMetadata, Commitment, NullifierPublicKey, NullifierSecretKey, + }, }; use risc0_zkvm::{InnerReceipt, Journal, Receipt}; use crate::{ - Address, V01State, + Address, V01State, merkle_tree::MerkleTree, privacy_preserving_transaction::circuit::prove_privacy_preserving_execution_circuit, - program::Program, + program::Program, state::CommitmentSet, }; use super::*; @@ -153,7 +155,7 @@ mod tests { let private_account_keys = vec![(NullifierPublicKey::from(&[1; 32]), [2; 32], [3; 32])]; let private_account_auth = vec![]; let visibility_mask = vec![0, 2]; - let commitment_set_digest = [99; 8]; + let commitment_set_digest = [99; 32]; let program = Program::authenticated_transfer_program(); let (proof, output) = prove_privacy_preserving_execution_circuit( &pre_states, @@ -190,6 +192,10 @@ mod tests { is_authorized: true, }; + let private_key = [1; 32]; + let Npk = NullifierPublicKey::from(&private_key); + let commitment = Commitment::new(&Npk, &sender.account); + let recipient = AccountWithMetadata { account: Account::default(), is_authorized: false, @@ -197,19 +203,22 @@ mod tests { let balance_to_move: u128 = 37; + let commitment_set = CommitmentSet(MerkleTree::new(vec![commitment.to_byte_array()])); + let expected_sender_pre = sender.clone(); - let pre_states = vec![sender, recipient]; + let pre_states = vec![sender.clone(), recipient]; let instruction_data = Program::serialize_instruction(balance_to_move).unwrap(); - let private_key = [1; 32]; let private_account_keys = vec![ - (NullifierPublicKey::from(&private_key), [2; 32], [3; 32]), + (Npk.clone(), [2; 32], [3; 32]), (NullifierPublicKey::from(&[2; 32]), [4; 32], [5; 32]), ]; - // TODO: Replace dummy authentication path when implemented - let private_account_auth = vec![(private_key, vec![])]; + let private_account_auth = vec![( + private_key, + commitment_set.get_proof_for(&commitment).unwrap(), + )]; let visibility_mask = vec![1, 2]; - let commitment_set_digest = [99; 8]; let program = Program::authenticated_transfer_program(); + let mut commitment_set_digest = commitment_set.digest(); let (proof, output) = prove_privacy_preserving_execution_circuit( &pre_states, &instruction_data, diff --git a/nssa/src/state.rs b/nssa/src/state.rs index 101fa25..5a94bc3 100644 --- a/nssa/src/state.rs +++ b/nssa/src/state.rs @@ -1,29 +1,32 @@ use crate::{ - address::Address, error::NssaError, + address::Address, error::NssaError, merkle_tree::MerkleTree, privacy_preserving_transaction::PrivacyPreservingTransaction, program::Program, public_transaction::PublicTransaction, }; use nssa_core::{ - CommitmentSetDigest, - account::{Account, Commitment, Nullifier}, - program::{DEFAULT_PROGRAM_ID, ProgramId}, + account::{Account, Commitment, Nullifier}, program::{ProgramId, DEFAULT_PROGRAM_ID}, CommitmentSetDigest, MembershipProof }; use std::collections::{HashMap, HashSet}; -struct CommitmentSet(HashSet); +pub(crate) struct CommitmentSet(pub(crate) MerkleTree); impl CommitmentSet { - fn extend(&mut self, commitments: Vec) { - self.0.extend(commitments) + fn extend(&mut self, commitments: &[Commitment]) { + for commitment in commitments { + self.0.insert(commitment.to_byte_array()); + } } - fn digest(&self) -> CommitmentSetDigest { - // TODO: implement - [0; 8] + pub fn digest(&self) -> CommitmentSetDigest { + self.0.root() + } + + pub fn get_proof_for(&self, commitment: &Commitment) -> Option { + self.0.get_authentication_path_for(&commitment.to_byte_array()) } fn contains(&self, commitment: &Commitment) -> bool { - self.0.contains(commitment) + self.0.contains(&commitment.to_byte_array()) } } @@ -54,7 +57,10 @@ impl V01State { let mut this = Self { public_state, - private_state: (CommitmentSet(HashSet::new()), NullifierSet::new()), + private_state: ( + CommitmentSet(MerkleTree::with_capacity(32)), + NullifierSet::new(), + ), builtin_programs: HashMap::new(), }; @@ -101,7 +107,7 @@ impl V01State { let message = tx.message(); // 2. Add new commitments - self.private_state.0.extend(message.new_commitments.clone()); + self.private_state.0.extend(&message.new_commitments); // 3. Add new nullifiers self.private_state.1.extend(message.new_nullifiers.clone());