Key length is checked when updating the SMT (#130)

🐛 Key length is checked when updating the SMT

src/sparse_merkle_tree/smt.rs

@@ -1483,6 +1483,12 @@ impl SparseMerkleTree {
             return Ok(Arc::clone(&self.root));
         }
         let (update_keys, update_values) = data.entries();
+        // check if all keys have the same length
+        if !utils::have_all_arrays_same_length(&update_keys, self.key_length.into()) {
+            return Err(SMTError::InvalidInput(String::from(
+                "all keys must have the same length",
+            )));
+        }
         // get the root subtree
         let root = self.get_subtree(db, &self.root.lock().unwrap())?;
         // update using the key-value pairs starting from the root (height: 0).
@@ -1608,6 +1614,52 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_key_length_invalid_size() {
+        let test_data = vec![
+            (
+                vec![
+                    "4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a",
+                    "6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa0", // invalid key size
+                ],
+                vec![
+                    "9c12cfdc04c74584d787ac3d23772132c18524bc7ab28dec4219b8fc5b425f70",
+                    "1406e05881e299367766d313e26c05564ec91bf721d31726bd6e46e60689539a",
+                ],
+            ),
+            (
+                vec![
+                    "4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a",
+                    "6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa07b3c", // invalid key size
+                ],
+                vec![
+                    "9c12cfdc04c74584d787ac3d23772132c18524bc7ab28dec4219b8fc5b425f70",
+                    "1406e05881e299367766d313e26c05564ec91bf721d31726bd6e46e60689539a",
+                ],
+            ),
+        ];
+
+        for (keys, values) in test_data {
+            let mut tree = SparseMerkleTree::new(&[], KeyLength(32), Default::default());
+            let mut data = UpdateData { data: Cache::new() };
+            for idx in 0..keys.len() {
+                data.data.insert(
+                    hex::decode(keys[idx]).unwrap(),
+                    hex::decode(values[idx]).unwrap(),
+                );
+            }
+            let mut db = smt_db::InMemorySmtDB::default();
+            let result = tree.commit(&mut db, &data);
+
+            assert_eq!(
+                result.err(),
+                Some(SMTError::InvalidInput(String::from(
+                    "all keys must have the same length"
+                )))
+            );
+        }
+    }
+
     #[test]
     fn test_small_tree_0() {
         let test_data = vec![(

src/utils.rs

@@ -120,6 +120,15 @@ pub fn array_equal_bool(a: &[bool], b: &[bool]) -> bool {
     true
 }
 
+pub fn have_all_arrays_same_length<T>(list: &[&[T]], length: usize) -> bool {
+    for v in list {
+        if v.len() != length {
+            return false;
+        }
+    }
+    true
+}
+
 pub fn binary_search<T>(list: &[T], callback: impl Fn(&T) -> bool) -> i32 {
     let mut lo = -1;
     let mut hi = list.len() as i32;
@@ -366,6 +375,28 @@ mod tests {
         assert!(!array_equal_bool(&[false, false], &[false, false, false]));
     }
 
+    #[test]
+    fn test_have_all_arrays_same_length() {
+        // arrays have the same length
+        assert!(have_all_arrays_same_length(&[&[1, 2, 3], &[1, 2, 3]], 3));
+        assert!(have_all_arrays_same_length(&[&[1, 2], &[3, 4]], 2));
+        assert!(have_all_arrays_same_length(
+            &[&[1, 2, 3], &[1, 2, 3], &[1, 2, 3]],
+            3
+        ));
+        assert!(have_all_arrays_same_length(&[&[1], &[2], &[3], &[4]], 1));
+
+        // arrays have different length
+        assert!(!have_all_arrays_same_length(&[&[1, 2, 3], &[1, 2, 3]], 2));
+        assert!(!have_all_arrays_same_length(&[&[1, 2], &[3, 4]], 3));
+        assert!(!have_all_arrays_same_length(
+            &[&[1, 2, 3], &[1, 2, 3], &[1, 2, 3]],
+            2
+        ));
+        assert!(!have_all_arrays_same_length(&[&[1, 2, 3], &[4, 5]], 2));
+        assert!(!have_all_arrays_same_length(&[&[1, 2, 3], &[4, 5]], 3));
+    }
+
     #[test]
     fn test_binary_search() {
         let test_data = vec![