From e1f24f3411089aaf3ad59d77feee7fdef96b92cf Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Fri, 18 Jul 2025 09:55:23 -0300 Subject: [PATCH] refactor client --- .../examples/happy_path.rs | 16 +++++----- .../examples/mocked_components/client/mod.rs | 25 +++++++++++++++ .../client/transfer_deshielded.rs | 27 ++++++---------- .../client/transfer_private.rs | 25 +++++---------- .../client/transfer_shielded.rs | 18 +++-------- .../mocked_components/sequencer/mod.rs | 2 -- risc0-selective-privacy-poc/src/lib.rs | 32 +++++++++++++++++-- 7 files changed, 84 insertions(+), 61 deletions(-) diff --git a/risc0-selective-privacy-poc/examples/happy_path.rs b/risc0-selective-privacy-poc/examples/happy_path.rs index 2ea921d..13402dc 100644 --- a/risc0-selective-privacy-poc/examples/happy_path.rs +++ b/risc0-selective-privacy-poc/examples/happy_path.rs @@ -17,7 +17,7 @@ fn main() { let mut sequencer = MockedSequencer::new(); let addresses = sequencer.addresses(); println!("addresses: {:?}", addresses); - println!("1️⃣ πŸš€ Initial balances"); + println!("πŸš€ Initial balances"); sequencer.print(); // A public execution of the Transfer Program @@ -26,7 +26,7 @@ fn main() { sequencer .invoke_public_execution::(&[sender_addr, receiver_addr], 10) .unwrap(); - println!("2️⃣ πŸš€ Balances after transfer"); + println!("πŸš€ Balances after transfer"); sequencer.print(); // A shielded execution of the Transfer Program @@ -36,25 +36,25 @@ fn main() { sequencer.print(); // A private execution of the Transfer Program - let private_account_1 = MockedClient::transfer_private( - &private_account_2, + let [_, private_account_1] = MockedClient::transfer_private( + private_account_2, &ACCOUNTS_PRIVATE_KEYS[1], // <-- this is shifted 🫠 &addresses[3], 8, &mut sequencer, ); - println!("3️⃣ πŸš€ Balances after shielded execution"); + println!("πŸš€ Balances after shielded execution"); sequencer.print(); // A deshielded execution of the Transfer Program MockedClient::transfer_deshielded( - &private_account_1, + private_account_1, &ACCOUNTS_PRIVATE_KEYS[0], &addresses[0], 1, &mut sequencer, ); - println!("4️⃣ πŸš€ Balances after deshielded execution"); + println!("πŸš€ Balances after deshielded execution"); sequencer.print(); // A public execution of the Pinata program @@ -62,6 +62,6 @@ fn main() { sequencer .invoke_public_execution::(&[addresses[0], addresses[3]], preimage) .unwrap(); - println!("5️⃣ πŸš€ Balances after public piΓ±ata execution"); + println!("πŸš€ Balances after public piΓ±ata execution"); sequencer.print(); } diff --git a/risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs b/risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs index 04bea16..54d3edb 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs @@ -20,4 +20,29 @@ impl MockedClient { let nonce = [0; 8]; Account::new(address, nonce) } + + fn prove_and_send_to_sequencer( + input_accounts: &[Account], + instruction_data: P::InstructionData, + visibilities: &[InputVisibiility], + commitment_tree_root: [u32; 8], + sequencer: &mut MockedSequencer, + ) -> Vec { + let (receipt, private_outputs) = nssa::invoke_privacy_execution::

( + input_accounts, + instruction_data, + visibilities, + commitment_tree_root, + ) + .unwrap(); + let output: (Vec, Vec, Vec, [u32; 8]) = + receipt.journal.decode().unwrap(); + + // Send to te sequencer + sequencer + .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2) + .unwrap(); + + private_outputs + } } diff --git a/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_deshielded.rs b/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_deshielded.rs index c89d4ad..919b775 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_deshielded.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_deshielded.rs @@ -4,7 +4,7 @@ use core::types::{Address, Commitment, Key, Nullifier}; use nssa::program::TransferProgram; -use super::{MockedSequencer, MockedClient}; +use super::{MockedClient, MockedSequencer}; impl MockedClient { /// A shielded execution of the Transfer program @@ -19,25 +19,16 @@ impl MockedClient { let commitment_tree_root = sequencer.get_commitment_tree_root(); let receiver_addr = to_address; let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr); - let visibilities = vec![InputVisibiility::Public, InputVisibiility::Private(None)]; - let (receipt, nonces) = nssa::invoke_privacy_execution::( - &[sender_account, receiver_account.clone()], + let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)]; + let private_outputs = Self::prove_and_send_to_sequencer::( + &[sender_account, receiver_account], balance_to_move, &visibilities, commitment_tree_root, - ) - .unwrap(); - - // Assemble the private account - receiver_account.nonce = nonces[1]; - receiver_account.balance = balance_to_move; - let output: (Vec, Vec, Vec, [u32; 8]) = - receipt.journal.decode().unwrap(); - - // Send to te sequencer - sequencer - .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2) - .unwrap(); - receiver_account + sequencer, + ); + let [receiver_private_account] = private_outputs.try_into().unwrap(); + receiver_private_account } + } diff --git a/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_private.rs b/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_private.rs index 6bf6a6f..c84d68a 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_private.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_private.rs @@ -9,12 +9,12 @@ use super::{MockedClient, MockedSequencer}; impl MockedClient { /// A private execution of the Transfer program pub fn transfer_private( - from_account: &Account, + from_account: Account, from_account_pk: &Key, to_address: &Address, balance_to_move: u128, sequencer: &mut MockedSequencer, - ) -> Account { + ) -> [Account; 2] { // All of this is executed locally by the sender let commitment_tree_root = sequencer.get_commitment_tree_root(); let receiver_addr = to_address; @@ -25,24 +25,15 @@ impl MockedClient { InputVisibiility::Private(Some((from_account_pk.clone(), sender_commitment_auth_path))), InputVisibiility::Private(None), ]; - let (receipt, nonces) = nssa::invoke_privacy_execution::( - &[from_account.clone(), receiver_account.clone()], + + let private_outputs = Self::prove_and_send_to_sequencer::( + &[from_account, receiver_account], balance_to_move, &visibilities, commitment_tree_root, - ) - .unwrap(); - let output: (Vec, Vec, Vec, [u32; 8]) = - receipt.journal.decode().unwrap(); + sequencer, + ); - // Send to te sequencer - sequencer - .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2) - .unwrap(); - - // Assemble the private account - receiver_account.nonce = nonces[1]; - receiver_account.balance = balance_to_move; - receiver_account + private_outputs.try_into().unwrap() } } diff --git a/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_shielded.rs b/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_shielded.rs index 03b61ce..69f9193 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_shielded.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/client/transfer_shielded.rs @@ -8,7 +8,7 @@ use super::{MockedClient, MockedSequencer}; impl MockedClient { pub fn transfer_deshielded( - from_account: &Account, + from_account: Account, from_account_pk: &Key, to_address: &Address, balance_to_move: u128, @@ -24,20 +24,12 @@ impl MockedClient { InputVisibiility::Private(Some((from_account_pk.clone(), sender_commitment_auth_path))), InputVisibiility::Public, ]; - let (receipt, nonces) = nssa::invoke_privacy_execution::( - &[from_account.clone(), to_account], + let private_outputs = Self::prove_and_send_to_sequencer::( + &[from_account, to_account], balance_to_move, &visibilities, commitment_tree_root, - ) - .unwrap(); - - let output: (Vec, Vec, Vec, [u32; 8]) = - receipt.journal.decode().unwrap(); - - // Send to te sequencer - sequencer - .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2) - .unwrap(); + sequencer, + ); } } diff --git a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs b/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs index 3a0fafd..26a457a 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs @@ -5,9 +5,7 @@ use core::{ }; use std::collections::{BTreeMap, HashSet}; -use nssa; use program_methods::{PINATA_ID, TRANSFER_ID, TRANSFER_MULTIPLE_ID}; -use risc0_zkvm::Receipt; use sparse_merkle_tree::SparseMerkleTree; use super::ACCOUNTS_PRIVATE_KEYS; diff --git a/risc0-selective-privacy-poc/src/lib.rs b/risc0-selective-privacy-poc/src/lib.rs index 1bef79c..abce4f0 100644 --- a/risc0-selective-privacy-poc/src/lib.rs +++ b/risc0-selective-privacy-poc/src/lib.rs @@ -1,7 +1,7 @@ use core::{ account::Account, input::InputVisibiility, - types::{Commitment, Nonce, Nullifier}, + types::{AuthenticationPath, Commitment, Key, Nonce, Nullifier}, }; use program_methods::{OUTER_ELF, OUTER_ID}; use rand::{rngs::OsRng, Rng}; @@ -66,12 +66,32 @@ pub fn execute( Ok(inputs_outputs) } +pub fn extract_private_outputs_from_inner_results( + inputs_outputs: &[Account], + num_inputs: usize, + visibilities: &[InputVisibiility], + nonces: &[Nonce], +) -> Vec { + inputs_outputs + .iter() + .skip(num_inputs) + .zip(visibilities) + .zip(nonces) + .filter(|((_, visibility), _)| matches!(visibility, InputVisibiility::Private(_))) + .map(|((account, _), nonce)| { + let mut this = account.clone(); + this.nonce = *nonce; + this + }) + .collect() +} + pub fn invoke_privacy_execution( inputs: &[Account], instruction_data: P::InstructionData, visibilities: &[InputVisibiility], commitment_tree_root: [u32; 8], -) -> Result<(Receipt, Vec), ()> { +) -> Result<(Receipt, Vec), ()> { // Prove inner program and get post state of the accounts let num_inputs = inputs.len(); let (inner_receipt, inputs_outputs) = execute_and_prove_inner::

(inputs, instruction_data)?; @@ -93,7 +113,13 @@ pub fn invoke_privacy_execution( let prover = default_prover(); let prove_info = prover.prove(env, OUTER_ELF).unwrap(); - Ok((prove_info.receipt, output_nonces)) + let private_outputs = extract_private_outputs_from_inner_results( + &inputs_outputs, + num_inputs, + &visibilities, + &output_nonces, + ); + Ok((prove_info.receipt, private_outputs)) } pub fn verify_privacy_execution(