From 362749b1b6bacd5189e64b5362ff878bf3246793 Mon Sep 17 00:00:00 2001 From: David Rusu Date: Fri, 23 Aug 2024 03:06:15 +0400 Subject: [PATCH] goas: MMR for the zone transaction log --- goas/atomic_asset_transfer/common/src/lib.rs | 65 +++---- goas/atomic_asset_transfer/common/src/mmr.rs | 162 ++++++++++++++++++ .../atomic_asset_transfer/executor/src/lib.rs | 56 +++--- .../executor/tests/atomic_transfer.rs | 13 +- .../executor/tests/deposit_ptx.rs | 10 +- .../executor/tests/withdraw_ptx.rs | 10 +- .../proof_statements/src/user_note.rs | 8 +- .../risc0_proofs/zone_state/src/main.rs | 2 +- 8 files changed, 234 insertions(+), 92 deletions(-) create mode 100644 goas/atomic_asset_transfer/common/src/mmr.rs diff --git a/goas/atomic_asset_transfer/common/src/lib.rs b/goas/atomic_asset_transfer/common/src/lib.rs index 1d15ccf..e351dcd 100644 --- a/goas/atomic_asset_transfer/common/src/lib.rs +++ b/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, - pub included_txs: Vec, + 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::(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 = 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::>(); - merkle::padded_leaves(&tx_bytes) - } } impl From 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, -} - -impl IncludedTxWitness { - pub fn tx_root(&self) -> [u8; 32] { - let leaf = merkle::leaf(&self.tx.to_bytes()); - merkle::path_root(leaf, &self.path) - } + pub proof: MMRProof, } #[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>(&[ - self.tx_root.to_vec(), - self.zone_id.to_vec(), - self.balance_root.to_vec(), - ]); - StateCommitment(cl::merkle::root(leaves)) + let mut hasher = Sha256::new(); + hasher.update(self.included_txs.commit()); + hasher.update(self.zone_id); + hasher.update(self.balance_root); + StateCommitment(hasher.finalize().into()) } } diff --git a/goas/atomic_asset_transfer/common/src/mmr.rs b/goas/atomic_asset_transfer/common/src/mmr.rs new file mode 100644 index 0000000..44b7423 --- /dev/null +++ b/goas/atomic_asset_transfer/common/src/mmr.rs @@ -0,0 +1,162 @@ +// use ckb_merkle_mountain_range::{util::MemStore, Merge, Result, MMR}; +// use sha2::{Digest, Sha256}; + +// #[derive(Eq, PartialEq, Clone, Debug, Default)] +// struct NumberHash(pub [u8; 32]); +// impl From for NumberHash { +// fn from(num: u32) -> Self { +// let mut hasher = Sha256::new(); +// hasher.update(num.to_le_bytes()); +// NumberHash(hasher.finalize().into()) +// } +// } + +// struct MergeNumberHash; + +// impl Merge for MergeNumberHash { +// type Item = NumberHash; +// fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Result { +// let mut hasher = Sha256::new(); +// hasher.update(lhs.0); +// hasher.update(rhs.0); +// Ok(NumberHash(hasher.finalize().into())) +// } +// } + +// fn prepare_mmr(count: u32) -> (u64, MemStore, Vec) { +// let store = MemStore::default(); +// let mut mmr = MMR::<_, MergeNumberHash, _>::new(0, &store); +// let positions: Vec = (0u32..count) +// .map(|i| mmr.push(NumberHash::from(i)).unwrap()) +// .collect(); +// let mmr_size = mmr.mmr_size(); +// mmr.commit().expect("write to store"); +// (mmr_size, store, positions) +// } + +use cl::merkle; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct MMR { + pub roots: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Root { + pub root: [u8; 32], + pub height: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct MMRProof { + pub path: Vec, +} + +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 u64 { + 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)); + } +} diff --git a/goas/atomic_asset_transfer/executor/src/lib.rs b/goas/atomic_asset_transfer/executor/src/lib.rs index a68d0d0..7cc8f33 100644 --- a/goas/atomic_asset_transfer/executor/src/lib.rs +++ b/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) -> Self { - for tx in txs { - self.state = self.state.apply(tx); - } + pub fn run(mut self, tx: Tx) -> (Self, IncludedTxWitness) { + let (new_state, included_tx) = 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 - .clone() - .run([signed_deposit.bound_tx.tx, signed_withdraw.bound_tx.tx]); + let zone_end = zone_start.clone().run(signed_withdraw.bound_tx.tx).0; let ptx = PartialTxWitness { inputs: vec![ @@ -276,10 +267,7 @@ mod tests { balance_blinding: BalanceWitness::random_blinding(&mut rng), }; - let txs = vec![ - (signed_deposit, ptx.input_witness(0)), - (signed_withdraw, ptx.input_witness(0)), - ]; + let txs = vec![(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)]); - zone_b = zone_b.run([Tx::Deposit(user_intent.deposit)]); + let (zone_a, withdraw_included_witnesss) = zone_a.run(Tx::Withdraw(user_intent.withdraw)); + 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), - deposit_tx: zone_b.state.included_tx_witness(0), + withdraw_tx: withdraw_included_witnesss, + deposit_tx: deposit_included_witnesss, }; let proof = prove_user_atomic_transfer(user_atomic_transfer); diff --git a/goas/atomic_asset_transfer/executor/tests/atomic_transfer.rs b/goas/atomic_asset_transfer/executor/tests/atomic_transfer.rs index 46791f2..29e0fa8 100644 --- a/goas/atomic_asset_transfer/executor/tests/atomic_transfer.rs +++ b/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 - .clone() - .run([Tx::Deposit(alice_intent.deposit)]); + let (zone_b_end, deposit_inclusion) = + zone_b_start.clone().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), - deposit_tx: zone_b_end.state.included_tx_witness(0), + withdraw_tx: withdraw_inclusion, + deposit_tx: deposit_inclusion, }), ), ( diff --git a/goas/atomic_asset_transfer/executor/tests/deposit_ptx.rs b/goas/atomic_asset_transfer/executor/tests/deposit_ptx.rs index b67b534..b535fb0 100644 --- a/goas/atomic_asset_transfer/executor/tests/deposit_ptx.rs +++ b/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() diff --git a/goas/atomic_asset_transfer/executor/tests/withdraw_ptx.rs b/goas/atomic_asset_transfer/executor/tests/withdraw_ptx.rs index c7718b2..eb9f266 100644 --- a/goas/atomic_asset_transfer/executor/tests/withdraw_ptx.rs +++ b/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() diff --git a/goas/atomic_asset_transfer/proof_statements/src/user_note.rs b/goas/atomic_asset_transfer/proof_statements/src/user_note.rs index dd73924..7a5036e 100644 --- a/goas/atomic_asset_transfer/proof_statements/src/user_note.rs +++ b/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_eq!(self.deposit_tx.tx_root(), self.zone_b_roots.tx_root); + assert!(self.zone_a_roots.verify_tx_inclusion(&self.withdraw_tx)); + assert!(self.zone_b_roots.verify_tx_inclusion(&self.deposit_tx)); // ensure the txs are the same ones the user requested assert_eq!( diff --git a/goas/atomic_asset_transfer/risc0_proofs/zone_state/src/main.rs b/goas/atomic_asset_transfer/risc0_proofs/zone_state/src/main.rs index b205142..cf8c3ba 100644 --- a/goas/atomic_asset_transfer/risc0_proofs/zone_state/src/main.rs +++ b/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);