diff --git a/src/merkle/smt/full/error.rs b/src/merkle/smt/full/error.rs index e1c3fad0..b6290432 100644 --- a/src/merkle/smt/full/error.rs +++ b/src/merkle/smt/full/error.rs @@ -12,16 +12,20 @@ use crate::{ #[derive(Clone, Debug, PartialEq, Eq)] pub enum SmtLeafError { - InconsistentKeysInEntries { + InconsistentKeys { entries: Vec<(RpoDigest, Word)>, key_1: RpoDigest, key_2: RpoDigest, }, InvalidNumEntriesForMultiple(usize), - KeyInconsistentWithLeafIndex { + SingleKeyInconsistentWithLeafIndex { key: RpoDigest, leaf_index: LeafIndex, }, + MultipleKeysInconsistentWithLeafIndex { + leaf_index_from_keys: LeafIndex, + leaf_index_supplied: LeafIndex, + }, } #[cfg(feature = "std")] @@ -34,16 +38,27 @@ impl fmt::Display for SmtLeafError { InvalidNumEntriesForMultiple(num_entries) => { write!(f, "Multiple leaf requires 2 or more entries. Got: {num_entries}") } - InconsistentKeysInEntries { entries, key_1, key_2 } => { + InconsistentKeys { entries, key_1, key_2 } => { write!(f, "Multiple leaf requires all keys to map to the same leaf index. Offending keys: {key_1} and {key_2}. Entries: {entries:?}.") } - KeyInconsistentWithLeafIndex { key, leaf_index } => { + SingleKeyInconsistentWithLeafIndex { key, leaf_index } => { write!( f, "Single key in leaf inconsistent with leaf index. Key: {key}, leaf index: {}", leaf_index.value() ) } + MultipleKeysInconsistentWithLeafIndex { + leaf_index_from_keys, + leaf_index_supplied, + } => { + write!( + f, + "Keys in entries map to leaf index {}, but leaf index {} was supplied", + leaf_index_from_keys.value(), + leaf_index_supplied.value() + ) + } } } } diff --git a/src/merkle/smt/full/leaf.rs b/src/merkle/smt/full/leaf.rs index 2c2aaa7d..3387fc6d 100644 --- a/src/merkle/smt/full/leaf.rs +++ b/src/merkle/smt/full/leaf.rs @@ -22,6 +22,8 @@ impl SmtLeaf { /// /// # Errors /// - Returns an error if 2 keys in `entries` map to a different leaf index + /// - Returns an error if 1 or more keys in `entries` map to a leaf index + /// different from `leaf_index` pub fn new( entries: Vec<(RpoDigest, Word)>, leaf_index: LeafIndex, @@ -32,12 +34,28 @@ impl SmtLeaf { let (key, value) = entries[0]; if LeafIndex::::from(key) != leaf_index { - return Err(SmtLeafError::KeyInconsistentWithLeafIndex { key, leaf_index }); + return Err(SmtLeafError::SingleKeyInconsistentWithLeafIndex { + key, + leaf_index, + }); } Ok(Self::new_single(key, value)) } - _ => Self::new_multiple(entries), + _ => { + let leaf = Self::new_multiple(entries)?; + + // `new_multiple()` checked that all keys map to the same leaf index. We still need + // to ensure that that leaf index is `leaf_index`. + if leaf.index() != leaf_index { + Err(SmtLeafError::MultipleKeysInconsistentWithLeafIndex { + leaf_index_from_keys: leaf.index(), + leaf_index_supplied: leaf_index, + }) + } else { + Ok(leaf) + } + } } } @@ -73,7 +91,7 @@ impl SmtLeaf { let next_leaf_index: LeafIndex = next_key.into(); if next_leaf_index != first_leaf_index { - return Err(SmtLeafError::InconsistentKeysInEntries { + return Err(SmtLeafError::InconsistentKeys { entries, key_1: first_key, key_2: next_key,