diff --git a/risc0-selective-privacy-poc/core/src/lib.rs b/risc0-selective-privacy-poc/core/src/lib.rs index 7650de1..0b8252d 100644 --- a/risc0-selective-privacy-poc/core/src/lib.rs +++ b/risc0-selective-privacy-poc/core/src/lib.rs @@ -2,7 +2,7 @@ pub mod account; pub mod types; pub mod visibility; -use crate::types::{AuthenticationPath, Commitment, Key, Nullifier}; +use crate::{account::Account, types::{AuthenticationPath, Commitment, Key, Nullifier}}; use risc0_zkvm::sha::{Impl, Sha256}; pub fn hash(bytes: &[u32]) -> [u32; 8] { @@ -49,3 +49,33 @@ pub fn bytes_to_words(bytes: &[u8; 32]) -> [u32; 8] { } words } + +/// Verifies that a program public execution didn't break the chain's rules. +/// `input_accounts` are the accounts provided as inputs to the program. +/// `output_accounts` are the accounts post states after execution of the program +pub fn inputs_outputs_preserve_invariants(input_accounts: &[Account], output_accounts: &[Account]) -> bool { + // Fail if the number of input and output accounts differ + if input_accounts.len() != output_accounts.len() { + return false; + } + + for (account_pre, account_post) in input_accounts.iter().zip(output_accounts) { + // Fail if the program modified the addresses of the input accounts + if account_pre.address != account_post.address { + return false; + } + // Fail if the program modified the nonces of the input accounts + if account_pre.nonce != account_post.nonce { + return false; + } + } + + // Fail if the execution didn't preserve the total supply. + let total_balance_pre: u128 = input_accounts.iter().map(|account| account.balance).sum(); + let total_balance_post: u128 = output_accounts.iter().map(|account| account.balance).sum(); + if total_balance_pre != total_balance_post { + return false; + } + + true +} diff --git a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs b/risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs index 808e29a..521c9d1 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs @@ -1,4 +1,4 @@ -use core::{account::Account, types::Address}; +use core::{account::Account, inputs_outputs_preserve_invariants, types::Address}; use crate::mocked_components::sequencer::error::Error; @@ -21,8 +21,18 @@ impl MockedSequencer { let program_output = nssa::execute_onchain::

(&input_accounts, instruction_data).map_err(|_| Error::BadInput)?; - // Perform consistency checks - if !self.program_output_is_valid(&input_accounts, &program_output.accounts_post) { + // Assert accounts pre- and post-states preserve chains invariants + if !inputs_outputs_preserve_invariants(&input_accounts, &program_output.accounts_post) { + return Err(Error::BadInput); + } + + // Fail if any of the output accounts is not yet registered. + // This is redundant with previous checks, but better make it explicit. + if !program_output + .accounts_post + .iter() + .all(|account| self.accounts.contains_key(&account.address)) + { return Err(Error::BadInput); } @@ -32,39 +42,4 @@ impl MockedSequencer { }); Ok(()) } - - /// Verifies that a program public execution didn't break the chain's rules. - /// `input_accounts` are the accounts provided as inputs to the program. - /// `output_accounts` are the accounts post states after execution of the program - fn program_output_is_valid(&self, input_accounts: &[Account], output_accounts: &[Account]) -> bool { - // Fail if the number of input and output accounts differ - if input_accounts.len() != output_accounts.len() { - return false; - } - - for (account_pre, account_post) in input_accounts.iter().zip(output_accounts) { - // Fail if the program modified the addresses of the input accounts - if account_pre.address != account_post.address { - return false; - } - // Fail if the program modified the nonces of the input accounts - if account_pre.nonce != account_post.nonce { - return false; - } - // Fail if any of the output accounts is not yet registered. - // (redundant with previous checks, but better make it explicit) - if !self.accounts.contains_key(&account_post.address) { - return false; - } - } - - // Fail if the execution didn't preserve the total supply. - let total_balance_pre: u128 = input_accounts.iter().map(|account| account.balance).sum(); - let total_balance_post: u128 = output_accounts.iter().map(|account| account.balance).sum(); - if total_balance_pre != total_balance_post { - 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 b415f29..0c5c747 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,5 +1,5 @@ use core::{ - compute_nullifier, hash, is_in_tree, + compute_nullifier, hash, inputs_outputs_preserve_invariants, is_in_tree, types::{Nonce, PrivacyExecutionOutput, ProgramId, ProgramOutput}, visibility::AccountVisibility, }; @@ -31,7 +31,6 @@ fn main() { // Read inner program output let inner_program_output: ProgramOutput = env::read(); let num_inputs = inner_program_output.accounts_pre.len(); - assert_eq!(inner_program_output.accounts_post.len(), num_inputs); // Read visibilities let account_visibilities: Vec = env::read(); @@ -45,14 +44,14 @@ fn main() { let commitment_tree_root: [u32; 8] = env::read(); let program_id: ProgramId = env::read(); - // Verify pre states and post states of accounts are consistent - // with the execution of the `program_id` program - env::verify(program_id, &to_vec(&inner_program_output).unwrap()).unwrap(); - - let inputs = inner_program_output.accounts_pre; - let mut outputs = inner_program_output.accounts_post; + // Authentication step: + // Check private accounts are owned by caller and that they are consistent + // with the commitments tree. let mut nullifiers = Vec::new(); - for (visibility, input_account) in account_visibilities.iter().zip(inputs.iter()) { + for (visibility, input_account) in account_visibilities + .iter() + .zip(inner_program_output.accounts_pre.iter()) + { match visibility { AccountVisibility::Private(Some((private_key, auth_path))) => { // Prove ownership of input accounts by proving knowledge of the pre-image of their addresses. @@ -75,28 +74,29 @@ fn main() { } } - // Assert that the inner program didn't modify address fields or nonces - for (account_pre, account_post) in inputs.iter().zip(outputs.iter()) { - assert_eq!(account_pre.address, account_post.address); - assert_eq!(account_pre.nonce, account_post.nonce); - } + // Verify pre states and post states of accounts are consistent + // with the execution of the `program_id` program + env::verify(program_id, &to_vec(&inner_program_output).unwrap()).unwrap(); - // Check that the program preserved the total supply - let total_balance_pre: u128 = inputs.iter().map(|account| account.balance).sum(); - let total_balance_post: u128 = outputs.iter().map(|account| account.balance).sum(); - assert_eq!(total_balance_pre, total_balance_post); + // Assert accounts pre- and post-states preserve chains invariants + assert!(inputs_outputs_preserve_invariants( + &inner_program_output.accounts_pre, + &inner_program_output.accounts_post + )); // From this point on the execution is considered valid // // Insert new nonces in outputs (including public ones) - 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); // Compute commitments for every private output let mut private_outputs = Vec::new(); - for (output, visibility) in outputs.iter().zip(account_visibilities.iter()) { + for (output, visibility) in accounts_post.iter().zip(account_visibilities.iter()) { match visibility { AccountVisibility::Public => continue, AccountVisibility::Private(_) => private_outputs.push(output), @@ -107,8 +107,10 @@ fn main() { // Get the list of public accounts pre and post states let mut public_accounts_pre = Vec::new(); let mut public_accounts_post = Vec::new(); - for ((account_pre, account_post), visibility) in - inputs.into_iter().zip(outputs.into_iter()).zip(account_visibilities) + for ((account_pre, account_post), visibility) in accounts_pre + .into_iter() + .zip(accounts_post.into_iter()) + .zip(account_visibilities) { match visibility { AccountVisibility::Public => {