diff --git a/nssa/core/src/account/commitment.rs b/nssa/core/src/account/commitment.rs index 60034cb..487789f 100644 --- a/nssa/core/src/account/commitment.rs +++ b/nssa/core/src/account/commitment.rs @@ -6,7 +6,8 @@ use serde::{Deserialize, Serialize}; use crate::account::{Account, NullifierPublicKey}; -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] +#[derive(Serialize, Deserialize)] +#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq, Hash))] pub struct Commitment(pub(super) [u8; 32]); impl Commitment { diff --git a/nssa/core/src/account/mod.rs b/nssa/core/src/account/mod.rs index 2bb2285..6f5e017 100644 --- a/nssa/core/src/account/mod.rs +++ b/nssa/core/src/account/mod.rs @@ -14,7 +14,8 @@ pub type Nonce = u128; type Data = Vec; /// Account to be used both in public and private contexts -#[derive(Default, Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Serialize, Deserialize, Clone, Default, PartialEq, Eq)] +#[cfg_attr(any(feature = "host", test), derive(Debug))] pub struct Account { pub program_owner: ProgramId, pub balance: u128, @@ -22,7 +23,8 @@ pub struct Account { pub nonce: Nonce, } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Serialize, Deserialize, Clone)] +#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct AccountWithMetadata { pub account: Account, pub is_authorized: bool, diff --git a/nssa/core/src/account/nullifier.rs b/nssa/core/src/account/nullifier.rs index bb83854..8faef0b 100644 --- a/nssa/core/src/account/nullifier.rs +++ b/nssa/core/src/account/nullifier.rs @@ -3,7 +3,8 @@ use serde::{Deserialize, Serialize}; use crate::account::Commitment; -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, Hash))] pub struct NullifierPublicKey(pub(super) [u8; 32]); impl From<&NullifierSecretKey> for NullifierPublicKey { @@ -22,7 +23,8 @@ impl From<&NullifierSecretKey> for NullifierPublicKey { pub type NullifierSecretKey = [u8; 32]; -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] +#[derive(Serialize, Deserialize)] +#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq, Hash))] pub struct Nullifier(pub(super) [u8; 32]); impl Nullifier { diff --git a/nssa/core/src/lib.rs b/nssa/core/src/lib.rs index dd2e893..3003fef 100644 --- a/nssa/core/src/lib.rs +++ b/nssa/core/src/lib.rs @@ -48,7 +48,8 @@ impl Tag { } } -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] +#[derive(Serialize, Deserialize)] +#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq))] pub struct EncryptedAccountData(u8); impl EncryptedAccountData { @@ -100,3 +101,15 @@ pub struct PrivacyPreservingCircuitOutput { pub new_nullifiers: Vec, pub commitment_set_digest: CommitmentSetDigest, } + +#[cfg(feature = "host")] +impl PrivacyPreservingCircuitOutput { + pub fn to_bytes(&self) -> Vec { + let words = to_vec(&self).unwrap(); + let mut result = Vec::with_capacity(4 * words.len()); + for word in &words { + result.extend_from_slice(&word.to_le_bytes()); + } + result + } +} diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs index b942536..f84bd61 100644 --- a/nssa/core/src/program.rs +++ b/nssa/core/src/program.rs @@ -15,7 +15,8 @@ pub struct ProgramInput { pub instruction: T, } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Serialize, Deserialize, Clone)] +#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct ProgramOutput { pub pre_states: Vec, pub post_states: Vec, diff --git a/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs b/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs index f165d86..f61d74a 100644 --- a/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs +++ b/nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs @@ -34,14 +34,22 @@ fn main() { validate_execution(&pre_states, &post_states, program_id); let n_accounts = pre_states.len(); - assert_eq!(visibility_mask.len(), n_accounts); + if visibility_mask.len() != n_accounts { + panic!(); + } let n_private_accounts = visibility_mask.iter().filter(|&&flag| flag != 0).count(); - assert_eq!(private_account_nonces.len(), n_private_accounts); - assert_eq!(private_account_keys.len(), n_private_accounts); + if private_account_nonces.len() != n_private_accounts { + panic!(); + } + if private_account_keys.len() != n_private_accounts { + panic!(); + } let n_auth_private_accounts = visibility_mask.iter().filter(|&&flag| flag == 1).count(); - assert_eq!(private_account_auth.len(), n_auth_private_accounts); + if private_account_auth.len() != n_auth_private_accounts { + panic!(); + } // These lists will be the public outputs of this circuit // and will be populated next. @@ -51,60 +59,71 @@ fn main() { let mut new_commitments: Vec = Vec::new(); let mut new_nullifiers: Vec = Vec::new(); + let mut private_nonces_iter = private_account_nonces.iter(); + let mut private_keys_iter = private_account_keys.iter(); + let mut private_auth_iter = private_account_auth.iter(); + for i in 0..n_accounts { - // visibility_mask[i] equal to 0 means public - if visibility_mask[i] == 0 { - // If the account is marked as public, add the pre and post - // states to the corresponding lists. - public_pre_states.push(pre_states[i].clone()); - public_post_states.push(post_states[i].clone()); - } else { - let new_nonce = &private_account_nonces[i]; - let (Npk, Ipk, esk) = &private_account_keys[i]; - - // Verify authentication - if visibility_mask[i] == 1 { - let (nsk, membership_proof) = &private_account_auth[i]; - - // 1. Compute Npk from the provided nsk and assert it is equal to the provided Npk - let expected_Npk = NullifierPublicKey::from(nsk); - assert_eq!(&expected_Npk, Npk); - // 2. Compute the commitment of the pre_state account using the provided Npk - let commitment_pre = Commitment::new(Npk, &pre_states[i].account); - // 3. Verify that the commitment belongs to the global commitment set - assert!(verify_membership_proof( - &commitment_pre, - membership_proof, - &commitment_set_digest, - )); - // At this point the account is correctly authenticated as a private account. - // Assert that `pre_states` marked this account as authenticated. - assert!(pre_states[i].is_authorized); - // Compute the nullifier of the pre state version of this private account - // and include it in the `new_nullifiers` list. - let nullifier = Nullifier::new(&commitment_pre, nsk); - new_nullifiers.push(nullifier); - } else if visibility_mask[i] == 2 { - assert_eq!(pre_states[i].account, Account::default()); - assert!(!pre_states[i].is_authorized); - } else { - panic!(); + match visibility_mask[i] { + 0 => { + // Public account + public_pre_states.push(pre_states[i].clone()); + public_post_states.push(post_states[i].clone()); } + 1 | 2 => { + let new_nonce = private_nonces_iter.next().expect("Missing private nonce"); + let (Npk, Ipk, esk) = private_keys_iter.next().expect("Missing private keys"); - // Update the nonce for the post state of this private account. - let mut post_with_updated_nonce = post_states[i].clone(); - post_with_updated_nonce.nonce = *new_nonce; + if visibility_mask[i] == 1 { + // Private account with authentication + let (nsk, membership_proof) = + private_auth_iter.next().expect("Missing private auth"); - // Compute the commitment of the post state of the private account, - // with the updated nonce, and include it in the `new_commitments` list. - let commitment_post = Commitment::new(Npk, &post_with_updated_nonce); - new_commitments.push(commitment_post); + // Verify Npk + let expected_Npk = NullifierPublicKey::from(nsk); + if &expected_Npk != Npk { + panic!("Npk mismatch"); + } - // Encrypt the post state of the private account with the updated - // nonce and include it in the `encrypted_private_post_states` list. - // - let encrypted_account = EncryptedAccountData::new(&post_with_updated_nonce, esk, Npk, Ipk); - encrypted_private_post_states.push(encrypted_account); + // Verify pre-state commitment membership + let commitment_pre = Commitment::new(Npk, &pre_states[i].account); + if !verify_membership_proof( + &commitment_pre, + membership_proof, + &commitment_set_digest, + ) { + panic!("Membership proof invalid"); + } + + // Check pre_state authorization + if !pre_states[i].is_authorized { + panic!("Pre-state not authorized"); + } + + // Compute nullifier + let nullifier = Nullifier::new(&commitment_pre, nsk); + new_nullifiers.push(nullifier); + } else { + // Private account marked as empty + if pre_states[i].account != Account::default() || pre_states[i].is_authorized { + panic!("Invalid empty private account pre-state"); + } + } + + // Update post-state with new nonce + let mut post_with_updated_nonce = post_states[i].clone(); + post_with_updated_nonce.nonce = *new_nonce; + + // Compute commitment and push + let commitment_post = Commitment::new(Npk, &post_with_updated_nonce); + new_commitments.push(commitment_post); + + // Encrypt and push post state + let encrypted_account = + EncryptedAccountData::new(&post_with_updated_nonce, esk, Npk, Ipk); + encrypted_private_post_states.push(encrypted_account); + } + _ => panic!("Invalid visibility mask value"), } } diff --git a/nssa/src/privacy_preserving_transaction/mod.rs b/nssa/src/privacy_preserving_transaction/mod.rs index 3618e3e..0caba33 100644 --- a/nssa/src/privacy_preserving_transaction/mod.rs +++ b/nssa/src/privacy_preserving_transaction/mod.rs @@ -3,6 +3,7 @@ mod message; mod transaction; mod witness_set; +pub use message::Message; pub use transaction::PrivacyPreservingTransaction; pub mod circuit { @@ -12,6 +13,7 @@ pub mod circuit { account::{Account, AccountWithMetadata, Nonce, NullifierPublicKey, NullifierSecretKey}, program::{InstructionData, ProgramOutput}, }; + use rand::{Rng, RngCore, rngs::OsRng}; use risc0_zkvm::{ExecutorEnv, Receipt, default_prover}; use crate::{error::NssaError, program::Program}; @@ -48,7 +50,7 @@ pub mod circuit { IncomingViewingPublicKey, EphemeralSecretKey, )], - private_account_auth: Vec<(NullifierSecretKey, MembershipProof)>, + private_account_auth: &[(NullifierSecretKey, MembershipProof)], visibility_mask: &[u8], commitment_set_digest: CommitmentSetDigest, program: &Program, @@ -93,13 +95,83 @@ pub mod circuit { Ok((proof, circuit_output)) } - fn new_random_nonce() -> Nonce { - todo!() + fn new_random_nonce() -> u128 { + let mut u128_bytes = [0u8; 16]; + OsRng.fill_bytes(&mut u128_bytes); + u128::from_le_bytes(u128_bytes) } } #[cfg(test)] mod tests { - use crate::program::Program; + use nssa_core::{ + EncryptedAccountData, + account::{Account, AccountWithMetadata, NullifierPublicKey, NullifierSecretKey}, + }; + use program_methods::PRIVACY_PRESERVING_CIRCUIT_ID; + use risc0_zkvm::{InnerReceipt, Journal, Receipt}; + use crate::{ + Address, V01State, + privacy_preserving_transaction::circuit::prove_privacy_preserving_execution_circuit, + program::Program, + }; + + use super::*; + + #[test] + fn test() { + let sender = AccountWithMetadata { + account: Account { + balance: 100, + ..Account::default() + }, + is_authorized: false, + }; + let recipient = AccountWithMetadata { + account: Account::default(), + is_authorized: false, + }; + + let balance_to_move: u128 = 37; + + let expected_sender_post = Account { + balance: 100 - balance_to_move, + ..Default::default() + }; + + 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_auth = vec![]; + let visibility_mask = vec![0, 2]; + let commitment_set_digest = [99; 8]; + let program = Program::simple_balance_transfer(); + let (proof, output) = prove_privacy_preserving_execution_circuit( + &pre_states, + &instruction_data, + &private_account_keys, + &private_account_auth, + &visibility_mask, + commitment_set_digest, + &program, + ) + .unwrap(); + + let inner: InnerReceipt = borsh::from_slice(&proof).unwrap(); + let receipt = Receipt::new(inner, output.to_bytes()); + receipt.verify(PRIVACY_PRESERVING_CIRCUIT_ID).unwrap(); + + let [sender_pre] = output.public_pre_states.try_into().unwrap(); + let [sender_post] = output.public_post_states.try_into().unwrap(); + assert_eq!(sender_pre, expected_sender_pre); + 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.encrypted_private_post_states.len(), 1); + // TODO: replace with real assert when encryption is implemented + assert_eq!(output.encrypted_private_post_states[0].to_bytes(), vec![0]); + } }