add program owners

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<ProgramId>,
     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],

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.

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();

risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs

@@ -22,7 +22,7 @@ impl MockedSequencer {
             nssa::execute_onchain::<P>(&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);
         }
 

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