Merge pull request #37 from logos-co/goas/mmr-in-zone-state

goas: MMR for the zone transaction log

goas/atomic_asset_transfer/common/src/lib.rs

@@ -1,8 +1,11 @@
+pub mod mmr;
+
 use cl::{balance::Unit, merkle, NoteCommitment};
 use ed25519_dalek::{
     ed25519::{signature::SignerMut, SignatureBytes},
     Signature, SigningKey, VerifyingKey, PUBLIC_KEY_LENGTH,
 };
+use mmr::{MMRProof, MMR};
 use once_cell::sync::Lazy;
 use rand_core::CryptoRngCore;
 use serde::{Deserialize, Serialize};
@@ -47,7 +50,7 @@ impl ZoneMetadata {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct StateWitness {
     pub balances: BTreeMap<AccountId, u64>,
-    pub included_txs: Vec<Tx>,
+    pub included_txs: MMR,
     pub zone_metadata: ZoneMetadata,
 }
 
@@ -58,21 +61,25 @@ impl StateWitness {
 
     pub fn state_roots(&self) -> StateRoots {
         StateRoots {
-            tx_root: self.included_txs_root(),
+            included_txs: self.included_txs.clone(),
             zone_id: self.zone_metadata.id(),
             balance_root: self.balances_root(),
         }
     }
 
-    pub fn apply(self, tx: Tx) -> Self {
+    pub fn apply(self, tx: Tx) -> (Self, IncludedTxWitness) {
         let mut state = match tx {
             Tx::Withdraw(w) => self.withdraw(w),
             Tx::Deposit(d) => self.deposit(d),
         };
 
-        state.included_txs.push(tx);
+        let inclusion_proof = state.included_txs.push(&tx.to_bytes());
+        let tx_inclusion_proof = IncludedTxWitness {
+            tx,
+            proof: inclusion_proof,
+        };
 
-        state
+        (state, tx_inclusion_proof)
     }
 
     fn withdraw(mut self, w: Withdraw) -> Self {
@@ -97,16 +104,6 @@ impl StateWitness {
         self
     }
 
-    pub fn included_txs_root(&self) -> [u8; 32] {
-        merkle::root::<MAX_TXS>(self.included_tx_merkle_leaves())
-    }
-
-    pub fn included_tx_witness(&self, idx: usize) -> IncludedTxWitness {
-        let tx = *self.included_txs.get(idx).unwrap();
-        let path = merkle::path(self.included_tx_merkle_leaves(), idx);
-        IncludedTxWitness { tx, path }
-    }
-
     pub fn balances_root(&self) -> [u8; 32] {
         let balance_bytes = Vec::from_iter(self.balances.iter().map(|(owner, balance)| {
             let mut bytes: Vec<u8> = vec![];
@@ -121,15 +118,6 @@ impl StateWitness {
     pub fn total_balance(&self) -> u64 {
         self.balances.values().sum()
     }
-
-    fn included_tx_merkle_leaves(&self) -> [[u8; 32]; MAX_TXS] {
-        let tx_bytes = self
-            .included_txs
-            .iter()
-            .map(|t| t.to_bytes())
-            .collect::<Vec<_>>();
-        merkle::padded_leaves(&tx_bytes)
-    }
 }
 
 impl From<StateCommitment> for [u8; 32] {
@@ -237,31 +225,28 @@ impl Tx {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct IncludedTxWitness {
     pub tx: Tx,
-    pub path: Vec<merkle::PathNode>,
+    pub proof: MMRProof,
-}
-
-impl IncludedTxWitness {
-    pub fn tx_root(&self) -> [u8; 32] {
-        let leaf = merkle::leaf(&self.tx.to_bytes());
-        merkle::path_root(leaf, &self.path)
-    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct StateRoots {
-    pub tx_root: [u8; 32],
+    pub included_txs: MMR,
     pub zone_id: [u8; 32],
     pub balance_root: [u8; 32],
 }
 
 impl StateRoots {
-    /// Merkle tree over: [txs, zoneid, balances]
+    pub fn verify_tx_inclusion(&self, tx_inclusion: &IncludedTxWitness) -> bool {
+        self.included_txs
+            .verify_proof(&tx_inclusion.tx.to_bytes(), &tx_inclusion.proof)
+    }
+
+    /// Commitment to the state roots
     pub fn commit(&self) -> StateCommitment {
-        let leaves = cl::merkle::padded_leaves::<4>(&[
+        let mut hasher = Sha256::new();
-            self.tx_root.to_vec(),
+        hasher.update(self.included_txs.commit());
-            self.zone_id.to_vec(),
+        hasher.update(self.zone_id);
-            self.balance_root.to_vec(),
+        hasher.update(self.balance_root);
-        ]);
+        StateCommitment(hasher.finalize().into())
-        StateCommitment(cl::merkle::root(leaves))
     }
 }

goas/atomic_asset_transfer/common/src/mmr.rs

@@ -0,0 +1,126 @@
+use cl::merkle;
+use serde::{Deserialize, Serialize};
+use sha2::{Digest, Sha256};
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct MMR {
+    pub roots: Vec<Root>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Root {
+    pub root: [u8; 32],
+    pub height: u8,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct MMRProof {
+    pub path: Vec<merkle::PathNode>,
+}
+
+impl MMR {
+    pub fn new() -> Self {
+        Self { roots: vec![] }
+    }
+
+    pub fn push(&mut self, elem: &[u8]) -> MMRProof {
+        let new_root = Root {
+            root: merkle::leaf(elem),
+            height: 1,
+        };
+        self.roots.push(new_root);
+
+        let mut path = vec![];
+
+        for i in (1..self.roots.len()).rev() {
+            if self.roots[i].height == self.roots[i - 1].height {
+                path.push(merkle::PathNode::Left(self.roots[i - 1].root));
+
+                self.roots[i - 1] = Root {
+                    root: merkle::node(self.roots[i - 1].root, self.roots[i].root),
+                    height: self.roots[i - 1].height + 1,
+                };
+
+                self.roots.remove(i);
+            } else {
+                break;
+            }
+        }
+
+        MMRProof { path }
+    }
+
+    pub fn verify_proof(&self, elem: &[u8], proof: &MMRProof) -> bool {
+        let path_len = proof.path.len();
+        let leaf = merkle::leaf(elem);
+        let root = merkle::path_root(leaf, &proof.path);
+
+        for mmr_root in self.roots.iter() {
+            if mmr_root.height == (path_len + 1) as u8 {
+                return mmr_root.root == root;
+            }
+        }
+
+        false
+    }
+
+    pub fn commit(&self) -> [u8; 32] {
+        let mut hasher = Sha256::new();
+        for mrr_root in self.roots.iter() {
+            hasher.update(mrr_root.root);
+            hasher.update(mrr_root.height.to_le_bytes());
+        }
+        hasher.finalize().into()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_mrr_push() {
+        let mut mmr = MMR::new();
+        let proof = mmr.push(b"hello");
+
+        assert_eq!(mmr.roots.len(), 1);
+        assert_eq!(mmr.roots[0].height, 1);
+        assert_eq!(mmr.roots[0].root, merkle::leaf(b"hello"));
+        assert!(mmr.verify_proof(b"hello", &proof));
+
+        let proof = mmr.push(b"world");
+
+        assert_eq!(mmr.roots.len(), 1);
+        assert_eq!(mmr.roots[0].height, 2);
+        assert_eq!(
+            mmr.roots[0].root,
+            merkle::node(merkle::leaf(b"hello"), merkle::leaf(b"world"))
+        );
+        assert!(mmr.verify_proof(b"world", &proof));
+
+        let proof = mmr.push(b"!");
+
+        assert_eq!(mmr.roots.len(), 2);
+        assert_eq!(mmr.roots[0].height, 2);
+        assert_eq!(
+            mmr.roots[0].root,
+            merkle::node(merkle::leaf(b"hello"), merkle::leaf(b"world"))
+        );
+        assert_eq!(mmr.roots[1].height, 1);
+        assert_eq!(mmr.roots[1].root, merkle::leaf(b"!"));
+        assert!(mmr.verify_proof(b"!", &proof));
+
+        let proof = mmr.push(b"!");
+
+        assert_eq!(mmr.roots.len(), 1);
+        assert_eq!(mmr.roots[0].height, 3);
+        assert_eq!(
+            mmr.roots[0].root,
+            merkle::node(
+                merkle::node(merkle::leaf(b"hello"), merkle::leaf(b"world")),
+                merkle::node(merkle::leaf(b"!"), merkle::leaf(b"!"))
+            )
+        );
+        assert!(mmr.verify_proof(b"!", &proof));
+    }
+}

goas/atomic_asset_transfer/executor/src/lib.rs

@@ -1,6 +1,8 @@
 use std::collections::BTreeMap;
 
-use common::{AccountId, SignedBoundTx, StateWitness, Tx, ZoneMetadata};
+use common::{
+    mmr::MMR, AccountId, IncludedTxWitness, SignedBoundTx, StateWitness, Tx, ZoneMetadata,
+};
 use goas_proof_statements::{
     user_note::UserAtomicTransfer, zone_funds::SpendFundsPrivate, zone_state::ZoneStatePrivate,
 };
@@ -21,7 +23,7 @@ impl ZoneNotes {
     ) -> Self {
         let state = StateWitness {
             balances,
-            included_txs: vec![],
+            included_txs: MMR::new(),
             zone_metadata: zone_metadata(zone_name),
         };
         let state_note = zone_state_utxo(&state, &mut rng);
@@ -41,10 +43,9 @@ impl ZoneNotes {
         cl::InputWitness::public(self.fund_note)
     }
 
-    pub fn run(mut self, txs: impl IntoIterator<Item = Tx>) -> Self {
+    pub fn run(mut self, tx: Tx) -> (Self, IncludedTxWitness) {
-        for tx in txs {
+        let (new_state, included_tx) = self.state.apply(tx);
-            self.state = self.state.apply(tx);
+        self.state = new_state;
-        }
 
         let state_in = self.state_input_witness();
         self.state_note = cl::OutputWitness::public(
@@ -63,7 +64,8 @@ impl ZoneNotes {
             },
             state_in.evolved_nonce(b"FUND_NONCE"),
         );
-        self
+
+        (self, included_tx)
     }
 }
 
@@ -231,26 +233,17 @@ mod tests {
     pub fn test_prove_zone_stf() {
         let mut rng = rand::thread_rng();
 
-        let zone_start = ZoneNotes::new_with_balances("ZONE", BTreeMap::from_iter([]), &mut rng);
+        let mut alice = common::new_account(&mut rng);
+        let alice_vk = alice.verifying_key().to_bytes();
+
+        let zone_start =
+            ZoneNotes::new_with_balances("ZONE", BTreeMap::from_iter([(alice_vk, 32)]), &mut rng);
 
         let bind = OutputWitness::public(
             NoteWitness::basic(32, *common::ZONE_CL_FUNDS_UNIT),
             cl::NullifierNonce::random(&mut rng),
         );
 
-        let mut alice = common::new_account(&mut rng);
-        let alice_vk = alice.verifying_key().to_bytes();
-
-        let signed_deposit = SignedBoundTx::sign(
-            BoundTx {
-                tx: Tx::Deposit(Deposit {
-                    to: alice_vk,
-                    amount: 32,
-                }),
-                bind: bind.commit_note(),
-            },
-            &mut alice,
-        );
         let signed_withdraw = SignedBoundTx::sign(
             BoundTx {
                 tx: Tx::Withdraw(Withdraw {
@@ -262,9 +255,7 @@ mod tests {
             &mut alice,
         );
 
-        let zone_end = zone_start
+        let zone_end = zone_start.clone().run(signed_withdraw.bound_tx.tx).0;
-            .clone()
-            .run([signed_deposit.bound_tx.tx, signed_withdraw.bound_tx.tx]);
 
         let ptx = PartialTxWitness {
             inputs: vec![
@@ -276,10 +267,7 @@ mod tests {
             balance_blinding: BalanceWitness::random_blinding(&mut rng),
         };
 
-        let txs = vec![
+        let txs = vec![(signed_withdraw, ptx.input_witness(0))];
-            (signed_deposit, ptx.input_witness(0)),
-            (signed_withdraw, ptx.input_witness(0)),
-        ];
 
         let proof = prove_zone_stf(
             zone_start.state.clone(),
@@ -322,9 +310,9 @@ mod tests {
         let alice = common::new_account(&mut rng);
         let alice_vk = alice.verifying_key().to_bytes();
 
-        let mut zone_a =
+        let zone_a =
             ZoneNotes::new_with_balances("ZONE_A", BTreeMap::from_iter([(alice_vk, 40)]), &mut rng);
-        let mut zone_b = ZoneNotes::new_with_balances("ZONE_B", BTreeMap::new(), &mut rng);
+        let zone_b = ZoneNotes::new_with_balances("ZONE_B", BTreeMap::new(), &mut rng);
 
         let user_intent = UserIntent {
             zone_a_meta: zone_a.state.zone_metadata,
@@ -343,8 +331,8 @@ mod tests {
             NullifierNonce::random(&mut rng),
         ));
 
-        zone_a = zone_a.run([Tx::Withdraw(user_intent.withdraw)]);
+        let (zone_a, withdraw_included_witnesss) = zone_a.run(Tx::Withdraw(user_intent.withdraw));
-        zone_b = zone_b.run([Tx::Deposit(user_intent.deposit)]);
+        let (zone_b, deposit_included_witnesss) = zone_b.run(Tx::Deposit(user_intent.deposit));
 
         let ptx = PartialTxWitness {
             inputs: vec![user_note],
@@ -359,8 +347,8 @@ mod tests {
             zone_b: ptx.output_witness(1),
             zone_a_roots: zone_a.state.state_roots(),
             zone_b_roots: zone_b.state.state_roots(),
-            withdraw_tx: zone_a.state.included_tx_witness(0),
+            withdraw_tx: withdraw_included_witnesss,
-            deposit_tx: zone_b.state.included_tx_witness(0),
+            deposit_tx: deposit_included_witnesss,
         };
 
         let proof = prove_user_atomic_transfer(user_atomic_transfer);

goas/atomic_asset_transfer/executor/tests/atomic_transfer.rs

@@ -46,13 +46,12 @@ fn test_atomic_transfer() {
         balance_blinding: BalanceWitness::random_blinding(&mut rng),
     };
 
-    let zone_a_end = zone_a_start
+    let (zone_a_end, withdraw_inclusion) = zone_a_start
         .clone()
-        .run([Tx::Withdraw(alice_intent.withdraw)]);
+        .run(Tx::Withdraw(alice_intent.withdraw));
 
-    let zone_b_end = zone_b_start
+    let (zone_b_end, deposit_inclusion) =
-        .clone()
+        zone_b_start.clone().run(Tx::Deposit(alice_intent.deposit));
-        .run([Tx::Deposit(alice_intent.deposit)]);
 
     let alice_intent_in = cl::InputWitness::public(alice_intent_out);
     let atomic_transfer_ptx = cl::PartialTxWitness {
@@ -97,8 +96,8 @@ fn test_atomic_transfer() {
                 zone_b: atomic_transfer_ptx.output_witness(2),
                 zone_a_roots: zone_a_end.state.state_roots(),
                 zone_b_roots: zone_b_end.state.state_roots(),
-                withdraw_tx: zone_a_end.state.included_tx_witness(0),
+                withdraw_tx: withdraw_inclusion,
-                deposit_tx: zone_b_end.state.included_tx_witness(0),
+                deposit_tx: deposit_inclusion,
             }),
         ),
         (

goas/atomic_asset_transfer/executor/tests/deposit_ptx.rs

@@ -1,7 +1,7 @@
 use std::collections::BTreeMap;
 
 use cl::{BalanceWitness, NoteWitness, NullifierSecret};
-use common::{new_account, BoundTx, SignedBoundTx, StateWitness, Tx, ZONE_CL_FUNDS_UNIT};
+use common::{mmr::MMR, new_account, BoundTx, SignedBoundTx, StateWitness, Tx, ZONE_CL_FUNDS_UNIT};
 use executor::ZoneNotes;
 use ledger::death_constraint::DeathProof;
 
@@ -20,7 +20,7 @@ fn test_deposit() {
         amount: 78,
     };
 
-    let zone_end = zone_start.clone().run([Tx::Deposit(deposit)]);
+    let zone_end = zone_start.clone().run(Tx::Deposit(deposit)).0;
 
     let alice_deposit = cl::InputWitness::from_output(
         cl::OutputWitness::random(
@@ -82,7 +82,11 @@ fn test_deposit() {
         zone_end.state_note.note.state,
         StateWitness {
             balances: BTreeMap::from_iter([(alice_vk, 78)]),
-            included_txs: vec![Tx::Deposit(deposit)],
+            included_txs: {
+                let mut mmr = MMR::new();
+                mmr.push(&Tx::Deposit(deposit).to_bytes());
+                mmr
+            },
             zone_metadata: zone_start.state.zone_metadata,
         }
         .commit()

goas/atomic_asset_transfer/executor/tests/withdraw_ptx.rs

@@ -1,7 +1,7 @@
 use std::collections::BTreeMap;
 
 use cl::{BalanceWitness, NoteWitness, NullifierSecret};
-use common::{new_account, BoundTx, SignedBoundTx, StateWitness, Tx, ZONE_CL_FUNDS_UNIT};
+use common::{mmr::MMR, new_account, BoundTx, SignedBoundTx, StateWitness, Tx, ZONE_CL_FUNDS_UNIT};
 use executor::ZoneNotes;
 use ledger::death_constraint::DeathProof;
 
@@ -30,7 +30,7 @@ fn test_withdrawal() {
         amount: 78,
     };
 
-    let zone_end = zone_start.clone().run([Tx::Withdraw(withdraw)]);
+    let zone_end = zone_start.clone().run(Tx::Withdraw(withdraw)).0;
 
     let alice_withdrawal = cl::OutputWitness::random(
         NoteWitness::stateless(
@@ -102,7 +102,11 @@ fn test_withdrawal() {
         zone_end.state_note.note.state,
         StateWitness {
             balances: BTreeMap::from_iter([(alice_vk, 22)]),
-            included_txs: vec![Tx::Withdraw(withdraw)],
+            included_txs: {
+                let mut mmr = MMR::new();
+                mmr.push(&Tx::Withdraw(withdraw).to_bytes());
+                mmr
+            },
             zone_metadata: zone_start.state.zone_metadata,
         }
         .commit()

goas/atomic_asset_transfer/proof_statements/src/user_note.rs

@@ -7,9 +7,9 @@
 ///
 /// The Alice will create a partial tx that looks like this:
 ///
-///     [fee note] -> [user note]
+///   [] -> [user note]
 ///
-/// The User Note will encode the logic that orchestrates the withdrawal from zone A
+/// Thep User Note will encode the logic that orchestrates the withdrawal from zone A
 /// and deposit to zone B.
 ///
 /// The User Notes death constraint requires the following statements to be satisfied
@@ -84,8 +84,8 @@ impl UserAtomicTransfer {
         );
 
         // ensure txs were included in the respective zones
-        assert_eq!(self.withdraw_tx.tx_root(), self.zone_a_roots.tx_root);
+        assert!(self.zone_a_roots.verify_tx_inclusion(&self.withdraw_tx));
-        assert_eq!(self.deposit_tx.tx_root(), self.zone_b_roots.tx_root);
+        assert!(self.zone_b_roots.verify_tx_inclusion(&self.deposit_tx));
 
         // ensure the txs are the same ones the user requested
         assert_eq!(

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

@@ -81,7 +81,7 @@ fn main() {
         assert_eq!(ptx_input_witness.input_root(), input_root);
 
         // apply the ptx
-        state = state.apply(bound_tx.tx)
+        state = state.apply(bound_tx.tx).0
     }
 
     validate_zone_transition(zone_in, zone_out, funds_out, in_state_cm, state);