Refactor zone auth (#19)

* Refactor zone auth

* remove redundant check

goas/atomic_asset_transfer/common/src/events.rs

@@ -1,36 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
-pub enum Event {
-    Spend(Spend),
-}
-
-impl Event {
-    pub fn to_bytes(&self) -> Vec<u8> {
-        // TODO: add variant tag to byte encoding
-        match self {
-            Event::Spend(spend) => spend.to_bytes().to_vec(),
-        }
-    }
-}
-
-/// An event that authorizes spending zone funds
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
-pub struct Spend {
-    pub amount: u64,
-    /// The public key of the recipient
-    pub to: cl::NullifierCommitment,
-    /// The nullifier of note that is being spent, this is to avoid using the spend event to
-    /// for multiple notes
-    pub fund_nf: cl::Nullifier,
-}
-
-impl Spend {
-    pub fn to_bytes(&self) -> [u8; 72] {
-        let mut bytes = [0; 72];
-        bytes[0..8].copy_from_slice(&self.amount.to_le_bytes());
-        bytes[8..40].copy_from_slice(self.to.as_bytes());
-        bytes[40..72].copy_from_slice(self.fund_nf.as_bytes());
-        bytes
-    }
-}

goas/atomic_asset_transfer/common/src/lib.rs

@@ -1,11 +1,4 @@
-pub mod events;
+use cl::{balance::Unit, nullifier::NullifierCommitment};
-
-use cl::{
-    balance::Unit,
-    input::InputWitness,
-    nullifier::{Nullifier, NullifierCommitment},
-    output::OutputWitness,
-};
 use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
@@ -46,8 +39,8 @@ impl ZoneMetadata {
 pub struct StateWitness {
     pub balances: BTreeMap<AccountId, u64>,
     pub included_txs: Vec<Input>,
-    pub output_events: Vec<events::Event>,
     pub zone_metadata: ZoneMetadata,
+    pub nonce: [u8; 32],
 }
 
 impl StateWitness {
@@ -56,10 +49,10 @@ impl StateWitness {
     ///              /        \
     ///            io          state
     ///          /   \        /     \
-    ///      events   txs   zoneid  balances
+    ///      nonce   txs   zoneid  balances
     pub fn commit(&self) -> StateCommitment {
         let root = cl::merkle::root([
-            self.events_root(),
+            self.nonce,
             self.included_txs_root(),
             self.zone_metadata.id(),
             self.balances_root(),
@@ -74,8 +67,7 @@ impl StateWitness {
         let Withdraw {
             from,
             amount,
-            to,
+            to: _,
-            fund_nf,
         } = w;
 
         let from_balance = self.balances.entry(from).or_insert(0);
@@ -83,20 +75,19 @@ impl StateWitness {
             .checked_sub(amount)
             .expect("insufficient funds in account");
 
-        let spend_auth = events::Spend {
-            amount,
-            to,
-            fund_nf,
-        };
-
-        self.output_events.push(events::Event::Spend(spend_auth));
         self
     }
 
-    pub fn events_root(&self) -> [u8; 32] {
+    pub fn deposit(mut self, d: Deposit) -> Self {
-        let event_bytes = Vec::from_iter(self.output_events.iter().map(events::Event::to_bytes));
+        self.included_txs.push(Input::Deposit(d));
-        let event_merkle_leaves = cl::merkle::padded_leaves(&event_bytes);
+
-        cl::merkle::root::<MAX_EVENTS>(event_merkle_leaves)
+        let Deposit { to, amount } = d;
+
+        let to_balance = self.balances.entry(to).or_insert(0);
+        *to_balance += to_balance
+            .checked_add(amount)
+            .expect("overflow in account balance");
+        self
     }
 
     pub fn included_txs_root(&self) -> [u8; 32] {
@@ -117,11 +108,21 @@ impl StateWitness {
         cl::merkle::root::<MAX_BALANCES>(balance_merkle_leaves)
     }
 
-    pub fn event_merkle_path(&self, event: events::Event) -> Vec<cl::merkle::PathNode> {
+    pub fn total_balance(&self) -> u64 {
-        let idx = self.output_events.iter().position(|e| e == &event).unwrap();
+        self.balances.values().sum()
-        let event_bytes = Vec::from_iter(self.output_events.iter().map(events::Event::to_bytes));
+    }
-        let event_merkle_leaves = cl::merkle::padded_leaves(&event_bytes);
+
-        cl::merkle::path::<MAX_EVENTS>(event_merkle_leaves, idx)
+    pub fn evolve_nonce(self) -> Self {
+        let updated_nonce = {
+            let mut hasher = Sha256::new();
+            hasher.update(&self.nonce);
+            hasher.update(b"NOMOS_ZONE_NONCE_EVOLVE");
+            hasher.finalize().into()
+        };
+        Self {
+            nonce: updated_nonce,
+            ..self
+        }
     }
 }
 
@@ -136,24 +137,14 @@ pub struct Withdraw {
     pub from: AccountId,
     pub amount: u64,
     pub to: NullifierCommitment,
-    pub fund_nf: Nullifier,
 }
 
 impl Withdraw {
-    pub fn to_event(&self) -> events::Spend {
+    pub fn to_bytes(&self) -> [u8; 44] {
-        events::Spend {
+        let mut bytes = [0; 44];
-            amount: self.amount,
-            to: self.to,
-            fund_nf: self.fund_nf,
-        }
-    }
-
-    pub fn to_bytes(&self) -> [u8; 72] {
-        let mut bytes = [0; 72];
         bytes[0..4].copy_from_slice(&self.from.to_le_bytes());
-        bytes[4..8].copy_from_slice(&self.amount.to_le_bytes());
+        bytes[4..12].copy_from_slice(&self.amount.to_le_bytes());
-        bytes[8..40].copy_from_slice(self.to.as_bytes());
+        bytes[12..44].copy_from_slice(self.to.as_bytes());
-        bytes[40..72].copy_from_slice(self.fund_nf.as_bytes());
         bytes
     }
 }
@@ -161,16 +152,17 @@ impl Withdraw {
 /// A deposit of funds into the zone
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Deposit {
-    /// The note that is used to deposit funds into the zone
+    pub to: AccountId,
-    pub deposit: InputWitness,
+    pub amount: u64,
+}
 
-    // This zone state note
+impl Deposit {
-    pub zone_note_in: InputWitness,
+    pub fn to_bytes(&self) -> [u8; 32] {
-    pub zone_note_out: OutputWitness,
+        let mut bytes = [0; 32];
-
+        bytes[0..4].copy_from_slice(&self.to.to_le_bytes());
-    // The zone funds note
+        bytes[4..12].copy_from_slice(&self.amount.to_le_bytes());
-    pub zone_funds_in: InputWitness,
+        bytes
-    pub zone_funds_out: OutputWitness,
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]

goas/atomic_asset_transfer/executor/src/lib.rs

@@ -1,17 +1,24 @@
-use common::{events::Event, Input, StateWitness};
+use common::{Input, StateWitness};
 use goas_proof_statements::{zone_funds::SpendFundsPrivate, zone_state::ZoneStatePrivate};
+use std::collections::VecDeque;
 
 pub fn prove_zone_stf(
     state: StateWitness,
     inputs: Vec<Input>,
     zone_in: cl::PartialTxInputWitness,
     zone_out: cl::PartialTxOutputWitness,
+    funds_out: cl::PartialTxOutputWitness,
+    withdrawals: VecDeque<cl::PartialTxOutputWitness>,
+    deposits: VecDeque<cl::PartialTxInputWitness>,
 ) -> ledger::DeathProof {
     let private_inputs = ZoneStatePrivate {
         state,
         inputs,
         zone_in,
         zone_out,
+        funds_out,
+        withdrawals,
+        deposits,
     };
 
     let env = risc0_zkvm::ExecutorEnv::builder()
@@ -36,21 +43,12 @@ pub fn prove_zone_stf(
 pub fn prove_zone_fund_withdraw(
     in_zone_funds: cl::PartialTxInputWitness,
     zone_note: cl::PartialTxOutputWitness,
-    out_zone_funds: cl::PartialTxOutputWitness,
-    spent_note: cl::PartialTxOutputWitness,
     out_zone_state: &StateWitness,
-    withdraw: common::Withdraw,
 ) -> ledger::DeathProof {
-    let spend_event = withdraw.to_event();
     let private_inputs = SpendFundsPrivate {
         in_zone_funds,
         zone_note,
-        out_zone_funds,
+        state_witness: out_zone_state.clone(),
-        spent_note,
-        spend_event,
-        spend_event_state_path: out_zone_state.event_merkle_path(Event::Spend(spend_event)),
-        balances_root: out_zone_state.balances_root(),
-        txs_root: out_zone_state.included_txs_root(),
     };
 
     let env = risc0_zkvm::ExecutorEnv::builder()

goas/atomic_asset_transfer/executor/tests/withdraw_ptx.rs

@@ -1,7 +1,7 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, VecDeque};
 
 use cl::{NoteWitness, NullifierNonce, NullifierSecret};
-use common::{events::Event, Input, StateWitness, ZoneMetadata, ZONE_CL_FUNDS_UNIT};
+use common::{Input, StateWitness, ZoneMetadata, ZONE_CL_FUNDS_UNIT};
 use ledger::death_constraint::DeathProof;
 use rand_core::CryptoRngCore;
 
@@ -53,26 +53,25 @@ fn test_withdrawal() {
     let init_state = StateWitness {
         balances: BTreeMap::from_iter([(alice, 100)]),
         included_txs: vec![],
-        output_events: vec![],
         zone_metadata: ZoneMetadata {
             zone_vk: zone_state_death_constraint(),
             funds_vk: zone_fund_death_constraint(),
             unit: cl::note::unit_point("ZONE_STATE"),
         },
+        nonce: [0; 32],
     };
 
     let zone_fund_in =
-        cl::InputWitness::public(zone_fund_utxo(35240, init_state.zone_metadata, &mut rng));
+        cl::InputWitness::public(zone_fund_utxo(100, init_state.zone_metadata, &mut rng));
     let zone_state_in = cl::InputWitness::public(zone_state_utxo(&init_state, &mut rng));
 
     let withdraw = common::Withdraw {
         from: alice,
         amount: 78,
         to: alice_sk.commit(),
-        fund_nf: zone_fund_in.nullifier(),
     };
 
-    let end_state = init_state.clone().withdraw(withdraw);
+    let end_state = init_state.clone().withdraw(withdraw).evolve_nonce();
 
     let zone_state_out = cl::OutputWitness::public(
         cl::NoteWitness {
@@ -86,7 +85,7 @@ fn test_withdrawal() {
             value: zone_fund_in.note.value - withdraw.amount,
             ..zone_fund_in.note
         },
-        zone_fund_in.evolved_nonce(),
+        NullifierNonce::from_bytes(end_state.nonce),
     );
 
     let alice_withdrawal = cl::OutputWitness::random(
@@ -112,6 +111,9 @@ fn test_withdrawal() {
                 vec![Input::Withdraw(withdraw)],
                 withdraw_ptx.input_witness(0), // input state note (input #0)
                 withdraw_ptx.output_witness(0), // output state note (output #0)
+                withdraw_ptx.output_witness(1), // output funds note (output #1)
+                VecDeque::from_iter([withdraw_ptx.output_witness(2)]), // alice withdrawal
+                VecDeque::new(),               // no deposits
             ),
         ),
         (
@@ -119,10 +121,7 @@ fn test_withdrawal() {
             executor::prove_zone_fund_withdraw(
                 withdraw_ptx.input_witness(1),  // input fund note (input #1)
                 withdraw_ptx.output_witness(0), // output state note (output #0)
-                withdraw_ptx.output_witness(1), // output state note (output #0)
-                withdraw_ptx.output_witness(2), // output state note (output #0)
                 &end_state,
-                withdraw,
             ),
         ),
     ]);
@@ -144,12 +143,8 @@ fn test_withdrawal() {
         StateWitness {
             balances: BTreeMap::from_iter([(alice, 22)]),
             included_txs: vec![Input::Withdraw(withdraw)],
-            output_events: vec![Event::Spend(common::events::Spend {
+            zone_metadata: init_state.zone_metadata,
-                amount: 78,
+            nonce: init_state.evolve_nonce().nonce,
-                to: alice_sk.commit(),
-                fund_nf: zone_fund_in.nullifier()
-            })],
-            zone_metadata: init_state.zone_metadata
         }
         .commit()
         .0

goas/atomic_asset_transfer/proof_statements/src/zone_funds.rs

@@ -6,16 +6,6 @@ pub struct SpendFundsPrivate {
     pub in_zone_funds: cl::PartialTxInputWitness,
     /// The zone note that is authorizing the spend
     pub zone_note: cl::PartialTxOutputWitness,
-    /// The note that is being created to send the change back to the zone
+    /// The state of the zone
-    pub out_zone_funds: cl::PartialTxOutputWitness,
+    pub state_witness: common::StateWitness,
-    /// The spent funds note
-    pub spent_note: cl::PartialTxOutputWitness,
-    /// The event emitted by the zone that authorizes the spend
-    pub spend_event: common::events::Spend,
-    /// Path to the zone output events root
-    pub spend_event_state_path: Vec<cl::merkle::PathNode>,
-    /// Merkle root of txs included in the zone
-    pub txs_root: [u8; 32],
-    /// Merkle root of balances in the zone
-    pub balances_root: [u8; 32],
 }

goas/atomic_asset_transfer/proof_statements/src/zone_state.rs

@@ -1,5 +1,6 @@
 use common::{Input, StateWitness};
 use serde::{Deserialize, Serialize};
+use std::collections::VecDeque;
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ZoneStatePrivate {
@@ -7,4 +8,13 @@ pub struct ZoneStatePrivate {
     pub inputs: Vec<Input>,
     pub zone_in: cl::PartialTxInputWitness,
     pub zone_out: cl::PartialTxOutputWitness,
+    /// While the absence of birth constraints does not guarantee uniqueness of a note that can be used as
+    /// zone funds, deposits and withdrawals make sure the funds are merged in a single note.
+    /// This means that while there's nothing to prevent creation of notes with the same characteristics of zone
+    /// funds, those would not be tracked by the zone state and can be ignored.
+    pub funds_out: cl::PartialTxOutputWitness,
+    /// Each note is the result of the execution of a withdrawal request
+    pub withdrawals: VecDeque<cl::PartialTxOutputWitness>,
+    /// Each note is providing funds for a deposit request
+    pub deposits: VecDeque<cl::PartialTxInputWitness>,
 }

goas/atomic_asset_transfer/risc0_proofs/spend_zone_funds/src/main.rs

@@ -10,87 +10,23 @@ use risc0_zkvm::guest::env;
 fn main() {
     let SpendFundsPrivate {
         in_zone_funds,
-        out_zone_funds,
         zone_note,
-        spent_note,
+        state_witness,
-        spend_event,
-        spend_event_state_path,
-        txs_root,
-        balances_root,
     } = env::read();
 
     let input_root = in_zone_funds.input_root();
-    let output_root = out_zone_funds.output_root();
+    let output_root = zone_note.output_root();
-
-    assert_eq!(output_root, zone_note.output_root());
-    assert_eq!(output_root, spent_note.output_root());
-    assert_eq!(output_root, out_zone_funds.output_root());
-
-    // ** Assert the spent event was an output of the correct zone stf **
-    // The zone state field is a merkle tree over:
-    //                  root
-    //              /        \
-    //            io          state
-    //          /   \        /     \
-    //      events   txs   zoneid  balances
-    // We need to check that:
-    // 1) There is a valid path from the spend event to the events root
-    // 2) The zone id matches the one in the current funds note state
-    // 3) The witnesses for spend path, txs and balances allow to calculate the correct root
-    let zone_id = in_zone_funds.input.note.state; // TODO: is there more state?
-    let spend_event_leaf = merkle::leaf(&spend_event.to_bytes());
-    let event_root = merkle::path_root(spend_event_leaf, &spend_event_state_path);
-
-    assert_eq!(
-        merkle::root([event_root, txs_root, zone_id, balances_root]),
-        zone_note.output.note.state
-    );
-
-    // Check we return the rest of the funds back to the zone
-    let change = in_zone_funds
-        .input
-        .note
-        .value
-        .checked_sub(spend_event.amount)
-        .unwrap();
-    assert_eq!(out_zone_funds.output.note.value, change);
-    // zone funds output should have the same death constraints as the zone funds input
-    assert_eq!(
-        in_zone_funds.input.note.death_constraint,
-        out_zone_funds.output.note.death_constraint
-    );
-    assert_eq!(
-        in_zone_funds.input.note.unit,
-        out_zone_funds.output.note.unit
-    );
-    // ensure zone fund sk's, blindings and nonces are propagated correctly.
-    assert_eq!(
-        in_zone_funds.input.nf_sk.commit(),
-        out_zone_funds.output.nf_pk
-    );
-    assert_eq!(
-        in_zone_funds.input.balance_blinding,
-        out_zone_funds.output.balance_blinding
-    );
-    assert_eq!(
-        in_zone_funds.input.evolved_nonce(),
-        out_zone_funds.output.nonce,
-    );
-    // the state is propagated
-    assert_eq!(
-        in_zone_funds.input.note.state,
-        out_zone_funds.output.note.state,
-    );
-
-    // check the correct amount of funds is being spent
-    assert_eq!(spent_note.output.note.value, spend_event.amount);
-    assert_eq!(spent_note.output.note.unit, in_zone_funds.input.note.unit);
-    // check the correct recipient is being paid
-    assert_eq!(spent_note.output.nf_pk, spend_event.to);
-
-    let nf = in_zone_funds.input.nullifier();
-    assert_eq!(nf, spend_event.fund_nf); // ensure this event was meant for this note.
 
     let ptx_root = PtxRoot(merkle::node(input_root, output_root));
+
+    // 1) Check the zone note is the correct one
+    assert_eq!(
+        in_zone_funds.input.note.state,
+        state_witness.zone_metadata.id()
+    );
+    assert_eq!(zone_note.output.note.state, state_witness.commit().0);
+
+    let nf = in_zone_funds.input.nullifier();
+
     env::commit(&DeathConstraintPublic { ptx_root, nf });
 }

goas/atomic_asset_transfer/risc0_proofs/zone_state/src/main.rs

@@ -1,7 +1,8 @@
 use cl::{
-    merkle,
+    note::NoteWitness,
-    nullifier::{Nullifier, NullifierSecret},
+    nullifier::NullifierNonce,
-    partial_tx::{MAX_INPUTS, MAX_OUTPUTS},
+    output::OutputWitness,
+    partial_tx::{PartialTxInputWitness, PartialTxOutputWitness},
     PtxRoot,
 };
 
@@ -10,114 +11,70 @@ use goas_proof_statements::zone_state::ZoneStatePrivate;
 use ledger_proof_statements::death_constraint::DeathConstraintPublic;
 use risc0_zkvm::guest::env;
 
-fn deposit(
+fn withdraw(
-    mut state: StateWitness,
+    state: StateWitness,
-    deposit: Deposit,
+    output_root: [u8; 32],
-    pub_inputs: DeathConstraintPublic,
+    withdrawal_req: Withdraw,
+    withdrawal: PartialTxOutputWitness,
 ) -> StateWitness {
-    state.included_txs.push(Input::Deposit(deposit.clone()));
+    // 1) check the correct amount of funds is being spent
+    assert_eq!(withdrawal.output.note.value, withdrawal_req.amount);
+    assert_eq!(withdrawal.output.note.unit, *ZONE_CL_FUNDS_UNIT);
+    // 2) check the correct recipient is being paid
+    assert_eq!(withdrawal.output.nf_pk, withdrawal_req.to);
 
-    let Deposit {
+    assert_eq!(output_root, withdrawal.output_root());
-        deposit,
-        zone_note_in,
-        zone_note_out,
-        zone_funds_in,
-        zone_funds_out,
-    } = deposit;
 
-    let funds_vk = state.zone_metadata.funds_vk;
+    state.withdraw(withdrawal_req)
+}
 
-    // 1) Check there are no more input/output notes than expected
+fn deposit(
-    let inputs = [
+    state: StateWitness,
-        deposit.commit().to_bytes().to_vec(),
+    input_root: [u8; 32],
-        zone_note_in.commit().to_bytes().to_vec(),
+    deposit_req: Deposit,
-        zone_funds_in.commit().to_bytes().to_vec(),
+    deposit: PartialTxInputWitness,
-    ];
+) -> StateWitness {
+    assert_eq!(deposit.input_root(), input_root);
 
-    let inputs_root = merkle::root(merkle::padded_leaves::<MAX_INPUTS>(&inputs));
+    // 1) Check the deposit note is not already under control of the zone
-
+    assert_ne!(
-    let outputs = [
+        deposit.input.note.death_constraint,
-        zone_note_out.commit().to_bytes().to_vec(),
+        state.zone_metadata.funds_vk
-        zone_funds_out.commit().to_bytes().to_vec(),
-    ];
-
-    let outputs_root = merkle::root(merkle::padded_leaves::<MAX_OUTPUTS>(&outputs));
-
-    let ptx_root = PtxRoot(merkle::node(inputs_root, outputs_root));
-    assert_eq!(ptx_root, pub_inputs.ptx_root);
-
-    // 2) Check the deposit note is not already under control of the zone
-    assert_ne!(deposit.note.death_constraint, funds_vk);
-
-    // 3) Check the ptx is balanced. This is not a requirement for standard ptxs, but we need it
-    //    in deposits (at least in a first version) to ensure fund tracking
-    assert_eq!(deposit.note.unit, *ZONE_CL_FUNDS_UNIT);
-    assert_eq!(zone_funds_in.note.unit, *ZONE_CL_FUNDS_UNIT);
-    assert_eq!(zone_funds_out.note.unit, *ZONE_CL_FUNDS_UNIT);
-
-    let in_sum = deposit.note.value + zone_funds_in.note.value;
-
-    let out_sum = zone_note_out.note.value;
-
-    assert_eq!(out_sum, in_sum, "deposit ptx is unbalanced");
-
-    // 4) Check the zone fund notes are correctly created
-    assert_eq!(zone_funds_in.note.death_constraint, funds_vk);
-    assert_eq!(zone_funds_out.note.death_constraint, funds_vk);
-    assert_eq!(zone_funds_in.note.state, state.zone_metadata.id());
-    assert_eq!(zone_funds_out.note.state, state.zone_metadata.id());
-    assert_eq!(zone_funds_in.nf_sk, NullifierSecret::from_bytes([0; 16])); // there is no secret in the zone funds
-    assert_eq!(zone_funds_out.nf_pk, zone_funds_in.nf_sk.commit()); // the sk is the same
-                                                                    // nonce is correctly evolved
-    assert_eq!(zone_funds_out.nonce, zone_funds_in.evolved_nonce());
-
-    // 5) Check zone state notes are correctly created
-    assert_eq!(
-        zone_note_in.note.death_constraint,
-        zone_note_out.note.death_constraint
     );
-    assert_eq!(zone_note_in.nf_sk, NullifierSecret::from_bytes([0; 16])); //// there is no secret in the zone state
-    assert_eq!(zone_note_out.nf_pk, zone_note_in.nf_sk.commit()); // the sk is the same
-    assert_eq!(zone_note_in.note.unit, zone_note_out.note.unit);
-    assert_eq!(zone_note_in.note.value, zone_note_out.note.value);
-    // nonce is correctly evolved
-    assert_eq!(zone_note_out.nonce, zone_note_in.evolved_nonce());
-    let nullifier = Nullifier::new(zone_note_in.nf_sk, zone_note_in.nonce);
-    assert_eq!(nullifier, pub_inputs.nf);
 
-    // 6) We're now ready to do the deposit!
+    // 2) Check the deposit note is for the correct amount
-    let amount = deposit.note.value;
+    assert_eq!(deposit.input.note.unit, *ZONE_CL_FUNDS_UNIT);
-    let to = AccountId::from_be_bytes(<[u8; 4]>::try_from(&deposit.note.state[0..4]).unwrap());
+    assert_eq!(deposit.input.note.value, deposit_req.amount);
 
-    let to_balance = state.balances.entry(to).or_insert(0);
+    // 3) Check the deposit note is for the correct recipient
-    *to_balance = to_balance
+    assert_eq!(
-        .checked_add(amount)
+        AccountId::from_le_bytes(<[u8; 4]>::try_from(&deposit.input.note.state[..4]).unwrap()),
-        .expect("overflow when depositing");
+        deposit_req.to
+    );
 
-    state
+    state.deposit(deposit_req)
 }
 
 fn validate_zone_transition(
     in_note: cl::PartialTxInputWitness,
     out_note: cl::PartialTxOutputWitness,
-    in_meta: ZoneMetadata,
+    out_funds: cl::PartialTxOutputWitness,
     in_state_cm: StateCommitment,
     out_state: StateWitness,
 ) {
+    let metadata = out_state.zone_metadata;
+    let out_state_cm = out_state.commit().0;
     // Ensure input/output notes are committing to the expected states.
     assert_eq!(in_note.input.note.state, in_state_cm.0);
-    assert_eq!(out_note.output.note.state, out_state.commit().0);
+    assert_eq!(out_note.output.note.state, out_state_cm);
-
-    // zone metadata is propagated
-    assert_eq!(out_state.zone_metadata.id(), in_meta.id());
 
     // ensure units match metadata
-    assert_eq!(in_note.input.note.unit, in_meta.unit);
+    assert_eq!(in_note.input.note.unit, metadata.unit);
-    assert_eq!(out_note.output.note.unit, in_meta.unit);
+    assert_eq!(out_note.output.note.unit, metadata.unit);
 
     // ensure constraints match metadata
-    assert_eq!(in_note.input.note.death_constraint, in_meta.zone_vk);
+    assert_eq!(in_note.input.note.death_constraint, metadata.zone_vk);
-    assert_eq!(out_note.output.note.death_constraint, in_meta.zone_vk);
+    assert_eq!(out_note.output.note.death_constraint, metadata.zone_vk);
 
     // nullifier secret is propagated
     assert_eq!(in_note.input.nf_sk.commit(), out_note.output.nf_pk);
@@ -130,6 +87,23 @@ fn validate_zone_transition(
 
     // the nonce is correctly evolved
     assert_eq!(in_note.input.evolved_nonce(), out_note.output.nonce);
+
+    // funds are still under control of the zone
+    let expected_note_witness = NoteWitness::new(
+        out_state.total_balance(),
+        *ZONE_CL_FUNDS_UNIT,
+        metadata.funds_vk,
+        metadata.id(),
+    );
+    assert_eq!(
+        out_funds.output,
+        OutputWitness::public(
+            expected_note_witness,
+            NullifierNonce::from_bytes(out_state.nonce)
+        )
+    );
+    // funds belong to the same partial tx
+    assert_eq!(out_funds.output_root(), out_note.output_root());
 }
 
 fn main() {
@@ -138,27 +112,30 @@ fn main() {
         inputs,
         zone_in,
         zone_out,
+        funds_out,
+        mut withdrawals,
+        mut deposits,
     } = env::read();
 
+    let input_root = zone_in.input_root();
+    let output_root = zone_out.output_root();
+
     let pub_inputs = DeathConstraintPublic {
-        ptx_root: PtxRoot(cl::merkle::node(
+        ptx_root: PtxRoot(cl::merkle::node(input_root, output_root)),
-            zone_in.input_root(),
-            zone_out.output_root(),
-        )),
         nf: zone_in.input.nullifier(),
     };
 
-    let in_meta = state.zone_metadata;
     let in_state_cm = state.commit();
 
     for input in inputs {
         state = match input {
-            Input::Withdraw(input) => state.withdraw(input),
+            Input::Withdraw(w) => withdraw(state, output_root, w, withdrawals.pop_front().unwrap()),
-            Input::Deposit(input) => deposit(state, input, pub_inputs),
+            Input::Deposit(d) => deposit(state, input_root, d, deposits.pop_front().unwrap()),
         }
     }
 
-    validate_zone_transition(zone_in, zone_out, in_meta, in_state_cm, state);
+    let state = state.evolve_nonce();
+    validate_zone_transition(zone_in, zone_out, funds_out, in_state_cm, state);
 
     env::commit(&pub_inputs);
 }