From 894e20fe0cc7218d74e8a33bda3e43469ee57918 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Wed, 22 Nov 2023 17:02:22 +0100 Subject: [PATCH] simplesmt: bugfix, index must be validated before modifying the tree --- src/merkle/index.rs | 21 ++++++--- src/merkle/simple_smt/mod.rs | 14 +++--- src/merkle/simple_smt/tests.rs | 86 ++++++++++++++++++++++++++++++++++ 3 files changed, 108 insertions(+), 13 deletions(-) diff --git a/src/merkle/index.rs b/src/merkle/index.rs index 25c9282d..c533fa0d 100644 --- a/src/merkle/index.rs +++ b/src/merkle/index.rs @@ -187,13 +187,20 @@ mod tests { #[test] fn test_node_index_value_too_high() { assert_eq!(NodeIndex::new(0, 0).unwrap(), NodeIndex { depth: 0, value: 0 }); - match NodeIndex::new(0, 1) { - Err(MerkleError::InvalidIndex { depth, value }) => { - assert_eq!(depth, 0); - assert_eq!(value, 1); - } - _ => unreachable!(), - } + let err = NodeIndex::new(0, 1).unwrap_err(); + assert_eq!(err, MerkleError::InvalidIndex { depth: 0, value: 1 }); + + assert_eq!(NodeIndex::new(1, 1).unwrap(), NodeIndex { depth: 1, value: 1 }); + let err = NodeIndex::new(1, 2).unwrap_err(); + assert_eq!(err, MerkleError::InvalidIndex { depth: 1, value: 2 }); + + assert_eq!(NodeIndex::new(2, 3).unwrap(), NodeIndex { depth: 2, value: 3 }); + let err = NodeIndex::new(2, 4).unwrap_err(); + assert_eq!(err, MerkleError::InvalidIndex { depth: 2, value: 4 }); + + assert_eq!(NodeIndex::new(3, 7).unwrap(), NodeIndex { depth: 3, value: 7 }); + let err = NodeIndex::new(3, 8).unwrap_err(); + assert_eq!(err, MerkleError::InvalidIndex { depth: 3, value: 8 }); } #[test] diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 1e3241c8..39f52930 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -213,6 +213,9 @@ impl SimpleSmt { /// # Errors /// Returns an error if the index is greater than the maximum tree capacity, that is 2^{depth}. pub fn update_leaf(&mut self, index: u64, value: Word) -> Result { + // validate the index before modifying the structure + let mut idx = NodeIndex::new(self.depth(), index)?; + let old_value = self.insert_leaf_node(index, value).unwrap_or(Self::EMPTY_VALUE); // if the old value and new value are the same, there is nothing to update @@ -220,14 +223,13 @@ impl SimpleSmt { return Ok(value); } - let mut index = NodeIndex::new(self.depth(), index)?; let mut value = RpoDigest::from(value); - for _ in 0..index.depth() { - let is_right = index.is_value_odd(); - index.move_up(); - let BranchNode { left, right } = self.get_branch_node(&index); + for _ in 0..idx.depth() { + let is_right = idx.is_value_odd(); + idx.move_up(); + let BranchNode { left, right } = self.get_branch_node(&idx); let (left, right) = if is_right { (left, value) } else { (value, right) }; - self.insert_branch_node(index, left, right); + self.insert_branch_node(idx, left, right); value = Rpo256::merge(&[left, right]); } self.root = value; diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index f2c55d19..0bd818b0 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -239,6 +239,92 @@ fn with_no_duplicates_empty_node() { assert!(smt.is_ok()); } +#[test] +fn test_simplesmt_update_nonexisting_leaf_with_zero() { + // TESTING WITH EMPTY WORD + // -------------------------------------------------------------------------------------------- + + // Depth 1 has 2 leaf. Position is 0-indexed, position 2 doesn't exist. + let mut smt = SimpleSmt::new(1).unwrap(); + let result = smt.update_leaf(2, EMPTY_WORD); + assert!(!smt.leaves.contains_key(&2)); + assert!(result.is_err()); + + // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist. + let mut smt = SimpleSmt::new(2).unwrap(); + let result = smt.update_leaf(4, EMPTY_WORD); + assert!(!smt.leaves.contains_key(&4)); + assert!(result.is_err()); + + // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist. + let mut smt = SimpleSmt::new(3).unwrap(); + let result = smt.update_leaf(8, EMPTY_WORD); + assert!(!smt.leaves.contains_key(&8)); + assert!(result.is_err()); + + // TESTING WITH A VALUE + // -------------------------------------------------------------------------------------------- + let value = int_to_node(1); + + // Depth 1 has 2 leaves. Position is 0-indexed, position 1 doesn't exist. + let mut smt = SimpleSmt::new(1).unwrap(); + let result = smt.update_leaf(2, *value); + assert!(!smt.leaves.contains_key(&2)); + assert!(result.is_err()); + + // Depth 2 has 4 leaves. Position is 0-indexed, position 2 doesn't exist. + let mut smt = SimpleSmt::new(2).unwrap(); + let result = smt.update_leaf(4, *value); + assert!(!smt.leaves.contains_key(&4)); + assert!(result.is_err()); + + // Depth 3 has 8 leaves. Position is 0-indexed, position 4 doesn't exist. + let mut smt = SimpleSmt::new(3).unwrap(); + let result = smt.update_leaf(8, *value); + assert!(!smt.leaves.contains_key(&8)); + assert!(result.is_err()); +} + +#[test] +fn test_simplesmt_with_leaves_nonexisting_leaf() { + // TESTING WITH EMPTY WORD + // -------------------------------------------------------------------------------------------- + + // Depth 1 has 2 leaf. Position is 0-indexed, position 2 doesn't exist. + let leaves = [(2, EMPTY_WORD)]; + let result = SimpleSmt::with_leaves(1, leaves); + assert!(result.is_err()); + + // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist. + let leaves = [(4, EMPTY_WORD)]; + let result = SimpleSmt::with_leaves(2, leaves); + assert!(result.is_err()); + + // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist. + let leaves = [(8, EMPTY_WORD)]; + let result = SimpleSmt::with_leaves(3, leaves); + assert!(result.is_err()); + + // TESTING WITH A VALUE + // -------------------------------------------------------------------------------------------- + let value = int_to_node(1); + + // Depth 1 has 2 leaves. Position is 0-indexed, position 2 doesn't exist. + let leaves = [(2, *value)]; + let result = SimpleSmt::with_leaves(1, leaves); + assert!(result.is_err()); + + // Depth 2 has 4 leaves. Position is 0-indexed, position 4 doesn't exist. + let leaves = [(4, *value)]; + let result = SimpleSmt::with_leaves(2, leaves); + assert!(result.is_err()); + + // Depth 3 has 8 leaves. Position is 0-indexed, position 8 doesn't exist. + let leaves = [(8, *value)]; + let result = SimpleSmt::with_leaves(3, leaves); + assert!(result.is_err()); +} + // HELPER FUNCTIONS // --------------------------------------------------------------------------------------------