Merge pull request #176 from 0xPolygonMiden/hacka-tsmt-error-codes

tsmt: return error code instead of panic

src/merkle/tiered_smt/error.rs

@@ -0,0 +1,49 @@
+use core::fmt::Display;
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum TieredSmtProofError {
+    EntriesEmpty,
+    PathTooLong,
+    NotATierPath(u8),
+    MultipleEntriesOutsideLastTier,
+    EmptyValueNotAllowed,
+    UnmatchingPrefixes(u64, u64),
+}
+
+impl Display for TieredSmtProofError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            TieredSmtProofError::EntriesEmpty => {
+                write!(f, "Missing entries for tiered sparse merkle tree proof")
+            }
+            TieredSmtProofError::PathTooLong => {
+                write!(
+                    f,
+                    "Path longer than maximum depth of 64 for tiered sparse merkle tree proof"
+                )
+            }
+            TieredSmtProofError::NotATierPath(got) => {
+                write!(
+                    f,
+                    "Path length does not correspond to a tier. Got {} Expected one of 16,32,48,64",
+                    got
+                )
+            }
+            TieredSmtProofError::MultipleEntriesOutsideLastTier => {
+                write!(f, "Multiple entries are only allowed for the last tier (depth 64)")
+            }
+            TieredSmtProofError::EmptyValueNotAllowed => {
+                write!(
+                    f,
+                    "The empty value [0,0,0,0] is not allowed inside a tiered sparse merkle tree"
+                )
+            }
+            TieredSmtProofError::UnmatchingPrefixes(first, second) => {
+                write!(f, "Not all leaves have the same prefix. First {} second {}", first, second)
+            }
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for TieredSmtProofError {}

src/merkle/tiered_smt/mod.rs

@@ -14,6 +14,9 @@ use values::ValueStore;
 mod proof;
 pub use proof::TieredSmtProof;
 
+mod error;
+pub use error::TieredSmtProofError;
+
 #[cfg(test)]
 mod tests;
 
@@ -159,7 +162,7 @@ impl TieredSmt {
             vec![(key, Self::EMPTY_VALUE)]
         };
 
-        TieredSmtProof::new(path, entries)
+        TieredSmtProof::new(path, entries).expect("Bug detected, TSMT produced invalid proof")
     }
 
     // STATE MUTATORS

src/merkle/tiered_smt/proof.rs

@@ -1,6 +1,6 @@
 use super::{
     get_common_prefix_tier_depth, get_key_prefix, hash_bottom_leaf, hash_upper_leaf,
-    EmptySubtreeRoots, LeafNodeIndex, MerklePath, RpoDigest, Vec, Word,
+    EmptySubtreeRoots, LeafNodeIndex, MerklePath, RpoDigest, TieredSmtProofError, Vec, Word,
 };
 
 // CONSTANTS
@@ -12,6 +12,9 @@ const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH;
 /// Value of an empty leaf.
 pub const EMPTY_VALUE: Word = super::TieredSmt::EMPTY_VALUE;
 
+/// Depths at which leaves can exist in a tiered SMT.
+pub const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS;
+
 // TIERED SPARSE MERKLE TREE PROOF
 // ================================================================================================
 
@@ -39,23 +42,38 @@ impl TieredSmtProof {
     /// - Entries contains more than 1 item, but the length of the path is not 64.
     /// - Entries contains more than 1 item, and one of the items has value set to [ZERO; 4].
     /// - Entries contains multiple items with keys which don't share the same 64-bit prefix.
-    pub fn new<I>(path: MerklePath, entries: I) -> Self
+    pub fn new<I>(path: MerklePath, entries: I) -> Result<Self, TieredSmtProofError>
     where
         I: IntoIterator<Item = (RpoDigest, Word)>,
     {
         let entries: Vec<(RpoDigest, Word)> = entries.into_iter().collect();
-        assert!(path.depth() <= MAX_DEPTH);
+
-        assert!(!entries.is_empty());
+        if !TIER_DEPTHS.into_iter().any(|e| e == path.depth()) {
+            return Err(TieredSmtProofError::NotATierPath(path.depth()));
+        }
+
+        if entries.is_empty() {
+            return Err(TieredSmtProofError::EntriesEmpty);
+        }
+
         if entries.len() > 1 {
-            assert!(path.depth() == MAX_DEPTH);
+            if path.depth() != MAX_DEPTH {
+                return Err(TieredSmtProofError::MultipleEntriesOutsideLastTier);
+            }
+
             let prefix = get_key_prefix(&entries[0].0);
             for entry in entries.iter().skip(1) {
-                assert_ne!(entry.1, EMPTY_VALUE);
+                if entry.1 == EMPTY_VALUE {
-                assert_eq!(prefix, get_key_prefix(&entry.0));
+                    return Err(TieredSmtProofError::EmptyValueNotAllowed);
+                }
+                let current = get_key_prefix(&entry.0);
+                if prefix != current {
+                    return Err(TieredSmtProofError::UnmatchingPrefixes(prefix, current));
+                }
             }
         }
 
-        Self { path, entries }
+        Ok(Self { path, entries })
     }
 
     // PROOF VERIFIER