From 6db06722c165c09fb9e4a91f5133ebd3d67b5a3b Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Fri, 18 Jul 2025 15:56:41 -0300 Subject: [PATCH] refactor --- .../core/src/account.rs | 28 +++--- risc0-selective-privacy-poc/core/src/lib.rs | 2 +- .../core/src/{input.rs => visibility.rs} | 2 +- .../examples/happy_path.rs | 85 ++++++++++++----- .../examples/mocked_components/client/mod.rs | 16 ++-- .../client/transfer_deshielded.rs | 24 +++-- .../client/transfer_private.rs | 9 +- .../client/transfer_shielded.rs | 25 +++-- .../mocked_components/sequencer/mod.rs | 91 ++++++++++++------- .../examples/private_execution.rs | 4 +- .../examples/public_execution.rs | 19 +--- .../program_methods/guest/src/bin/outer.rs | 5 +- risc0-selective-privacy-poc/src/lib.rs | 2 +- 13 files changed, 185 insertions(+), 127 deletions(-) rename risc0-selective-privacy-poc/core/src/{input.rs => visibility.rs} (100%) diff --git a/risc0-selective-privacy-poc/core/src/account.rs b/risc0-selective-privacy-poc/core/src/account.rs index cdc56f2..fd0549a 100644 --- a/risc0-selective-privacy-poc/core/src/account.rs +++ b/risc0-selective-privacy-poc/core/src/account.rs @@ -5,6 +5,7 @@ use crate::{ use risc0_zkvm::{serde::to_vec, sha::Impl}; use serde::{Deserialize, Serialize}; +/// Account to be used both in public and private contexts #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] pub struct Account { pub address: Address, @@ -13,25 +14,26 @@ pub struct Account { } impl Account { - /// Creates a new account with address = hash(private_key) and balance = 0 - pub fn new_from_private_key(private_key: Key, nonce: Nonce) -> Self { - let address = Self::address_for_key(&private_key); - Self::new(address, nonce) + pub fn new(address: Address, balance: u128) -> Self { + Self { + address, + balance, + nonce: [0; 8], + } } + /// Creates a new account with address = hash(private_key) and balance = 0 + pub fn new_from_private_key(private_key: Key) -> Self { + let address = Self::address_for_key(&private_key); + Self::new(address, 0) + } + + /// Computes the address corresponding to the given private key pub fn address_for_key(private_key: &Key) -> Address { hash(private_key) } - pub fn new(address: Address, nonce: Nonce) -> Self { - Self { - address, - balance: 0, - nonce, - } - } - - /// Returns Hash(Account)[0] (only first word for this POC) + /// Returns (first 8 bytes of) SHA256(Account) pub fn commitment(&self) -> Commitment { hash(&to_vec(&self).unwrap())[0] } diff --git a/risc0-selective-privacy-poc/core/src/lib.rs b/risc0-selective-privacy-poc/core/src/lib.rs index 87d255f..57f40bc 100644 --- a/risc0-selective-privacy-poc/core/src/lib.rs +++ b/risc0-selective-privacy-poc/core/src/lib.rs @@ -1,6 +1,6 @@ pub mod account; -pub mod input; pub mod types; +pub mod visibility; use crate::types::{AuthenticationPath, Commitment, Key, Nullifier}; use risc0_zkvm::sha::{Impl, Sha256}; diff --git a/risc0-selective-privacy-poc/core/src/input.rs b/risc0-selective-privacy-poc/core/src/visibility.rs similarity index 100% rename from risc0-selective-privacy-poc/core/src/input.rs rename to risc0-selective-privacy-poc/core/src/visibility.rs index 944164a..dc04fd0 100644 --- a/risc0-selective-privacy-poc/core/src/input.rs +++ b/risc0-selective-privacy-poc/core/src/visibility.rs @@ -1,5 +1,5 @@ -use serde::{Deserialize, Serialize}; use crate::types::{AuthenticationPath, Key}; +use serde::{Deserialize, Serialize}; #[derive(Serialize, Deserialize)] pub enum InputVisibiility { diff --git a/risc0-selective-privacy-poc/examples/happy_path.rs b/risc0-selective-privacy-poc/examples/happy_path.rs index 580b2e1..d25abd3 100644 --- a/risc0-selective-privacy-poc/examples/happy_path.rs +++ b/risc0-selective-privacy-poc/examples/happy_path.rs @@ -1,14 +1,14 @@ use core::{ account::Account, bytes_to_words, hash, - input::InputVisibiility, types::{Address, Commitment, Key, Nullifier}, + visibility::InputVisibiility, }; use nssa::program::{PinataProgram, TransferProgram}; use risc0_zkvm::Receipt; -use crate::mocked_components::sequencer::MockedSequencer; +use crate::mocked_components::sequencer::{print_accounts, MockedSequencer}; use crate::mocked_components::{client::MockedClient, USER_CLIENTS}; mod mocked_components; @@ -16,44 +16,81 @@ mod mocked_components; fn main() { let mut sequencer = MockedSequencer::new(); let addresses: [Address; 3] = USER_CLIENTS.map(|client| client.user_address()); - - println!("addresses: {:?}", addresses); - println!("๐Ÿš€ Initial balances"); - sequencer.print(); + println!("๐Ÿ“ Initial balances"); + print_accounts(&sequencer, &[]); // A public execution of the Transfer Program - USER_CLIENTS[0] - .transfer_public(&addresses[1], 10, &mut sequencer) + USER_CLIENTS[1] + .transfer_public(&[0xcafe; 8], 51, &mut sequencer) .unwrap(); - println!("๐Ÿš€ Balances after transfer"); - sequencer.print(); + println!("๐Ÿ“ Balances after transfer"); + print_accounts(&sequencer, &[]); // A shielded execution of the Transfer Program - let private_account_1 = USER_CLIENTS[0] + let private_account_user_1 = USER_CLIENTS[0] .transfer_shielded(&addresses[1], 15, &mut sequencer) .unwrap(); - println!("Balances after shielded execution"); - sequencer.print(); + println!("๐Ÿ“ Balances after shielded execution"); + print_accounts(&sequencer, &[&private_account_user_1]); // A private execution of the Transfer Program - let [_, private_account_2] = USER_CLIENTS[1] - .transfer_private(private_account_1, &addresses[2], 8, &mut sequencer) + let [private_account_user_1, private_account_user_2] = USER_CLIENTS[1] + .transfer_private(private_account_user_1, &addresses[2], 8, &mut sequencer) .unwrap(); - println!("๐Ÿš€ Balances after shielded execution"); - sequencer.print(); + println!("๐Ÿ“ Balances after shielded execution"); + print_accounts( + &sequencer, + &[&private_account_user_1, &private_account_user_2], + ); // A deshielded execution of the Transfer Program - USER_CLIENTS[2] - .transfer_deshielded(private_account_2, &addresses[0], 1, &mut sequencer) + let private_acount_user_2 = USER_CLIENTS[2] + .transfer_deshielded(private_account_user_2, &addresses[0], 1, &mut sequencer) .unwrap(); - println!("๐Ÿš€ Balances after deshielded execution"); - sequencer.print(); + println!("๐Ÿ“ Balances after deshielded execution"); + print_accounts( + &sequencer, + &[&private_account_user_1, &private_acount_user_2], + ); - // A public execution of the Pinata program + // A public execution of the Piรฑata program let preimage = bytes_to_words(b"NSSA Selective privacy is great!").to_vec(); sequencer .process_public_execution::(&[[0xcafe; 8], addresses[2]], preimage) .unwrap(); - println!("๐Ÿš€ Balances after public piรฑata execution"); - sequencer.print(); + println!("๐Ÿ“ Balances after public piรฑata execution"); + print_accounts( + &sequencer, + &[&private_account_user_1, &private_acount_user_2], + ); + + // A deshielded execution of the Piรฑata program + let private_account_user_0 = { + // All of this is executed locally by the sender + let receiver_addr = USER_CLIENTS[1].user_address(); + let pinata_account = sequencer.get_account(&[0xcafe; 8]).unwrap(); + let mut receiver_account = MockedClient::fresh_account_for_mint(receiver_addr); + let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)]; + let preimage = bytes_to_words(b"NSSA Selective privacy is great!").to_vec(); + + let private_outputs = MockedClient::prove_and_send_to_sequencer::( + &[pinata_account, receiver_account], + preimage, + &visibilities, + sequencer.get_commitment_tree_root(), + &mut sequencer, + ) + .unwrap(); + let [private_account_user_0] = private_outputs.try_into().unwrap(); + private_account_user_0 + }; + println!("๐Ÿ“ Balances after private piรฑata execution"); + print_accounts( + &sequencer, + &[ + &private_account_user_1, + &private_acount_user_2, + &private_account_user_0, + ], + ); } 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 1fd754f..a6a9cdc 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs @@ -1,7 +1,7 @@ use core::{ account::Account, - input::InputVisibiility, types::{Address, Commitment, Key, Nullifier}, + visibility::InputVisibiility, }; use nssa::program::TransferProgram; @@ -15,19 +15,22 @@ pub mod transfer_shielded; pub struct MockedClient { user_private_key: Key, + private_accounts: Vec, } impl MockedClient { pub const fn new(user_private_key: Key) -> Self { - Self { user_private_key } + Self { + user_private_key, + private_accounts: Vec::new(), + } } pub fn user_address(&self) -> Address { - let address = Account::address_for_key(&self.user_private_key); - address + Account::address_for_key(&self.user_private_key) } - fn prove_and_send_to_sequencer( + pub fn prove_and_send_to_sequencer( input_accounts: &[Account], instruction_data: P::InstructionData, visibilities: &[InputVisibiility], @@ -48,7 +51,6 @@ impl MockedClient { } pub fn fresh_account_for_mint(address: Address) -> Account { - let nonce = [0; 8]; - Account::new(address, nonce) + Account::new(address, 0) } } 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 df62ce4..b09c173 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 @@ -1,5 +1,5 @@ use core::account::Account; -use core::input::InputVisibiility; +use core::visibility::InputVisibiility; use core::types::{Address, Commitment, Key, Nullifier}; use nssa::program::TransferProgram; @@ -7,28 +7,32 @@ use nssa::program::TransferProgram; use super::{MockedClient, MockedSequencer}; impl MockedClient { - /// A shielded execution of the Transfer program - pub fn transfer_shielded( + pub fn transfer_deshielded( &self, + from_account: Account, to_address: &Address, balance_to_move: u128, sequencer: &mut MockedSequencer, ) -> Result { // All of this is executed locally by the sender - let sender_account = sequencer.get_account(&self.user_address()).ok_or(())?; 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 = [InputVisibiility::Public, InputVisibiility::Private(None)]; - + // let from_account = sequencer.get_account(&self.user_address()).ok_or(())?; + let sender_commitment_auth_path = + sequencer.get_authentication_path_for(&from_account.commitment()); + let to_account = sequencer.get_account(&to_address).unwrap(); + let visibilities = vec![ + InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))), + InputVisibiility::Public, + ]; let private_outputs = Self::prove_and_send_to_sequencer::( - &[sender_account, receiver_account], + &[from_account, to_account], balance_to_move, &visibilities, commitment_tree_root, sequencer, )?; - let [receiver_private_account] = private_outputs.try_into().unwrap(); - Ok(receiver_private_account) + let [sender_private_account] = private_outputs.try_into().unwrap(); + Ok(sender_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 92fe405..7e8a6f8 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 @@ -1,5 +1,5 @@ use core::account::Account; -use core::input::InputVisibiility; +use core::visibility::InputVisibiility; use core::types::{Address, Commitment, Key, Nullifier}; use nssa::program::TransferProgram; @@ -10,7 +10,7 @@ impl MockedClient { /// A private execution of the Transfer program pub fn transfer_private( &self, - from_account: Account, + owned_private_account: Account, to_address: &Address, balance_to_move: u128, sequencer: &mut MockedSequencer, @@ -18,9 +18,8 @@ impl MockedClient { // All of this is executed locally by the sender let commitment_tree_root = sequencer.get_commitment_tree_root(); let receiver_addr = to_address; - // let from_account = sequencer.get_account(&self.user_address()).ok_or(())?; let sender_commitment_auth_path = - sequencer.get_authentication_path_for(&from_account.commitment()); + sequencer.get_authentication_path_for(&owned_private_account.commitment()); let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr); let visibilities = vec![ InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))), @@ -28,7 +27,7 @@ impl MockedClient { ]; let private_outputs = Self::prove_and_send_to_sequencer::( - &[from_account, receiver_account], + &[owned_private_account, receiver_account], balance_to_move, &visibilities, commitment_tree_root, 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 816ff26..1aecf3b 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 @@ -1,37 +1,34 @@ use core::account::Account; -use core::input::InputVisibiility; use core::types::{Address, Commitment, Key, Nullifier}; +use core::visibility::InputVisibiility; use nssa::program::TransferProgram; use super::{MockedClient, MockedSequencer}; impl MockedClient { - pub fn transfer_deshielded( + /// A shielded execution of the Transfer program + pub fn transfer_shielded( &self, - from_account: Account, to_address: &Address, balance_to_move: u128, sequencer: &mut MockedSequencer, - ) -> Result<(), ()> { + ) -> Result { // All of this is executed locally by the sender + let sender_account = sequencer.get_account(&self.user_address()).ok_or(())?; let commitment_tree_root = sequencer.get_commitment_tree_root(); let receiver_addr = to_address; - // let from_account = sequencer.get_account(&self.user_address()).ok_or(())?; - let sender_commitment_auth_path = - sequencer.get_authentication_path_for(&from_account.commitment()); - let to_account = sequencer.get_account(&to_address).unwrap(); - let visibilities = vec![ - InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))), - InputVisibiility::Public, - ]; + let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr); + let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)]; + let private_outputs = Self::prove_and_send_to_sequencer::( - &[from_account, to_account], + &[sender_account, receiver_account], balance_to_move, &visibilities, commitment_tree_root, sequencer, )?; - Ok(()) + let [receiver_private_account] = private_outputs.try_into().unwrap(); + Ok(receiver_private_account) } } 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 358d2ed..0880c75 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs @@ -20,28 +20,20 @@ pub struct MockedSequencer { deployed_program_ids: HashSet, } -const ACCOUNTS_INITIAL_BALANCES: [u128; 3] = [100, 1337, 37]; const DEPLOYED_PROGRAM_IDS: [ProgramId; 3] = [TRANSFER_ID, TRANSFER_MULTIPLE_ID, PINATA_ID]; +const INITIAL_BALANCE: u128 = 150; +const PINATA_ADDRESS: Address = [0xcafe; 8]; impl MockedSequencer { pub fn new() -> Self { let mut accounts: BTreeMap = USER_CLIENTS .iter() .map(|client| client.user_address()) - .zip(ACCOUNTS_INITIAL_BALANCES) - .map(|(address, initial_balance)| { - let mut this = Account::new(address, [0; 8]); - this.balance = initial_balance; - this - }) + .map(|address| Account::new(address, INITIAL_BALANCE)) .map(|account| (account.address, account)) .collect(); - let pinata_account = { - let mut this = Account::new([0xcafe; 8], [0; 8]); - this.balance = 100; - this - }; + let pinata_account = Account::new(PINATA_ADDRESS, INITIAL_BALANCE); accounts.insert(pinata_account.address, pinata_account); let commitment_tree = SparseMerkleTree::new_empty(); @@ -75,23 +67,60 @@ impl MockedSequencer { pub fn addresses(&self) -> Vec
{ self.accounts.keys().cloned().collect() } - - pub fn print(&self) { - println!("{:<20} | {:>10}", "Address (first u32)", "Balance"); - println!("{:-<20}-+-{:-<10}", "", ""); - - for account in self.accounts.values() { - println!("{:<20x} | {:>10}", account.address[0], account.balance); - } - println!("{:-<20}-+-{:-<10}", "", ""); - println!("Commitments: {:?}", self.commitment_tree.values()); - let formatted: Vec = self - .nullifier_set - .iter() - .map(|arr| format!("0x{:x}", arr[0])) - .collect(); - println!("Nullifiers: [{}]", formatted.join(", ")); - println!(""); - println!(""); - } +} + +pub fn print_accounts(sequencer: &MockedSequencer, private_accounts: &[&Account]) { + println!("\n====================== ACCOUNT SNAPSHOT ======================\n"); + + println!(">> Public Accounts:"); + println!("{:<20} | {:>10} |", "Address (first u32)", "Balance"); + println!("{:-<20}-+-{:-<10}", "", ""); + + for account in sequencer.accounts.values() { + println!("0x{:<20x} | {:>10} |", account.address[0], account.balance); + } + + println!("{:-<20}-+-{:-<10}\n", "", ""); + + println!(">> Commitments:"); + println!("{:-<20}", ""); + + for commitment in sequencer.commitment_tree.values().iter() { + println!("{:<20x}", commitment); + } + + println!("{:-<20}\n", ""); + + let formatted: Vec = sequencer + .nullifier_set + .iter() + .map(|nullifier| format!("0x{:x}", nullifier[0])) + .collect(); + + println!(">> Nullifiers (first u32):"); + println!("{:-<20}", ""); + + for entry in formatted { + println!("{:<20}", entry); + } + + println!("{:-<20}\n", ""); + + println!(">> Private Accounts:"); + println!( + "{:<20} | {:>10} | {:>10}", + "Address (first u32)", "Nonce", "Balance" + ); + println!("{:-<20}-+-{:-<10}-+-{:-<10}", "", "", ""); + + for account in private_accounts.iter() { + println!( + "{:<20x} | {:>10x}| {:>10} | ", + account.address[0], account.nonce[0], account.balance, + ); + } + + println!("{:-<20}-+-{:-<10}-+-{:-<10}", "", "", ""); + + println!("\n=============================================================\n"); } diff --git a/risc0-selective-privacy-poc/examples/private_execution.rs b/risc0-selective-privacy-poc/examples/private_execution.rs index 6a6edf2..b32fc9c 100644 --- a/risc0-selective-privacy-poc/examples/private_execution.rs +++ b/risc0-selective-privacy-poc/examples/private_execution.rs @@ -1,7 +1,7 @@ use core::{ account::Account, bytes_to_words, - input::InputVisibiility, + visibility::InputVisibiility, types::{Address, AuthenticationPath, Commitment, Nullifier}, }; use nssa::program::TransferMultipleProgram; @@ -22,7 +22,7 @@ fn main() { let sender_private_key = [1, 2, 3, 4, 4, 3, 2, 1]; let sender = { // Creating it now but it's supposed to be already created by other previous transactions. - let mut account = Account::new_from_private_key(sender_private_key, [1; 8]); + let mut account = Account::new_from_private_key(sender_private_key); account.balance = 150; account }; diff --git a/risc0-selective-privacy-poc/examples/public_execution.rs b/risc0-selective-privacy-poc/examples/public_execution.rs index e2fada7..0b22a3a 100644 --- a/risc0-selective-privacy-poc/examples/public_execution.rs +++ b/risc0-selective-privacy-poc/examples/public_execution.rs @@ -10,24 +10,13 @@ use nssa::program::TransferMultipleProgram; /// the initiating transaction includes the sender's signature. pub fn main() { // Account fetched from the chain state with 150 in its balance. - let sender = { - let mut account = Account::new([5; 8], [98; 8]); - account.balance = 150; - account - }; + let sender = Account::new([5; 8], 150); // Account fetched from the chain state with 900 in its balance. - let receiver_1 = { - let mut account = Account::new([6; 8], [99; 8]); - account.balance = 900; - account - }; + let receiver_1 = Account::new([6; 8], 900); - let receiver_2 = { - let mut account = Account::new([6; 8], [99; 8]); - account.balance = 500; - account - }; + // Account fetched from the chain state with 500 in its balance. + let receiver_2 = Account::new([6; 8], 500); let balance_to_move = vec![10, 20]; 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 8fadb3e..72f8831 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 @@ -1,9 +1,8 @@ use core::{ account::Account, - compute_nullifier, hash, - input::InputVisibiility, - is_in_tree, + compute_nullifier, hash, is_in_tree, types::{Nonce, ProgramId}, + visibility::InputVisibiility, }; use risc0_zkvm::{guest::env, serde::to_vec}; diff --git a/risc0-selective-privacy-poc/src/lib.rs b/risc0-selective-privacy-poc/src/lib.rs index abce4f0..ab0cabd 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::{AuthenticationPath, Commitment, Key, Nonce, Nullifier}, + visibility::InputVisibiility, }; use program_methods::{OUTER_ELF, OUTER_ID}; use rand::{rngs::OsRng, Rng};