From 745d65404865e35ca6a1ef841068412923d4bfc2 Mon Sep 17 00:00:00 2001 From: arnaucube Date: Thu, 24 Jul 2025 12:02:44 +0200 Subject: [PATCH] MerkleTree insertion proofs (#344) * implement merkletree insert & insert-proof-verification * add merkletree circuit to verify insertion proof wip * fix merkletree's GraphViz generation for cases with empty siblings * implement tree insert-verif circuit siblings checks Note: I've implemented also an alternative version which instead of inputting a witness value 'divergence_level' it inputs a bitmask. Both approaches (divergence_level and divergence_bitmask) take the same amount of constraints (336 constraints for a tree of 32 levels, and for an hybrid approach it takes 331 constraints but the code gets a bit less readable). So I've kept with the current implementation (using divergence_level) which is more easy to follow. * [tree] modify the strategy for the insert-proof (out-circuit) * re-implement insert-proof verification circuit * add pending checks and polish * add tests for disabled(&enabled) cases that should fail * update typos.toml config * Add test with tampering * add check 5.3, to prevent tampering (at insertion proof circuit) * move old_leaf_hash computation outside the loop, simplify check 5.3 booleans * apply @ed255 review suggestions --------- Co-authored-by: Ahmad --- .github/workflows/typos.toml | 1 + .../plonky2/primitives/merkletree/circuit.rs | 481 +++++++++++++++++- .../plonky2/primitives/merkletree/error.rs | 5 + .../plonky2/primitives/merkletree/mod.rs | 239 ++++++++- 4 files changed, 695 insertions(+), 31 deletions(-) diff --git a/.github/workflows/typos.toml b/.github/workflows/typos.toml index 0c36a42..605ff57 100644 --- a/.github/workflows/typos.toml +++ b/.github/workflows/typos.toml @@ -8,3 +8,4 @@ nin = "nin" # not in kow = "kow" # key or wildcard KOW = "KOW" # Key Or Wildcard datas = "datas" # plural (for 'verifier_datas', a vector of 'verifier_data') +typ = "typ" # from 'type', which can not be used as variable name since it is a keyword diff --git a/src/backends/plonky2/primitives/merkletree/circuit.rs b/src/backends/plonky2/primitives/merkletree/circuit.rs index 52699ab..aa147f5 100644 --- a/src/backends/plonky2/primitives/merkletree/circuit.rs +++ b/src/backends/plonky2/primitives/merkletree/circuit.rs @@ -10,6 +10,7 @@ //! use std::iter; +use itertools::zip_eq; use plonky2::{ field::types::Field, hash::{ @@ -28,7 +29,7 @@ use crate::{ basetypes::D, circuits::common::{CircuitBuilderPod, ValueTarget}, error::Result, - primitives::merkletree::MerkleClaimAndProof, + primitives::merkletree::{MerkleClaimAndProof, MerkleProofStateTransition}, }, measure_gates_begin, measure_gates_end, middleware::{EMPTY_HASH, EMPTY_VALUE, F, HASH_SIZE}, @@ -49,9 +50,10 @@ pub struct MerkleClaimAndProofTarget { pub(crate) other_value: ValueTarget, } -/// Allows to verify both proofs of existence and proofs non-existence with the same circuit. If -/// only proofs of existence are needed, use `verify_merkle_proof_existence_circuit`, which -/// requires less amount of constraints. +/// Allows to verify both proofs of existence and proofs non-existence with the +/// same circuit. If only proofs of existence are needed, use +/// `verify_merkle_proof_existence_circuit`, which requires less amount of +/// constraints. pub fn verify_merkle_proof_circuit( builder: &mut CircuitBuilder, proof: &MerkleClaimAndProofTarget, @@ -202,12 +204,14 @@ pub struct MerkleProofExistenceTarget { pub(crate) siblings: Vec, } -/// Allows to verify proofs of existence only. If proofs of non-existence are needed, use -/// `verify_merkle_proof_circuit`. +/// Allows to verify proofs of existence only. If proofs of non-existence are +/// needed, use `verify_merkle_proof_circuit`. +/// It returns the computed path, in case is needed at other parts of the upper +/// logic to avoid recomputing it again. pub fn verify_merkle_proof_existence_circuit( builder: &mut CircuitBuilder, proof: &MerkleProofExistenceTarget, -) { +) -> Vec { let max_depth = proof.max_depth; let measure = measure_gates_begin!(builder, format!("MerkleProofExist_{}", max_depth)); @@ -233,6 +237,8 @@ pub fn verify_merkle_proof_existence_circuit( builder.connect(computed_root[j], expected_root[j]); } measure_gates_end!(builder, measure); + + path } impl MerkleProofExistenceTarget { @@ -385,6 +391,229 @@ fn kv_hash_target( builder.hash_n_to_hash_no_pad::(inputs) } +/// Verifies that the merkletree state transition (from old_root to new_root) +/// has been done correctly for the given new_key. This will allow verifying +/// correct new leaf insertion, and leaf edition&deletion (if needed). +pub struct MerkleProofStateTransitionTarget { + pub(crate) max_depth: usize, + // `enabled` determines if the merkleproof state transition verification is enabled + pub(crate) enabled: BoolTarget, + pub(crate) typ: Target, + + pub(crate) old_root: HashOutTarget, + pub(crate) proof_non_existence: MerkleClaimAndProofTarget, + + pub(crate) new_root: HashOutTarget, + pub(crate) new_key: ValueTarget, + pub(crate) new_value: ValueTarget, + pub(crate) new_siblings: Vec, + + // auxiliary witness + pub(crate) divergence_level: Target, +} +/// creates the targets and defines the logic of the circuit +pub fn verify_merkle_state_transition_circuit( + builder: &mut CircuitBuilder, + proof: &MerkleProofStateTransitionTarget, +) { + let measure = measure_gates_begin!( + builder, + format!("MerkleProofStateTransition_{}", proof.max_depth) + ); + let zero = builder.zero(); + + // for now, only type=0 (insertion proof) is supported + builder.connect(proof.typ, zero); + + // 1) check that for the old_root, the new_key does not exist in the tree + verify_merkle_proof_circuit(builder, &proof.proof_non_existence); + + // 2) check that for the new_root, the new_key does exist in the tree + let new_key_proof = MerkleProofExistenceTarget { + max_depth: proof.max_depth, + enabled: proof.enabled, + root: proof.new_root, + key: proof.new_key, + value: proof.new_value, + siblings: proof.new_siblings.clone(), + }; + let new_leaf_path = verify_merkle_proof_existence_circuit(builder, &new_key_proof); + + // 3.1) assert that proof_non_existence.existence==false + builder.conditional_assert_eq( + proof.enabled.target, + proof.proof_non_existence.existence.target, + zero, + ); + // 3.2) assert that proof.enabled matches with proof_non_existence.enabled + builder.connect( + proof.proof_non_existence.enabled.target, + proof.enabled.target, + ); + + // 4) assert proof_non_existence.root==old_root, and that it uses new_key + for j in 0..HASH_SIZE { + // 4.1) assert that proof.proof_non_existence.root == proof.old_root + builder.conditional_assert_eq( + proof.enabled.target, + proof.proof_non_existence.root.elements[j], + proof.old_root.elements[j], + ); + // 4.2) assert that the non-existence proof uses the new_key (value not needed for + // non-existence) + builder.conditional_assert_eq( + proof.enabled.target, + proof.proof_non_existence.key.elements[j], + proof.new_key.elements[j], + ); + } + + // prepare value for check 5.2) + let old_leaf_hash = kv_hash_target( + builder, + &proof.proof_non_existence.other_key, + &proof.proof_non_existence.other_value, + ); + // prepare values for check 5.3) + let old_leaf_path = keypath_target( + proof.max_depth, + builder, + &proof.proof_non_existence.other_key, + ); + + // 5) check that old_siblings & new_siblings match as expected. Let + // d=divergence_level, assert that: + // 5.1) old_siblings[i] == new_siblings[i] ∀ i \ {d} + // 5.2) at i==d, if old_siblings[i] != new_siblings[i]: old_siblings[i] == + // EMPTY_HASH new_siblings[i] == old_leaf_hash + // 5.3) assert that if old_key!=empty, both old_leaf_path&new_leaf_path + // should diverge at the inputted divergence level + let old_siblings = proof.proof_non_existence.siblings.clone(); + let new_siblings = proof.new_siblings.clone(); + for i in 0..proof.max_depth { + let i_targ = builder.constant(F::from_canonical_u64(i as u64)); + let is_divergence_level = builder.is_equal(i_targ, proof.divergence_level); + + // 5.1) for all i except for i==divergence_level, assert that the + // siblings are the same + let old_sibling_i: Vec = (0..HASH_SIZE) + .map(|j| builder.select(is_divergence_level, zero, old_siblings[i].elements[j])) + .collect(); + let new_sibling_i: Vec = (0..HASH_SIZE) + .map(|j| builder.select(is_divergence_level, zero, new_siblings[i].elements[j])) + .collect(); + for j in 0..HASH_SIZE { + builder.conditional_assert_eq(proof.enabled.target, old_sibling_i[j], new_sibling_i[j]); + } + + // 5.2) when i==d && if old_siblings[i] != new_siblings[i], check that: + // old_siblings[i] == EMPTY_HASH && new_siblings[i] == old_leaf_hash + + // in_case_5_2=true if: i==d (= is_divergence_level) && old_siblings[i]!=new_siblings[i] + let old_is_eq_new = zip_eq(old_siblings[i].elements, new_siblings[i].elements).fold( + builder._true(), + |acc, (old, new)| { + let eq_at_i = builder.is_equal(old, new); + builder.and(acc, eq_at_i) + }, + ); + let old_is_noteq_new = builder.not(old_is_eq_new); + let in_case_5_2 = builder.and(old_is_noteq_new, is_divergence_level); + + // do the case2's checks + let sel = builder.and(proof.enabled, in_case_5_2); + for j in 0..HASH_SIZE { + builder.conditional_assert_eq(sel.target, old_siblings[i].elements[j], zero); + builder.conditional_assert_eq( + sel.target, + new_siblings[i].elements[j], + old_leaf_hash.elements[j], + ); + } + + // 5.3) assert that if old_key!=empty, both old_leaf_path&new_leaf_path + // should diverge at the inputted divergence level. We can check it + // without having into account old_key!=empty, since if + // old_key==empty, the paths would still diverge. + let paths_eq_at_d = builder.is_equal(old_leaf_path[i].target, new_leaf_path[i].target); + builder.conditional_assert_eq( + is_divergence_level.target, + // expect them to not be equal, ie. the is_equal check to be 0 + paths_eq_at_d.target, + zero, + ); + } + + measure_gates_end!(builder, measure); +} + +impl MerkleProofStateTransitionTarget { + pub fn new_virtual(max_depth: usize, builder: &mut CircuitBuilder) -> Self { + Self { + max_depth, + enabled: builder.add_virtual_bool_target_safe(), + typ: builder.add_virtual_target(), + + old_root: builder.add_virtual_hash(), + proof_non_existence: MerkleClaimAndProofTarget::new_virtual(max_depth, builder), + new_root: builder.add_virtual_hash(), + new_key: builder.add_virtual_value(), + new_value: builder.add_virtual_value(), + + // siblings are padded till max_depth length + new_siblings: builder.add_virtual_hashes(max_depth), + + divergence_level: builder.add_virtual_target(), + } + } + + /// assigns the given values to the targets + #[allow(clippy::too_many_arguments)] + pub fn set_targets( + &self, + pw: &mut PartialWitness, + enabled: bool, + mp: &MerkleProofStateTransition, + ) -> Result<()> { + pw.set_bool_target(self.enabled, enabled)?; + pw.set_target(self.typ, F::from_canonical_u8(mp.typ))?; + + pw.set_hash_target(self.old_root, HashOut::from_vec(mp.old_root.0.to_vec()))?; + self.proof_non_existence.set_targets( + pw, + enabled, + &MerkleClaimAndProof { + root: mp.old_root, + key: mp.new_key, + value: EMPTY_VALUE, // not needed for non-existence + proof: mp.proof_non_existence.clone(), + }, + )?; + + pw.set_hash_target(self.new_root, HashOut::from_vec(mp.new_root.0.to_vec()))?; + pw.set_target_arr(&self.new_key.elements, &mp.new_key.0)?; + pw.set_target_arr(&self.new_value.elements, &mp.new_value.0)?; + + let new_siblings = mp.siblings.clone(); + + assert!(new_siblings.len() <= self.max_depth); + for (i, sibling) in new_siblings + .iter() + .chain(iter::repeat(&EMPTY_HASH)) + .take(self.max_depth) + .enumerate() + { + pw.set_hash_target(self.new_siblings[i], HashOut::from_vec(sibling.0.to_vec()))?; + } + pw.set_target( + self.divergence_level, + F::from_canonical_u64((new_siblings.len() - 1) as u64), + )?; + + Ok(()) + } +} + #[cfg(test)] pub mod tests { use std::collections::HashMap; @@ -395,7 +624,7 @@ pub mod tests { use crate::{ backends::plonky2::{ basetypes::C, - primitives::merkletree::{keypath, kv_hash, MerkleTree}, + primitives::merkletree::{keypath, kv_hash, MerkleProof, MerkleTree}, }, middleware::{hash_value, RawValue}, }; @@ -726,4 +955,240 @@ pub mod tests { Ok(()) } + + fn run_state_transition_circuit( + expect_pass: bool, + max_depth: usize, + state_transition_proof: &MerkleProofStateTransition, + ) -> Result<()> { + // sanity check, run the out-circuit proof verification + if expect_pass { + MerkleTree::verify_state_transition(max_depth, state_transition_proof)?; + } else { + // expect out-circuit verification to fail + let _ = MerkleTree::verify_state_transition(max_depth, state_transition_proof).is_err(); + } + + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + let mut pw = PartialWitness::::new(); + + let targets = MerkleProofStateTransitionTarget::new_virtual(max_depth, &mut builder); + verify_merkle_state_transition_circuit(&mut builder, &targets); + targets.set_targets(&mut pw, true, state_transition_proof)?; + + // generate & verify proof + let data = builder.build::(); + if expect_pass { + let proof = data.prove(pw)?; + data.verify(proof)?; + } else { + assert!(data.prove(pw).is_err()); // expect prove to fail + } + Ok(()) + } + + #[test] + fn test_state_transition_gadget() -> Result<()> { + let max_depth: usize = 32; + let mut kvs = HashMap::new(); + for i in 0..8 { + kvs.insert(RawValue::from(i), RawValue::from(1000 + i)); + } + let mut tree = MerkleTree::new(max_depth, &kvs)?; + + // key=37 shares path with key=5, till the level 6, needing 2 extra + // 'empty' nodes between the original position of key=5 with the new + // position of key=5 and key=37. + let old_root = tree.root(); + let key = RawValue::from(37); + let value = RawValue::from(1037); + let state_transition_proof = tree.insert(&key, &value)?; + run_state_transition_circuit(true, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + + // add a new leaf, which shares path with the previous one, but diverges + // one level before, where there is no leaf yet to be pushed down. + let old_root = tree.root(); + let key = RawValue::from(21); + let value = RawValue::from(1021); + let state_transition_proof = tree.insert(&key, &value)?; + run_state_transition_circuit(true, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + + // another leaf which will push further down the leaf with key=37 + let old_root = tree.root(); + let key = RawValue::from(101); + let value = RawValue::from(1101); + let state_transition_proof = tree.insert(&key, &value)?; + run_state_transition_circuit(true, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + + // insert two more leaves, which share almost all the path except for + // the last level (max_depth-1) + + let max_depth: usize = 32; + let mut kvs = HashMap::new(); + for i in 0..8 { + kvs.insert(RawValue::from(i), RawValue::from(1000 + i)); + } + let mut tree = MerkleTree::new(max_depth, &kvs)?; + + let old_root = tree.root(); + let key = RawValue::from(4294967295); // 0xffffffff + let value = RawValue::from(4294967295); + let state_transition_proof = tree.insert(&key, &value)?; + run_state_transition_circuit(true, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + + // insert a leaf that shares the path with the previous one, except for + // the last level (in max_depth); which would force both leaves to be + // pushed down till max_depth-1 level. + let old_root = tree.root(); + let key = RawValue::from(4026531839); // 0xefffffff + let value = RawValue::from(4026531839); + let state_transition_proof = tree.insert(&key, &value)?; + run_state_transition_circuit(true, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + + Ok(()) + } + + #[test] + fn test_state_transition_gadget_with_alteration() -> Result<()> { + let max_depth: usize = 32; + let mut kvs = HashMap::new(); + for i in 0..8 { + kvs.insert(RawValue::from(i), RawValue::from(1000 + i)); + } + let mut tree = MerkleTree::new(max_depth, &kvs)?; + + // key=37 shares path with key=5, till the level 6, needing 2 extra + // 'empty' nodes between the original position of key=5 with the new + // position of key=5 and key=37. + let old_root = tree.root(); + let key = RawValue::from(37); + let value = RawValue::from(1037); + let state_transition_proof = tree.insert(&key, &value)?; + run_state_transition_circuit(true, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + + // add a new leaf, which shares path with the previous one, but diverges + // one level before, where there is no leaf yet to be pushed down. + let old_root = tree.root(); + let key = RawValue::from(21); + let value = RawValue::from(1021); + let state_transition_proof = tree.insert(&key, &value)?; + run_state_transition_circuit(true, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + + // another leaf which will push further down the leaf with key=37 + let old_root = tree.root(); + let key = RawValue::from(101); + let value = RawValue::from(1101); + let mut state_transition_proof = tree.insert(&key, &value)?; + + // Tamper with state transition. + const OFFSET: usize = 20; + let other_leaf = state_transition_proof + .proof_non_existence + .other_leaf + .unwrap(); + let altered_proof = MerkleProof { + existence: true, + siblings: [ + state_transition_proof.proof_non_existence.siblings.clone(), + vec![EMPTY_HASH; OFFSET], + vec![kv_hash(&other_leaf.0, Some(other_leaf.1))], + ] + .concat(), + other_leaf: None, + }; + let altered_root = altered_proof.compute_root_from_leaf( + max_depth, + &state_transition_proof.new_key, + Some(state_transition_proof.new_value), + )?; + state_transition_proof.siblings = altered_proof.siblings; + state_transition_proof.new_root = altered_root; + + run_state_transition_circuit(false, max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_ne!(state_transition_proof.new_root, tree.root()); // Tamper check + Ok(()) + } + + #[test] + fn test_state_transition_gadget_disabled() -> Result<()> { + let max_depth: usize = 32; + let mut kvs = HashMap::new(); + for i in 0..8 { + kvs.insert(RawValue::from(i), RawValue::from(1000 + i)); + } + let mut tree = MerkleTree::new(max_depth, &kvs)?; + + let key = RawValue::from(37); + let value = RawValue::from(1037); + let _ = tree.insert(&key, &value)?; + + let key = RawValue::from(21); + let value = RawValue::from(1021); + let original_state_transition_proof = tree.insert(&key, &value)?; + + let mut state_transition_proof = original_state_transition_proof.clone(); + + // modify the proof, so that it should fail when `enabled=true`, by + // changing the new_root + state_transition_proof.new_root = state_transition_proof.old_root; + + run_circuit_disabled(max_depth, &state_transition_proof)?; + + // modify the proof, so that it should fail when `enabled=true`, by + // changing the new_sibling at the divergence level, which should not + // pass the verification in the case where we're inserting key=21 + let mut state_transition_proof = original_state_transition_proof.clone(); + state_transition_proof.siblings[4] = EMPTY_HASH; + + run_circuit_disabled(max_depth, &state_transition_proof)?; + + Ok(()) + } + + fn run_circuit_disabled( + max_depth: usize, + state_transition_proof: &MerkleProofStateTransition, + ) -> Result<()> { + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + let mut pw = PartialWitness::::new(); + + let targets = MerkleProofStateTransitionTarget::new_virtual(max_depth, &mut builder); + verify_merkle_state_transition_circuit(&mut builder, &targets); + targets.set_targets(&mut pw, true, state_transition_proof)?; + + // generate proof, and expect it to fail + let data = builder.build::(); + assert!(data.prove(pw).is_err()); // expect prove to fail + + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + let mut pw = PartialWitness::::new(); + + let targets = MerkleProofStateTransitionTarget::new_virtual(max_depth, &mut builder); + verify_merkle_state_transition_circuit(&mut builder, &targets); + targets.set_targets(&mut pw, false, state_transition_proof)?; + + // generate and expect it to pass + let data = builder.build::(); + let proof = data.prove(pw)?; + data.verify(proof)?; + Ok(()) + } } diff --git a/src/backends/plonky2/primitives/merkletree/error.rs b/src/backends/plonky2/primitives/merkletree/error.rs index f856883..36f3ed0 100644 --- a/src/backends/plonky2/primitives/merkletree/error.rs +++ b/src/backends/plonky2/primitives/merkletree/error.rs @@ -18,6 +18,8 @@ pub enum TreeInnerError { ProofFail(String), // inclusion / exclusion #[error("invalid {0} proof")] InvalidProof(String), + #[error("state transition proof does not verify, reason: {0}")] + StateTransitionProofFail(String), #[error("key too short (key length: {0}) for the max_depth: {1}")] TooShortKey(usize, usize), } @@ -73,6 +75,9 @@ impl TreeError { pub(crate) fn invalid_proof(obj: String) -> Self { new!(InvalidProof(obj)) } + pub(crate) fn state_transition_fail(reason: String) -> Self { + new!(StateTransitionProofFail(reason)) + } pub(crate) fn too_short_key(depth: usize, max_depth: usize) -> Self { new!(TooShortKey(depth, max_depth)) } diff --git a/src/backends/plonky2/primitives/merkletree/mod.rs b/src/backends/plonky2/primitives/merkletree/mod.rs index 0d763f9..d59e406 100644 --- a/src/backends/plonky2/primitives/merkletree/mod.rs +++ b/src/backends/plonky2/primitives/merkletree/mod.rs @@ -2,6 +2,7 @@ //! https://0xparc.github.io/pod2/merkletree.html . use std::{collections::HashMap, fmt, iter::IntoIterator}; +use itertools::zip_eq; use plonky2::field::types::Field; use serde::{Deserialize, Serialize}; @@ -55,7 +56,7 @@ impl MerkleTree { /// returns the value at the given key pub fn get(&self, key: &RawValue) -> TreeResult { let path = keypath(self.max_depth, *key)?; - let key_resolution = self.root.down(0, self.max_depth, path, None)?; + let (key_resolution, _) = self.root.down(0, self.max_depth, path, None)?; match key_resolution { Some((k, v)) if &k == key => Ok(v), _ => Err(TreeError::key_not_found()), @@ -66,7 +67,7 @@ impl MerkleTree { pub fn contains(&self, key: &RawValue) -> TreeResult { let path = keypath(self.max_depth, *key)?; match self.root.down(0, self.max_depth, path, None) { - Ok(Some((k, _))) => { + Ok((Some((k, _)), _)) => { if &k == key { Ok(true) } else { @@ -77,6 +78,34 @@ impl MerkleTree { } } + pub fn insert( + &mut self, + key: &RawValue, + value: &RawValue, + ) -> TreeResult { + let proof_non_existence = self.prove_nonexistence(key)?; + + let old_root: Hash = self.root.hash(); + self.root + .add_leaf(0, self.max_depth, Leaf::new(self.max_depth, *key, *value)?)?; + let new_root = self.root.compute_hash(); + + let (v, proof) = self.prove(key)?; + assert!(proof.existence); + assert_eq!(v, *value); + assert!(proof.other_leaf.is_none()); + + Ok(MerkleProofStateTransition { + typ: 0, // insertion + old_root, + proof_non_existence, + new_root, + new_key: *key, + new_value: *value, + siblings: proof.siblings, + }) + } + /// returns a proof of existence, which proves that the given key exists in /// the tree. It returns the `value` of the leaf at the given `key`, and the /// `MerkleProof`. @@ -89,7 +118,7 @@ impl MerkleTree { .root .down(0, self.max_depth, path, Some(&mut siblings))? { - Some((k, v)) if &k == key => Ok(( + (Some((k, v)), _) if &k == key => Ok(( v, MerkleProof { existence: true, @@ -116,20 +145,20 @@ impl MerkleTree { .down(0, self.max_depth, path, Some(&mut siblings))? { // case i) the expected leaf does not exist - None => Ok(MerkleProof { + (None, _) => Ok(MerkleProof { existence: false, siblings, other_leaf: None, }), // case ii) the expected leaf does exist in the tree, but it has a different `key` - Some((k, v)) if &k != key => Ok(MerkleProof { + (Some((k, v)), _) if &k != key => Ok(MerkleProof { existence: false, siblings, other_leaf: Some((k, v)), }), _ => Err(TreeError::key_not_found()), } - // both cases prove that the given key don't exist in the tree. ∎ + // both cases prove that the given key don't exist in the tree. } /// verifies an inclusion proof for the given `key` and `value` @@ -175,6 +204,90 @@ impl MerkleTree { } } + pub fn verify_state_transition( + max_depth: usize, + proof: &MerkleProofStateTransition, + ) -> TreeResult<()> { + let mut old_siblings = proof.proof_non_existence.siblings.clone(); + let new_siblings = proof.siblings.clone(); + + // check that for the old_root, the new_key does not exist in the tree + Self::verify_nonexistence( + max_depth, + proof.old_root, + &proof.proof_non_existence, + &proof.new_key, + )?; + + // check that new_siblings verify with the new_root + Self::verify( + max_depth, + proof.new_root, + &MerkleProof { + existence: true, + siblings: new_siblings.clone(), + other_leaf: None, + }, + &proof.new_key, + &proof.new_value, + )?; + + // if other_leaf exists, check path divergence + if let Some((other_key, _)) = proof.proof_non_existence.other_leaf { + let old_path = keypath(max_depth, other_key)?; + let new_path = keypath(max_depth, proof.new_key)?; + + let divergence_lvl: usize = match zip_eq(old_path, new_path).position(|(x, y)| x != y) { + Some(d) => d, + None => return Err(TreeError::max_depth()), + }; + + if divergence_lvl != new_siblings.len() - 1 { + return Err(TreeError::state_transition_fail( + "paths divergence does not match".to_string(), + )); + } + } + + // let d=divergence_level, assert that: + // 1) old_siblings[i] == new_siblings[i] ∀ i \ {d} + // 2) at i==d, if old_siblings[i] != new_siblings[i]: + // old_siblings[i] == EMPTY_HASH + // new_siblings[i] == old_leaf_hash + let d = new_siblings.len() - 1; + old_siblings.resize(d + 1, EMPTY_HASH); + for i in 0..d { + if old_siblings[i] != new_siblings[i] { + return Err(TreeError::state_transition_fail( + "siblings don't match: old[i]!=new[i] ∀ i (except at i==d)".to_string(), + )); + } + } + if old_siblings[d] != new_siblings[d] { + if old_siblings[d] != EMPTY_HASH { + return Err(TreeError::state_transition_fail( + "siblings don't match: old[d]!=empty".to_string(), + )); + } + let k = proof + .proof_non_existence + .other_leaf + .map(|(k, _)| k) + .ok_or(TreeError::state_transition_fail( + "proof.proof_non_existence.other_leaf can not be empty for the case old_siblings[d]!=new_siblings[d]".to_string() + ))?; + let v: Option = proof.proof_non_existence.other_leaf.map(|(_, v)| v); + let old_leaf_hash = kv_hash(&k, v); + if new_siblings[d] != old_leaf_hash { + return Err(TreeError::state_transition_fail( + "siblings don't match: new[d]!=old_leaf_hash".to_string(), + )); + } + } + + Ok(()) + } + /// returns an iterator over the leaves of the tree pub fn iter(&self) -> Iter { Iter { @@ -245,13 +358,22 @@ impl MerkleProof { max_depth: usize, key: &RawValue, value: Option, + ) -> TreeResult { + let path = keypath(max_depth, *key)?; + let h = kv_hash(key, value); + self.compute_root_from_node(max_depth, &h, path) + } + fn compute_root_from_node( + &self, + max_depth: usize, + node_hash: &Hash, + path: Vec, ) -> TreeResult { if self.siblings.len() >= max_depth { return Err(TreeError::max_depth()); } - let path = keypath(max_depth, *key)?; - let mut h = kv_hash(key, value); + let mut h = *node_hash; for (i, sibling) in self.siblings.iter().enumerate().rev() { let mut input: Vec = if path[i] { [sibling.0, h.0].concat() @@ -302,6 +424,22 @@ impl fmt::Display for MerkleClaimAndProof { } } +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct MerkleProofStateTransition { + // type: 0:insertion, 1:update, 2:deletion + pub(crate) typ: u8, + + pub(crate) old_root: Hash, + // proof of non-existence of the new_key for the old_root + pub(crate) proof_non_existence: MerkleProof, + + pub(crate) new_root: Hash, + pub(crate) new_key: RawValue, + pub(crate) new_value: RawValue, + + pub(crate) siblings: Vec, +} + #[derive(Clone, Debug)] enum Node { None, @@ -313,13 +451,33 @@ impl fmt::Display for Node { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Intermediate(n) => { - writeln!( - f, - "\"{}\" -> {{ \"{}\" \"{}\" }}", - n.hash(), - n.left.hash(), - n.right.hash() - )?; + let left_hash: String = if n.left.is_empty() { + writeln!( + f, + "\"{}_child_of_{}\" [label=\"{}\"]", + n.left.hash(), + n.hash(), + n.left.hash() + )?; + format!("\"{}_child_of_{}\"", n.left.hash(), n.hash()) + } else { + writeln!(f, "\"{}\"", n.left.hash(),)?; + format!("\"{}\"", n.left.hash()) + }; + let right_hash = if n.right.is_empty() { + writeln!( + f, + "\"{}_child_of_{}\" [label=\"{}\"]", + n.right.hash(), + n.hash(), + n.right.hash() + )?; + format!("\"{}_child_of_{}\"", n.right.hash(), n.hash()) + } else { + writeln!(f, "\"{}\"", n.right.hash(),)?; + format!("\"{}\"", n.right.hash()) + }; + writeln!(f, "\"{}\" -> {{ {} {} }}", n.hash(), left_hash, right_hash,)?; write!(f, "{}", n.left)?; write!(f, "{}", n.right) } @@ -377,7 +535,7 @@ impl Node { max_depth: usize, path: Vec, mut siblings: Option<&mut Vec>, - ) -> TreeResult> { + ) -> TreeResult<(Option<(RawValue, RawValue)>, usize)> { if lvl >= max_depth { return Err(TreeError::max_depth()); } @@ -401,8 +559,8 @@ impl Node { value, path: _p, hash: _h, - }) => Ok(Some((*key, *value))), - _ => Ok(None), + }) => Ok((Some((*key, *value)), lvl)), + _ => Ok((None, lvl)), } } @@ -619,6 +777,7 @@ pub mod tests { #[test] fn test_merkletree() -> TreeResult<()> { + let max_depth: usize = 32; let mut kvs = HashMap::new(); for i in 0..8 { if i == 1 { @@ -630,7 +789,7 @@ pub mod tests { let value = RawValue::from(1013); kvs.insert(key, value); - let tree = MerkleTree::new(32, &kvs)?; + let tree = MerkleTree::new(max_depth, &kvs)?; // when printing the tree, it should print the same tree as in // https://0xparc.github.io/pod2/merkletree.html#example-2 println!("{}", tree); @@ -640,7 +799,7 @@ pub mod tests { assert_eq!(v, RawValue::from(1013)); println!("{}", proof); - MerkleTree::verify(32, tree.root(), &proof, &key, &value)?; + MerkleTree::verify(max_depth, tree.root(), &proof, &key, &value)?; // Exclusion checks let key = RawValue::from(12); @@ -651,14 +810,14 @@ pub mod tests { ); println!("{}", proof); - MerkleTree::verify_nonexistence(32, tree.root(), &proof, &key)?; + MerkleTree::verify_nonexistence(max_depth, tree.root(), &proof, &key)?; let key = RawValue::from(1); let proof = tree.prove_nonexistence(&RawValue::from(1))?; assert_eq!(proof.other_leaf, None); println!("{}", proof); - MerkleTree::verify_nonexistence(32, tree.root(), &proof, &key)?; + MerkleTree::verify_nonexistence(max_depth, tree.root(), &proof, &key)?; // Check iterator let collected_kvs: Vec<_> = tree.into_iter().collect::>(); @@ -686,11 +845,45 @@ pub mod tests { let sorted_kvs = kvs .iter() - .sorted_by(|(k1, _), (k2, _)| cmp(32)(**k1, **k2)) + .sorted_by(|(k1, _), (k2, _)| cmp(max_depth)(**k1, **k2)) .collect::>(); assert_eq!(collected_kvs, sorted_kvs); Ok(()) } + + #[test] + fn test_state_transition() -> TreeResult<()> { + let max_depth: usize = 32; + let mut kvs = HashMap::new(); + for i in 0..8 { + kvs.insert(RawValue::from(i), RawValue::from(1000 + i)); + } + + let mut tree = MerkleTree::new(max_depth, &kvs)?; + let old_root = tree.root(); + + // key=37 shares path with key=5, till the level 6, needing 2 extra + // 'empty' nodes between the original position of key=5 with the new + // position of key=5 and key=37. + let key = RawValue::from(37); + let value = RawValue::from(1037); + let state_transition_proof = tree.insert(&key, &value)?; + + MerkleTree::verify_state_transition(max_depth, &state_transition_proof)?; + assert_eq!(state_transition_proof.old_root, old_root); + assert_eq!(state_transition_proof.new_root, tree.root()); + assert_eq!(state_transition_proof.new_key, key); + assert_eq!(state_transition_proof.new_value, value); + + // 2nd part of the test. Add a new leaf + let key = RawValue::from(21); + let value = RawValue::from(1021); + let state_transition_proof = tree.insert(&key, &value)?; + + MerkleTree::verify_state_transition(max_depth, &state_transition_proof)?; + + Ok(()) + } }