From 04def6e82b9fb1a91cae0f06f92f36980ee1aec0 Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Tue, 15 Jul 2025 08:07:27 -0300 Subject: [PATCH] add public pre and post states to the output of the outer program --- risc0-selective-privacy-poc/Cargo.toml | 10 +- .../core/src/account.rs | 5 - .../outer_methods/guest/src/bin/outer.rs | 91 +++++++++++++------ .../src/private_execution.rs | 39 ++++---- .../src/public_execution.rs | 18 +++- .../guest/src/bin/transfer.rs | 2 +- 6 files changed, 103 insertions(+), 62 deletions(-) diff --git a/risc0-selective-privacy-poc/Cargo.toml b/risc0-selective-privacy-poc/Cargo.toml index e817b5d..a32e53f 100644 --- a/risc0-selective-privacy-poc/Cargo.toml +++ b/risc0-selective-privacy-poc/Cargo.toml @@ -5,13 +5,15 @@ edition = "2021" [dependencies] risc0-zkvm = "2.2" -toy-example-core = {path = "core"} -transfer-methods = {path = "transfer_methods"} -outer-methods = {path = "outer_methods"} +toy-example-core = { path = "core" } +transfer-methods = { path = "transfer_methods" } +outer-methods = { path = "outer_methods" } serde = "1.0" tracing-subscriber = { version = "0.3", features = ["env-filter"] } +rand = "0.8" [features] cuda = ["risc0-zkvm/cuda"] default = [] -prove = ["risc0-zkvm/prove"] \ No newline at end of file +prove = ["risc0-zkvm/prove"] + diff --git a/risc0-selective-privacy-poc/core/src/account.rs b/risc0-selective-privacy-poc/core/src/account.rs index 9cfe075..c4da8dd 100644 --- a/risc0-selective-privacy-poc/core/src/account.rs +++ b/risc0-selective-privacy-poc/core/src/account.rs @@ -52,8 +52,3 @@ pub fn compute_nullifier(commitment: &[u32; 8], private_key: &[u32; 8]) -> [u32; bytes_to_hash[8..].copy_from_slice(private_key); hash(&bytes_to_hash) } - -/// Dummy implementation -pub fn new_random_nonce() -> [u32; 8] { - [0xcc, 0xaa, 0xff, 0xee, 0xcc, 0xaa, 0xff, 0xff] -} diff --git a/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs b/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs index dc5a1b4..cc93293 100644 --- a/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs +++ b/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs @@ -1,8 +1,4 @@ -use risc0_zkvm::{ - guest::env, - serde::to_vec, - sha::{Impl, Sha256}, -}; +use risc0_zkvm::{guest::env, serde::to_vec}; use toy_example_core::{ account::{compute_nullifier, hash, is_in_commitment_tree, Account}, input::InputVisibiility, @@ -26,25 +22,32 @@ use toy_example_core::{ fn main() { let num_inputs: u32 = env::read(); // Read inputs and outputs - let mut inputs_outputs = Vec::new(); - for _ in 0..(2 * num_inputs) { - let account: Account = env::read(); - inputs_outputs.push(account); - } + let mut inputs_outputs: Vec = env::read(); + assert_eq!(inputs_outputs.len() as u32, num_inputs * 2); // Read visibilities - let mut input_visibilities = Vec::new(); - for _ in 0..num_inputs { - let input_visibility: InputVisibiility = env::read(); - input_visibilities.push(input_visibility); - } + let input_visibilities: Vec = env::read(); + assert_eq!(input_visibilities.len() as u32, num_inputs); + + // Read nonces for outputs + let output_nonces: Vec<[u32; 8]> = env::read(); + assert_eq!(output_nonces.len() as u32, num_inputs); let commitment_tree_root: [u32; 8] = env::read(); let program_id: [u32; 8] = env::read(); - let inputs = inputs_outputs.iter().take(num_inputs as usize); + // Verify pre states and post states of accounts are consistent + // with the execution of the `program_id`` program + env::verify(program_id, &to_vec(&inputs_outputs).unwrap()).unwrap(); + + // Split inputs_outputs into two separate vectors + let (inputs, mut outputs) = { + let outputs = inputs_outputs.split_off(num_inputs as usize); + (inputs_outputs, outputs) + }; + let mut nullifiers = Vec::new(); - for (visibility, input_account) in input_visibilities.iter().zip(inputs) { + for (visibility, input_account) in input_visibilities.iter().zip(inputs.iter()) { match visibility { InputVisibiility::Private(Some(private_key)) => { // Prove ownership of input accounts by proving @@ -55,7 +58,7 @@ fn main() { let commitment = input_account.commitment(); assert!(is_in_commitment_tree(commitment, commitment_tree_root)); // Compute nullifier to nullify this private input account. - let nullifier = compute_nullifier(&commitment, &private_key); + let nullifier = compute_nullifier(&commitment, private_key); nullifiers.push(nullifier); } InputVisibiility::Private(None) => { @@ -67,22 +70,52 @@ fn main() { InputVisibiility::Public => continue, } } - let outputs = inputs_outputs.iter().skip(num_inputs as usize); - let output_commitments: Vec<_> = outputs.map(|account| account.commitment()).collect(); // Assert `program_id` program didn't modify address fields - for (account_pre, account_post) in inputs_outputs - .iter() - .take(num_inputs as usize) - .zip(inputs_outputs.iter().skip(num_inputs as usize)) - { + for (account_pre, account_post) in inputs.iter().zip(outputs.iter()) { assert_eq!(account_pre.address, account_post.address); } - // Verify pre states and post states of accounts are consistent - // with the execution of the `program_id`` program - env::verify(program_id, &to_vec(&inputs_outputs).unwrap()).unwrap(); + // Insert new nonces in outputs (including public ones (?!)) + outputs + .iter_mut() + .zip(output_nonces) + .for_each(|(account, new_nonce)| account.nonce = new_nonce); + + // Compute private outputs commitments + let mut private_outputs = Vec::new(); + for (output, visibility) in outputs.iter().zip(input_visibilities.iter()) { + match visibility { + InputVisibiility::Public => continue, + InputVisibiility::Private(_) => private_outputs.push(output), + } + } + + // Get the list of public inputs pre states and their post states + let mut public_inputs_outputs = Vec::new(); + for (account, visibility) in inputs + .iter() + .chain(outputs.iter()) + .zip(input_visibilities.iter().chain(input_visibilities.iter())) + { + match visibility { + InputVisibiility::Public => { + public_inputs_outputs.push(account); + } + InputVisibiility::Private(_) => continue, + } + } + + // Compute commitments for every private output + let private_output_commitments: Vec<_> = private_outputs + .iter() + .map(|account| account.commitment()) + .collect(); // Output nullifier of consumed input accounts and commitments of new output private accounts - env::commit(&(nullifiers, output_commitments)); + env::commit(&( + public_inputs_outputs, + nullifiers, + private_output_commitments, + )); } diff --git a/risc0-selective-privacy-poc/src/private_execution.rs b/risc0-selective-privacy-poc/src/private_execution.rs index a4b11c6..65e1a97 100644 --- a/risc0-selective-privacy-poc/src/private_execution.rs +++ b/risc0-selective-privacy-poc/src/private_execution.rs @@ -1,13 +1,16 @@ use outer_methods::{OUTER_ELF, OUTER_ID}; +use rand::{rngs::OsRng, Rng}; use risc0_zkvm::{default_prover, ExecutorEnv, Receipt}; -use toy_example_core::{ - account::{new_random_nonce, Account}, - input::InputVisibiility, -}; +use toy_example_core::{account::Account, input::InputVisibiility}; use transfer_methods::{TRANSFER_ELF, TRANSFER_ID}; const COMMITMENT_TREE_ROOT: [u32; 8] = [0xdd, 0xee, 0xaa, 0xdd, 0xbb, 0xee, 0xee, 0xff]; +pub fn new_random_nonce() -> [u32; 8] { + let mut rng = OsRng; + std::array::from_fn(|_| rng.gen()) +} + fn mint_fresh_account(address: [u32; 8]) -> Account { let nonce = new_random_nonce(); Account::new(address, nonce) @@ -35,19 +38,16 @@ fn run_private_execution_of_transfer_program() { let receiver = mint_fresh_account(receiver_address); // Prove inner program and get post state of the accounts - let (inner_receipt, outputs) = prove_inner(&sender, &receiver, balance_to_move); + let (inner_receipt, inputs_outputs) = prove_inner(&sender, &receiver, balance_to_move); let visibilities = vec![ InputVisibiility::Private(Some(sender_private_key)), InputVisibiility::Private(None), ]; - let inputs_outputs = { - let mut vec = vec![sender, receiver]; - vec.extend_from_slice(&outputs); - vec - }; let num_inputs: u32 = inputs_outputs.len() as u32 / 2; + let output_nonces: Vec<_> = (0..num_inputs).map(|_| new_random_nonce()).collect(); + println!("output nonces {output_nonces:?}"); // Prove outer program. // This computes the nullifier for the input account @@ -57,6 +57,7 @@ fn run_private_execution_of_transfer_program() { env_builder.write(&num_inputs).unwrap(); env_builder.write(&inputs_outputs).unwrap(); env_builder.write(&visibilities).unwrap(); + env_builder.write(&output_nonces).unwrap(); env_builder.write(&COMMITMENT_TREE_ROOT).unwrap(); env_builder.write(&TRANSFER_ID).unwrap(); let env = env_builder.build().unwrap(); @@ -69,10 +70,10 @@ fn run_private_execution_of_transfer_program() { // Sanity check receipt.verify(OUTER_ID).unwrap(); - let output: [[u32; 8]; 3] = receipt.journal.decode().unwrap(); - println!("nullifier: {:?}", output[0]); - println!("commitment_1: {:?}", output[1]); - println!("commitment_2: {:?}", output[2]); + let output: (Vec, Vec<[u32; 8]>, Vec<[u32; 8]>) = receipt.journal.decode().unwrap(); + println!("public_outputs: {:?}", output.0); + println!("nullifiers: {:?}", output.1); + println!("commitments: {:?}", output.2); } fn prove_inner( @@ -91,22 +92,22 @@ fn prove_inner( let receipt = prove_info.receipt; - let output: [Account; 4] = receipt.journal.decode().unwrap(); - let [_, _, sender_post, receiver_post] = output; + let inputs_outputs: Vec = receipt.journal.decode().unwrap(); + assert_eq!(inputs_outputs.len(), 4); println!( "sender_before: {:?}, sender_after: {:?}", - sender, sender_post + inputs_outputs[0], inputs_outputs[2] ); println!( "receiver_before: {:?}, receiver_after: {:?}", - receiver, receiver_post + inputs_outputs[1], inputs_outputs[3] ); // Sanity check receipt.verify(TRANSFER_ID).unwrap(); - (receipt, vec![sender_post, receiver_post]) + (receipt, inputs_outputs) } #[cfg(test)] diff --git a/risc0-selective-privacy-poc/src/public_execution.rs b/risc0-selective-privacy-poc/src/public_execution.rs index bc0c536..6a2889c 100644 --- a/risc0-selective-privacy-poc/src/public_execution.rs +++ b/risc0-selective-privacy-poc/src/public_execution.rs @@ -29,11 +29,21 @@ pub fn run_public_execution_of_transfer_program() { let env = env_builder.build().unwrap(); let executor = default_executor(); - let result: [Account; 4] = executor.execute(env, TRANSFER_ELF).unwrap().journal.decode().unwrap(); - let [_, _, sender_post, receiver_post] = result; + let inputs_outputs: Vec = executor + .execute(env, TRANSFER_ELF) + .unwrap() + .journal + .decode() + .unwrap(); - println!("sender_before: {:?}, sender_after: {:?}", sender, sender_post); - println!("receiver_before: {:?}, receiver_after: {:?}", receiver, receiver_post); + println!( + "sender_before: {:?}, sender_after: {:?}", + inputs_outputs[0], inputs_outputs[2] + ); + println!( + "receiver_before: {:?}, receiver_after: {:?}", + inputs_outputs[1], inputs_outputs[3], + ); } #[cfg(test)] diff --git a/risc0-selective-privacy-poc/transfer_methods/guest/src/bin/transfer.rs b/risc0-selective-privacy-poc/transfer_methods/guest/src/bin/transfer.rs index 2f33c0f..1fe7de8 100644 --- a/risc0-selective-privacy-poc/transfer_methods/guest/src/bin/transfer.rs +++ b/risc0-selective-privacy-poc/transfer_methods/guest/src/bin/transfer.rs @@ -18,5 +18,5 @@ fn main() { sender_post.balance -= balance_to_move; receiver_post.balance += balance_to_move; - env::commit(&(sender, receiver, sender_post, receiver_post)); + env::commit(&vec![sender, receiver, sender_post, receiver_post]); }