diff --git a/risc0-selective-privacy-poc/core/src/lib.rs b/risc0-selective-privacy-poc/core/src/lib.rs index dbd4b5b..853ab40 100644 --- a/risc0-selective-privacy-poc/core/src/lib.rs +++ b/risc0-selective-privacy-poc/core/src/lib.rs @@ -2,6 +2,8 @@ pub mod account; pub mod types; pub mod visibility; +use std::collections::HashSet; + use crate::{ account::Account, types::{AuthenticationPath, Commitment, Key, Nullifier, ProgramId}, @@ -59,7 +61,7 @@ pub fn bytes_to_words(bytes: &[u8; 32]) -> [u32; 8] { /// - does not change account addresses /// - does not change account nonces /// - does not change the `program_owner` field -/// - only reduces the balance of accounts it owns +/// - only reduces the balance of accounts it owns /// - preserves the total token supply across all accounts /// /// This function does **not** check that the output accounts are the result of correctly @@ -104,6 +106,20 @@ pub fn check_well_behaved_account_transition( if total_balance_pre != total_balance_post { return false; } + // The previous check is only meaningful if the accounts involved in the sum are all different + // otherwise the same balance is counted more than once in the total balance check. Therefore, we need to check + // that the set of input accounts doesn't have repeated pairs of (address, nonce). + // Note: Checking only that addresses are unique would be too restrictive, since different private accounts can have + // the same address (but different nonces). + if input_accounts + .iter() + .map(|account| (account.address, account.nonce)) + .collect::>() + .len() + != input_accounts.len() + { + return false; + } true } diff --git a/risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs b/risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs index 2f4fcd7..bc65957 100644 --- a/risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs +++ b/risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs @@ -10,13 +10,10 @@ use risc0_zkvm::{guest::env, serde::to_vec}; /// It also verifies that the chain's invariants are not violated and the program is well-behaved. /// /// Inputs: -/// - ProgramOuptut: The output of the inner program. This is includes the accounts pre and -/// post-states of the execution of the inner program. +/// - ProgramOuptut: The output of the inner program. This is includes the accounts pre and post-states of the execution of the inner program. /// - Vec: A vector indicating which accounts are private and which are public. -/// - Vec: The vector of nonces to be used for the output accounts. This is assumed to be -/// sampled at random by the host program. -/// - [u32; 8]: The root of the commitment tree. Commitments of input private accounts will be -/// checked against this to prove that they belong to the tree. +/// - Vec: The vector of nonces to be used for the output accounts. This is assumed to be sampled at random by the host program. +/// - [u32; 8]: The root of the commitment tree. Commitments of input private accounts will be checked against this to prove that they belong to the tree. /// - ProgamId: The ID of the inner program. /// /// Public outputs: @@ -75,13 +72,18 @@ fn assemble_privacy_execution_output( commitment_tree_root: [u32; 8], nullifiers: Vec<[u32; 8]>, ) -> PrivacyExecutionOutput { - // Insert new nonces in outputs (including public ones) + // Insert new nonces in private outputs let accounts_pre = inner_program_output.accounts_pre; let mut accounts_post = inner_program_output.accounts_post; accounts_post .iter_mut() .zip(output_nonces) - .for_each(|(account, new_nonce)| account.nonce = new_nonce); + .zip(visibilities.iter()) + .for_each(|((account, new_nonce), visibility)| { + if matches!(visibility, AccountVisibility::Private(_)) { + account.nonce = new_nonce; + } + }); // Compute commitments for every private output let mut private_outputs = Vec::new(); @@ -110,20 +112,19 @@ fn assemble_privacy_execution_output( } } - let output = PrivacyExecutionOutput { + PrivacyExecutionOutput { public_accounts_pre, public_accounts_post, private_output_commitments, nullifiers, commitment_tree_root, - }; - output + } } /// Compute nullifiers of private accounts pre states and check that their commitments belong to the commitments tree fn verify_and_nullify_private_inputs( inner_program_output: &ProgramOutput, - account_visibilities: &Vec, + account_visibilities: &[AccountVisibility], commitment_tree_root: [u32; 8], ) -> Vec<[u32; 8]> { let mut nullifiers = Vec::new();