From 623839f58116c0828bc5406adbd1dd1b68e7bb3d Mon Sep 17 00:00:00 2001 From: Aaron Feickert <66188213+AaronFeickert@users.noreply.github.com> Date: Mon, 3 Jul 2023 07:00:53 -0500 Subject: [PATCH] fix: sparse Merkle tree key querying (#5566) Description --- Fixes key querying on sparse Merkle trees. Adds a test. Closes #5565. Motivation and Context --- It's possible to ask a sparse Merkle tree if it [contains](https://github.com/tari-project/tari/blob/87c070305951152c62a0179e13fadc55065cc318/base_layer/mmr/src/sparse_merkle_tree/tree.rs#L256-L259) a particular key. However, the query returns `true` whether or not the key is in the tree. This PR fixes the incorrect behavior. The query returns `true` if and only if the tree search yields a leaf node with the given key. If the tree is malformed such that the search returns an error, the query fails safe and returns `false`. Interestingly, the existing unit tests did not detect the incorrect behavior, as no test checked for the case where a key did not exist in a tree. We add a test that checks this for a few cases of interest. How Has This Been Tested? --- Existing tests pass. A new test asserts the expected behavior (and would have caught the previous incorrect behavior). What process can a PR reviewer use to test or verify this change? --- Confirm that the detection logic in the `contains` function is correct. Breaking Changes --- None. --- base_layer/mmr/src/sparse_merkle_tree/tree.rs | 43 ++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/base_layer/mmr/src/sparse_merkle_tree/tree.rs b/base_layer/mmr/src/sparse_merkle_tree/tree.rs index f5bf96dddc..bc3378ac59 100644 --- a/base_layer/mmr/src/sparse_merkle_tree/tree.rs +++ b/base_layer/mmr/src/sparse_merkle_tree/tree.rs @@ -255,7 +255,14 @@ impl> SparseMerkleTree { /// Returns true if the tree contains the key `key`. pub fn contains(&self, key: &NodeKey) -> bool { - self.search_node(key).ok().is_some() + match self.search_node(key) { + // In the case of a malformed tree where the search fails unexpectedly, play it safe + Err(_) => false, + // The node is either empty or is a leaf with an unexpected key + Ok(None) => false, + // The node is a leaf with the expected key + Ok(Some(_)) => true, + } } /// Returns the value at location `key` if it exists, or `None` otherwise. @@ -806,4 +813,38 @@ mod test { let _ = tree.delete(&short_key(65)).unwrap(); assert!(tree.is_empty()); } + + #[test] + fn contains() { + let mut tree = SparseMerkleTree::::default(); + + // An empty tree contains no keys + assert!(!tree.contains(&short_key(0))); + assert!(!tree.contains(&short_key(1))); + + // Add a key, which the tree must then contain + tree.upsert(short_key(1), ValueHash::from([1u8; 32])).unwrap(); + assert!(!tree.contains(&short_key(0))); + assert!(tree.contains(&short_key(1))); + + // Delete the key, which the tree must not contain + tree.delete(&short_key(1)).unwrap(); + assert!(!tree.contains(&short_key(0))); + assert!(!tree.contains(&short_key(1))); + + // Build a more complex tree with two keys, which the tree must then contain + tree.upsert(short_key(0), ValueHash::from([0u8; 32])).unwrap(); + tree.upsert(short_key(1), ValueHash::from([1u8; 32])).unwrap(); + assert!(tree.contains(&short_key(0))); + assert!(tree.contains(&short_key(1))); + + // Delete each key in turn + tree.delete(&short_key(0)).unwrap(); + assert!(!tree.contains(&short_key(0))); + tree.delete(&short_key(1)).unwrap(); + assert!(!tree.contains(&short_key(1))); + + // Sanity check that the tree is now empty + assert!(tree.is_empty()); + } }