From 87dc9d9d3e1d51fdc9314026a4563e4902874024 Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Sun, 20 Jul 2025 21:55:21 -0300 Subject: [PATCH] add program owners --- .../core/src/account.rs | 4 +++- risc0-selective-privacy-poc/core/src/lib.rs | 22 +++++++++++++++++-- .../mocked_components/sequencer/mod.rs | 10 ++++++++- .../sequencer/process_public_execution.rs | 2 +- .../program_methods/guest/src/bin/outer.rs | 3 ++- 5 files changed, 35 insertions(+), 6 deletions(-) diff --git a/risc0-selective-privacy-poc/core/src/account.rs b/risc0-selective-privacy-poc/core/src/account.rs index edcbc44..088b7ee 100644 --- a/risc0-selective-privacy-poc/core/src/account.rs +++ b/risc0-selective-privacy-poc/core/src/account.rs @@ -1,6 +1,6 @@ use crate::{ hash, - types::{Address, Commitment, Key, Nonce}, + types::{Address, Commitment, Key, Nonce, ProgramId}, }; use risc0_zkvm::serde::to_vec; use serde::{Deserialize, Serialize}; @@ -8,6 +8,7 @@ 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 program_owner: Option, pub address: Address, pub balance: u128, pub nonce: Nonce, @@ -16,6 +17,7 @@ pub struct Account { impl Account { pub fn new(address: Address, balance: u128) -> Self { Self { + program_owner: None, address, balance, nonce: [0; 8], diff --git a/risc0-selective-privacy-poc/core/src/lib.rs b/risc0-selective-privacy-poc/core/src/lib.rs index 0b8252d..85b6b1f 100644 --- a/risc0-selective-privacy-poc/core/src/lib.rs +++ b/risc0-selective-privacy-poc/core/src/lib.rs @@ -2,7 +2,10 @@ pub mod account; pub mod types; pub mod visibility; -use crate::{account::Account, types::{AuthenticationPath, Commitment, Key, Nullifier}}; +use crate::{ + account::Account, + types::{AuthenticationPath, Commitment, Key, Nullifier, ProgramId}, +}; use risc0_zkvm::sha::{Impl, Sha256}; pub fn hash(bytes: &[u32]) -> [u32; 8] { @@ -53,7 +56,11 @@ pub fn bytes_to_words(bytes: &[u8; 32]) -> [u32; 8] { /// 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 { +pub fn inputs_outputs_preserve_invariants( + input_accounts: &[Account], + output_accounts: &[Account], + program_id: ProgramId, +) -> bool { // Fail if the number of input and output accounts differ if input_accounts.len() != output_accounts.len() { return false; @@ -68,6 +75,17 @@ pub fn inputs_outputs_preserve_invariants(input_accounts: &[Account], output_acc if account_pre.nonce != account_post.nonce { return false; } + + // Fail if the program modified the program owner + if account_pre.program_owner != account_post.program_owner { + return false; + } + + // Fail if the program subtracted balance from an account it doesn't own. + // (This check always passes if `program_owner` is `None`) + if account_pre.balance > account_post.balance && account_pre.program_owner.unwrap_or(program_id) != program_id { + return false; + } } // Fail if the execution didn't preserve the total supply. 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 ce950f0..e66bf83 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs @@ -5,6 +5,7 @@ use core::{ }; use std::collections::{BTreeMap, HashSet}; +use nssa::{program::PinataProgram, Program}; use sparse_merkle_tree::SparseMerkleTree; use crate::mocked_components::USER_CLIENTS; @@ -33,7 +34,14 @@ impl MockedSequencer { .map(|account| (account.address, account)) .collect(); - let pinata_account = Account::new(PINATA_ADDRESS, INITIAL_BALANCE); + let pinata_account = { + let mut this = Account::new(PINATA_ADDRESS, INITIAL_BALANCE); + // Set the owner of the Pinata account so that only the Pinata program + // can reduce its balance. + this.program_owner = Some(PinataProgram::PROGRAM_ID); + this + }; + accounts.insert(pinata_account.address, pinata_account); let commitment_tree = SparseMerkleTree::new_empty(); 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 7fc0df9..7a67911 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 @@ -22,7 +22,7 @@ impl MockedSequencer { nssa::execute_onchain::

(&input_accounts, instruction_data).map_err(|_| Error::BadInput)?; // Assert accounts pre- and post-states preserve chains invariants - if !inputs_outputs_preserve_invariants(&input_accounts, &program_output.accounts_post) { + if !inputs_outputs_preserve_invariants(&input_accounts, &program_output.accounts_post, P::PROGRAM_ID) { return Err(Error::BadInput); } 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 97c35b6..7c877f7 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 @@ -54,7 +54,8 @@ fn main() { // Assert accounts pre- and post-states preserve chains invariants assert!(inputs_outputs_preserve_invariants( &inner_program_output.accounts_pre, - &inner_program_output.accounts_post + &inner_program_output.accounts_post, + program_id )); // From this point on the execution is considered valid