From bda50f1d2fa8e868fbd1ac75767a8486163b6e70 Mon Sep 17 00:00:00 2001 From: moudyellaz Date: Thu, 7 May 2026 16:54:57 +0200 Subject: [PATCH] refactor(privacy_preserving_circuit): extract output module Refs: #454 --- .../bin/privacy_preserving_circuit/main.rs | 297 ++---------------- .../bin/privacy_preserving_circuit/output.rs | 283 +++++++++++++++++ 2 files changed, 301 insertions(+), 279 deletions(-) create mode 100644 program_methods/guest/src/bin/privacy_preserving_circuit/output.rs diff --git a/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs b/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs index f658ea53..6f245937 100644 --- a/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs +++ b/program_methods/guest/src/bin/privacy_preserving_circuit/main.rs @@ -4,11 +4,8 @@ use std::{ }; use nssa_core::{ - Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Identifier, - InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey, - PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, SharedSecretKey, - account::{Account, AccountId, AccountWithMetadata, Nonce}, - compute_digest_for_path, + InputAccountIdentity, NullifierPublicKey, PrivacyPreservingCircuitInput, + account::{Account, AccountId, AccountWithMetadata}, program::{ AccountPostState, BlockValidityWindow, ChainedCall, Claim, DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, PdaSeed, ProgramId, ProgramOutput, TimestampValidityWindow, @@ -17,7 +14,7 @@ use nssa_core::{ }; use risc0_zkvm::{guest::env, serde::to_vec}; -const PRIVATE_PDA_FIXED_IDENTIFIER: Identifier = u128::MAX; +mod output; /// State of the involved accounts before and after program execution. struct ExecutionState { @@ -386,17 +383,26 @@ impl ExecutionState { } } - /// Get an iterator over pre and post states of each account involved in the execution. - pub fn into_states_iter( + /// Consume self and yield the validity windows alongside an iterator over pre and post + /// states of each account involved in the execution. Returning the windows here keeps the + /// fields module-private rather than forcing them visible to downstream consumers. + pub fn into_parts( mut self, - ) -> impl ExactSizeIterator { - self.pre_states.into_iter().map(move |pre| { + ) -> ( + BlockValidityWindow, + TimestampValidityWindow, + impl ExactSizeIterator, + ) { + let block_validity_window = self.block_validity_window; + let timestamp_validity_window = self.timestamp_validity_window; + let states_iter = self.pre_states.into_iter().map(move |pre| { let post = self .post_states .remove(&pre.account_id) .expect("Account from pre states should exist in state diff"); (pre, post) - }) + }); + (block_validity_window, timestamp_validity_window, states_iter) } } @@ -476,273 +482,6 @@ fn resolve_authorization_and_record_bindings( previous_is_authorized || matched_caller_seed.is_some() } -fn compute_circuit_output( - execution_state: ExecutionState, - account_identities: &[InputAccountIdentity], -) -> PrivacyPreservingCircuitOutput { - let mut output = PrivacyPreservingCircuitOutput { - public_pre_states: Vec::new(), - public_post_states: Vec::new(), - ciphertexts: Vec::new(), - new_commitments: Vec::new(), - new_nullifiers: Vec::new(), - block_validity_window: execution_state.block_validity_window, - timestamp_validity_window: execution_state.timestamp_validity_window, - }; - - let states_iter = execution_state.into_states_iter(); - assert_eq!( - account_identities.len(), - states_iter.len(), - "Invalid account_identities length" - ); - - let mut output_index = 0; - for (account_identity, (pre_state, post_state)) in account_identities.iter().zip(states_iter) { - match account_identity { - InputAccountIdentity::Public => { - output.public_pre_states.push(pre_state); - output.public_post_states.push(post_state); - } - InputAccountIdentity::PrivateAuthorizedInit { - ssk, - nsk, - identifier, - } => { - assert_ne!( - *identifier, PRIVATE_PDA_FIXED_IDENTIFIER, - "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA." - ); - let npk = NullifierPublicKey::from(nsk); - let account_id = AccountId::from((&npk, *identifier)); - - assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); - assert!( - pre_state.is_authorized, - "Pre-state not authorized for authenticated private account" - ); - assert_eq!( - pre_state.account, - Account::default(), - "Found new private account with non default values" - ); - - let new_nullifier = ( - Nullifier::for_account_initialization(&account_id), - DUMMY_COMMITMENT_HASH, - ); - let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); - - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - *identifier, - ssk, - new_nullifier, - new_nonce, - ); - } - InputAccountIdentity::PrivateAuthorizedUpdate { - ssk, - nsk, - membership_proof, - identifier, - } => { - assert_ne!( - *identifier, PRIVATE_PDA_FIXED_IDENTIFIER, - "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA." - ); - let npk = NullifierPublicKey::from(nsk); - let account_id = AccountId::from((&npk, *identifier)); - - assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); - assert!( - pre_state.is_authorized, - "Pre-state not authorized for authenticated private account" - ); - - let new_nullifier = compute_update_nullifier_and_set_digest( - membership_proof, - &pre_state.account, - &account_id, - nsk, - ); - let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); - - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - *identifier, - ssk, - new_nullifier, - new_nonce, - ); - } - InputAccountIdentity::PrivateUnauthorized { - npk, - ssk, - identifier, - } => { - assert_ne!( - *identifier, PRIVATE_PDA_FIXED_IDENTIFIER, - "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA." - ); - let account_id = AccountId::from((npk, *identifier)); - - assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); - assert_eq!( - pre_state.account, - Account::default(), - "Found new private account with non default values", - ); - assert!( - !pre_state.is_authorized, - "Found new private account marked as authorized." - ); - - let new_nullifier = ( - Nullifier::for_account_initialization(&account_id), - DUMMY_COMMITMENT_HASH, - ); - let new_nonce = Nonce::private_account_nonce_init(&account_id); - - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - *identifier, - ssk, - new_nullifier, - new_nonce, - ); - } - InputAccountIdentity::PrivatePdaInit { npk: _, ssk } => { - // The npk-to-account_id binding is established upstream in - // `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds` - // match. Here we only enforce the init pre-conditions. The supplied npk on - // the variant has been recorded into `private_pda_npk_by_position` and used - // for the binding check; we use `pre_state.account_id` directly for nullifier - // and commitment derivation. - assert!( - !pre_state.is_authorized, - "PrivatePdaInit requires unauthorized pre_state" - ); - assert_eq!( - pre_state.account, - Account::default(), - "New private PDA must be default" - ); - - let new_nullifier = ( - Nullifier::for_account_initialization(&pre_state.account_id), - DUMMY_COMMITMENT_HASH, - ); - let new_nonce = Nonce::private_account_nonce_init(&pre_state.account_id); - - let account_id = pre_state.account_id; - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - PRIVATE_PDA_FIXED_IDENTIFIER, - ssk, - new_nullifier, - new_nonce, - ); - } - InputAccountIdentity::PrivatePdaUpdate { - ssk, - nsk, - membership_proof, - } => { - // The npk binding is established upstream. Authorization must already be set; - // an unauthorized PrivatePdaUpdate would mean the prover supplied an nsk for an - // unbound PDA, which the upstream binding check would have rejected anyway, - // but we assert here to fail fast and document the precondition. - assert!( - pre_state.is_authorized, - "PrivatePdaUpdate requires authorized pre_state" - ); - - let new_nullifier = compute_update_nullifier_and_set_digest( - membership_proof, - &pre_state.account, - &pre_state.account_id, - nsk, - ); - let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); - - let account_id = pre_state.account_id; - emit_private_output( - &mut output, - &mut output_index, - post_state, - &account_id, - PRIVATE_PDA_FIXED_IDENTIFIER, - ssk, - new_nullifier, - new_nonce, - ); - } - } - } - - output -} - -#[expect( - clippy::too_many_arguments, - reason = "All seven inputs are distinct concerns from the variant arms; bundling would be artificial" -)] -fn emit_private_output( - output: &mut PrivacyPreservingCircuitOutput, - output_index: &mut u32, - post_state: Account, - account_id: &AccountId, - identifier: Identifier, - shared_secret: &SharedSecretKey, - new_nullifier: (Nullifier, CommitmentSetDigest), - new_nonce: Nonce, -) { - output.new_nullifiers.push(new_nullifier); - - let mut post_with_updated_nonce = post_state; - post_with_updated_nonce.nonce = new_nonce; - - let commitment_post = Commitment::new(account_id, &post_with_updated_nonce); - let encrypted_account = EncryptionScheme::encrypt( - &post_with_updated_nonce, - identifier, - shared_secret, - &commitment_post, - *output_index, - ); - - output.new_commitments.push(commitment_post); - output.ciphertexts.push(encrypted_account); - *output_index = output_index - .checked_add(1) - .unwrap_or_else(|| panic!("Too many private accounts, output index overflow")); -} - -fn compute_update_nullifier_and_set_digest( - membership_proof: &MembershipProof, - pre_account: &Account, - account_id: &AccountId, - nsk: &NullifierSecretKey, -) -> (Nullifier, CommitmentSetDigest) { - let commitment_pre = Commitment::new(account_id, pre_account); - let set_digest = compute_digest_for_path(&commitment_pre, membership_proof); - let nullifier = Nullifier::for_account_update(&commitment_pre, nsk); - (nullifier, set_digest) -} - fn main() { let PrivacyPreservingCircuitInput { program_outputs, @@ -753,7 +492,7 @@ fn main() { let execution_state = ExecutionState::derive_from_outputs(&account_identities, program_id, program_outputs); - let output = compute_circuit_output(execution_state, &account_identities); + let output = output::compute_circuit_output(execution_state, &account_identities); env::commit(&output); } diff --git a/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs b/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs new file mode 100644 index 00000000..6d4962d5 --- /dev/null +++ b/program_methods/guest/src/bin/privacy_preserving_circuit/output.rs @@ -0,0 +1,283 @@ +use nssa_core::{ + Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Identifier, + InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey, + PrivacyPreservingCircuitOutput, SharedSecretKey, + account::{Account, AccountId, Nonce}, + compute_digest_for_path, +}; + +use super::ExecutionState; + +// SECURITY: the non-PDA private variants below assert that the prover-supplied `identifier` is +// not equal to this constant; the PDA variants pass it as the fixed identifier. This keeps the +// `(npk, identifier)` account-id space disjoint from private-PDA accounts. Single source of +// truth, do not redefine in another module. +const PRIVATE_PDA_FIXED_IDENTIFIER: Identifier = u128::MAX; + +pub(super) fn compute_circuit_output( + execution_state: ExecutionState, + account_identities: &[InputAccountIdentity], +) -> PrivacyPreservingCircuitOutput { + let (block_validity_window, timestamp_validity_window, states_iter) = + execution_state.into_parts(); + let mut output = PrivacyPreservingCircuitOutput { + public_pre_states: Vec::new(), + public_post_states: Vec::new(), + ciphertexts: Vec::new(), + new_commitments: Vec::new(), + new_nullifiers: Vec::new(), + block_validity_window, + timestamp_validity_window, + }; + + assert_eq!( + account_identities.len(), + states_iter.len(), + "Invalid account_identities length" + ); + + let mut output_index = 0; + for (account_identity, (pre_state, post_state)) in account_identities.iter().zip(states_iter) { + match account_identity { + InputAccountIdentity::Public => { + output.public_pre_states.push(pre_state); + output.public_post_states.push(post_state); + } + InputAccountIdentity::PrivateAuthorizedInit { + ssk, + nsk, + identifier, + } => { + assert_ne!( + *identifier, PRIVATE_PDA_FIXED_IDENTIFIER, + "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA." + ); + let npk = NullifierPublicKey::from(nsk); + let account_id = AccountId::from((&npk, *identifier)); + + assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); + assert!( + pre_state.is_authorized, + "Pre-state not authorized for authenticated private account" + ); + assert_eq!( + pre_state.account, + Account::default(), + "Found new private account with non default values" + ); + + let new_nullifier = ( + Nullifier::for_account_initialization(&account_id), + DUMMY_COMMITMENT_HASH, + ); + let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); + + emit_private_output( + &mut output, + &mut output_index, + post_state, + &account_id, + *identifier, + ssk, + new_nullifier, + new_nonce, + ); + } + InputAccountIdentity::PrivateAuthorizedUpdate { + ssk, + nsk, + membership_proof, + identifier, + } => { + assert_ne!( + *identifier, PRIVATE_PDA_FIXED_IDENTIFIER, + "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA." + ); + let npk = NullifierPublicKey::from(nsk); + let account_id = AccountId::from((&npk, *identifier)); + + assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); + assert!( + pre_state.is_authorized, + "Pre-state not authorized for authenticated private account" + ); + + let new_nullifier = compute_update_nullifier_and_set_digest( + membership_proof, + &pre_state.account, + &account_id, + nsk, + ); + let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); + + emit_private_output( + &mut output, + &mut output_index, + post_state, + &account_id, + *identifier, + ssk, + new_nullifier, + new_nonce, + ); + } + InputAccountIdentity::PrivateUnauthorized { + npk, + ssk, + identifier, + } => { + assert_ne!( + *identifier, PRIVATE_PDA_FIXED_IDENTIFIER, + "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA." + ); + let account_id = AccountId::from((npk, *identifier)); + + assert_eq!(account_id, pre_state.account_id, "AccountId mismatch"); + assert_eq!( + pre_state.account, + Account::default(), + "Found new private account with non default values", + ); + assert!( + !pre_state.is_authorized, + "Found new private account marked as authorized." + ); + + let new_nullifier = ( + Nullifier::for_account_initialization(&account_id), + DUMMY_COMMITMENT_HASH, + ); + let new_nonce = Nonce::private_account_nonce_init(&account_id); + + emit_private_output( + &mut output, + &mut output_index, + post_state, + &account_id, + *identifier, + ssk, + new_nullifier, + new_nonce, + ); + } + InputAccountIdentity::PrivatePdaInit { npk: _, ssk } => { + // The npk-to-account_id binding is established upstream in + // `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds` + // match. Here we only enforce the init pre-conditions. The supplied npk on + // the variant has been recorded into `private_pda_npk_by_position` and used + // for the binding check; we use `pre_state.account_id` directly for nullifier + // and commitment derivation. + assert!( + !pre_state.is_authorized, + "PrivatePdaInit requires unauthorized pre_state" + ); + assert_eq!( + pre_state.account, + Account::default(), + "New private PDA must be default" + ); + + let new_nullifier = ( + Nullifier::for_account_initialization(&pre_state.account_id), + DUMMY_COMMITMENT_HASH, + ); + let new_nonce = Nonce::private_account_nonce_init(&pre_state.account_id); + + let account_id = pre_state.account_id; + emit_private_output( + &mut output, + &mut output_index, + post_state, + &account_id, + PRIVATE_PDA_FIXED_IDENTIFIER, + ssk, + new_nullifier, + new_nonce, + ); + } + InputAccountIdentity::PrivatePdaUpdate { + ssk, + nsk, + membership_proof, + } => { + // The npk binding is established upstream. Authorization must already be set; + // an unauthorized PrivatePdaUpdate would mean the prover supplied an nsk for an + // unbound PDA, which the upstream binding check would have rejected anyway, + // but we assert here to fail fast and document the precondition. + assert!( + pre_state.is_authorized, + "PrivatePdaUpdate requires authorized pre_state" + ); + + let new_nullifier = compute_update_nullifier_and_set_digest( + membership_proof, + &pre_state.account, + &pre_state.account_id, + nsk, + ); + let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk); + + let account_id = pre_state.account_id; + emit_private_output( + &mut output, + &mut output_index, + post_state, + &account_id, + PRIVATE_PDA_FIXED_IDENTIFIER, + ssk, + new_nullifier, + new_nonce, + ); + } + } + } + + output +} + +#[expect( + clippy::too_many_arguments, + reason = "All seven inputs are distinct concerns from the variant arms; bundling would be artificial" +)] +fn emit_private_output( + output: &mut PrivacyPreservingCircuitOutput, + output_index: &mut u32, + post_state: Account, + account_id: &AccountId, + identifier: Identifier, + shared_secret: &SharedSecretKey, + new_nullifier: (Nullifier, CommitmentSetDigest), + new_nonce: Nonce, +) { + output.new_nullifiers.push(new_nullifier); + + let mut post_with_updated_nonce = post_state; + post_with_updated_nonce.nonce = new_nonce; + + let commitment_post = Commitment::new(account_id, &post_with_updated_nonce); + let encrypted_account = EncryptionScheme::encrypt( + &post_with_updated_nonce, + identifier, + shared_secret, + &commitment_post, + *output_index, + ); + + output.new_commitments.push(commitment_post); + output.ciphertexts.push(encrypted_account); + *output_index = output_index + .checked_add(1) + .unwrap_or_else(|| panic!("Too many private accounts, output index overflow")); +} + +fn compute_update_nullifier_and_set_digest( + membership_proof: &MembershipProof, + pre_account: &Account, + account_id: &AccountId, + nsk: &NullifierSecretKey, +) -> (Nullifier, CommitmentSetDigest) { + let commitment_pre = Commitment::new(account_id, pre_account); + let set_digest = compute_digest_for_path(&commitment_pre, membership_proof); + let nullifier = Nullifier::for_account_update(&commitment_pre, nsk); + (nullifier, set_digest) +}