simplesmt: bugfix, index must be validated before modifying the tree

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]

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<Word, MerkleError> {
+        // 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;

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
 // --------------------------------------------------------------------------------------------