fix repeated accounts bug

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::<HashSet<_>>()
+        .len()
+        != input_accounts.len()
+    {
+        return false;
+    }
 
     true
 }

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<AccountVisibility>: A vector indicating which accounts are private and which are public.
-/// - Vec<Nonce>: 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<Nonce>: 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<AccountVisibility>,
+    account_visibilities: &[AccountVisibility],
     commitment_tree_root: [u32; 8],
 ) -> Vec<[u32; 8]> {
     let mut nullifiers = Vec::new();