diff --git a/src/merkle/tiered_smt/proof.rs b/src/merkle/tiered_smt/proof.rs index 28ac2880..1de17d52 100644 --- a/src/merkle/tiered_smt/proof.rs +++ b/src/merkle/tiered_smt/proof.rs @@ -85,18 +85,26 @@ impl TieredSmtProof { /// Note: this method cannot be used to assert non-membership. That is, if false is returned, /// it does not mean that the provided key-value pair is not in the tree. pub fn verify_membership(&self, key: &RpoDigest, value: &Word, root: &RpoDigest) -> bool { - if self.is_value_empty() { - if value != &EMPTY_VALUE { - return false; - } - // if the proof is for an empty value, we can verify it against any key which has a - // common prefix with the key storied in entries, but the prefix must be greater than - // the path length - let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0); - if common_prefix_tier < self.path.depth() { - return false; - } - } else if !self.entries.contains(&(*key, *value)) { + // Handles the following scenarios: + // - the value is set + // - empty leaf, there is an explicit entry for the key with the empty value + // - shared 64-bit prefix, the target key is not included in the entries list, the value is implicitly the empty word + let v = match self.entries.iter().find(|(k, _)| k == key) { + Some((_, v)) => v, + None => &EMPTY_VALUE, + }; + + // The value must match for the proof to be valid + if v != value { + return false; + } + + // If the proof is for an empty value, we can verify it against any key which has a common + // prefix with the key storied in entries, but the prefix must be greater than the path + // length + if self.is_value_empty() + && get_common_prefix_tier_depth(key, &self.entries[0].0) < self.path.depth() + { return false; } diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index 560db47c..788ba288 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -715,6 +715,38 @@ fn tsmt_bottom_tier_two() { // GET PROOF TESTS // ================================================================================================ +/// Tests the membership and non-membership proof for a single at depth 64 +#[test] +fn tsmt_get_proof_single_element_64() { + let mut smt = TieredSmt::default(); + + let raw_a = 0b_00000000_00000001_00000000_00000001_00000000_00000001_00000000_00000001_u64; + let key_a = [ONE, ONE, ONE, raw_a.into()].into(); + let value_a = [ONE, ONE, ONE, ONE]; + smt.insert(key_a, value_a); + + // push element `a` to depth 64, by inserting another value that shares the 48-bit prefix + let raw_b = 0b_00000000_00000001_00000000_00000001_00000000_00000001_00000000_00000000_u64; + let key_b = [ONE, ONE, ONE, raw_b.into()].into(); + smt.insert(key_b, [ONE, ONE, ONE, ONE]); + + // verify the proof for element `a` + let proof = smt.prove(key_a); + assert!(proof.verify_membership(&key_a, &value_a, &smt.root())); + + // check that a value that is not inserted in the tree produces a valid membership proof for the + // empty word + let key = [ZERO, ZERO, ZERO, ZERO].into(); + let proof = smt.prove(key); + assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root())); + + // check that a key that shared the 64-bit prefix with `a`, but is not inserted, also has a + // valid membership proof for the empty word + let key = [ONE, ONE, ZERO, raw_a.into()].into(); + let proof = smt.prove(key); + assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root())); +} + #[test] fn tsmt_get_proof() { let mut smt = TieredSmt::default();