Fix out-of-bounds key nibbles in concatenation

Author: paberr

primitives/src/key_nibbles.rs

@@ -226,6 +226,16 @@ impl KeyNibbles {
             (false, false) => self.cmp(other),
         }
     }
+
+    /// Concatenates two keys, returning `None` if the combined nibble length would exceed
+    /// the storage capacity (`MAX_BYTES * 2`). Use this at trust boundaries where one of
+    /// the operands is deserialized from an untrusted source.
+    pub fn checked_add(&self, other: &KeyNibbles) -> Option<KeyNibbles> {
+        if self.len() + other.len() > Self::MAX_BYTES * 2 {
+            return None;
+        }
+        Some(self + other)
+    }
 }
 
 impl Default for KeyNibbles {
@@ -512,6 +522,25 @@ mod tests {
         assert_eq!((&key1 + &key2).to_string(), "cfb986");
     }
 
+    #[test]
+    fn checked_add_rejects_overflow() {
+        // Each operand is individually valid (length <= MAX_BYTES * 2), but their
+        // combined nibble length exceeds the 63-byte storage capacity.
+        let parent_even: KeyNibbles = "a".repeat(40).parse().unwrap();
+        let suffix_even: KeyNibbles = "b".repeat(90).parse().unwrap();
+        assert_eq!(parent_even.checked_add(&suffix_even), None);
+
+        let parent_odd: KeyNibbles = "a".repeat(41).parse().unwrap();
+        let suffix_odd: KeyNibbles = "b".repeat(87).parse().unwrap();
+        assert_eq!(parent_odd.checked_add(&suffix_odd), None);
+
+        // At the limit (sum == 126) the addition still succeeds.
+        let lhs: KeyNibbles = "a".repeat(63).parse().unwrap();
+        let rhs: KeyNibbles = "b".repeat(63).parse().unwrap();
+        let combined = lhs.checked_add(&rhs).expect("sum at the limit should fit");
+        assert_eq!(combined.len(), 126);
+    }
+
     #[test]
     fn nibbles_get_works() {
         let key: KeyNibbles = "cfb98637bcae43c13323eaa1731ced2b716962fd".parse().unwrap();

primitives/src/trie/trie_node.rs

@@ -48,9 +48,12 @@ impl TrieNodeChild {
         parent_key: &KeyNibbles,
         missing_range: &Option<RangeFrom<KeyNibbles>>,
     ) -> bool {
+        let Some(combined) = parent_key.checked_add(&self.suffix) else {
+            return false;
+        };
         missing_range
             .as_ref()
-            .map(|range| range.contains(&(parent_key + &self.suffix)))
+            .map(|range| range.contains(&combined))
             .unwrap_or(false)
     }
 
@@ -63,10 +66,13 @@ impl TrieNodeChild {
         parent_key: &KeyNibbles,
         missing_range: &Option<RangeFrom<KeyNibbles>>,
     ) -> Result<KeyNibbles, MerkleRadixTrieError> {
+        let combined = parent_key
+            .checked_add(&self.suffix)
+            .ok_or(MerkleRadixTrieError::WrongPrefix)?;
         if self.is_stump(parent_key, missing_range) {
             return Err(MerkleRadixTrieError::ChildIsStump);
         }
-        Ok(parent_key + &self.suffix)
+        Ok(combined)
     }
 }
 
@@ -723,4 +729,32 @@ mod tests {
         let root_node = TrieNode::new_root();
         assert_eq!(root_node.value, None);
     }
+
+    #[test]
+    fn key_rejects_oversized_suffix() {
+        // Both operands deserialize as valid `KeyNibbles` individually (length <= 126),
+        // but their concatenation would overflow the 63-byte buffer in `KeyNibbles::Add`.
+        // `TrieNodeChild::key` must return an error instead of panicking.
+        let cases = [
+            ("a".repeat(40), "b".repeat(90)), // even+even
+            ("a".repeat(41), "b".repeat(87)), // odd+odd
+        ];
+
+        for (parent_str, suffix_str) in &cases {
+            let parent: KeyNibbles = parent_str.parse().unwrap();
+            let suffix: KeyNibbles = suffix_str.parse().unwrap();
+            let child = TrieNodeChild {
+                suffix,
+                hash: Blake2bHash::default(),
+            };
+
+            assert_eq!(
+                child.key(&parent, &None),
+                Err(MerkleRadixTrieError::WrongPrefix),
+            );
+            // `is_stump` is panic-safe too: oversized combined keys cannot lie inside
+            // any `missing_range`, so it must return `false` without invoking `Add`.
+            assert!(!child.is_stump(&parent, &None));
+        }
+    }
 }

primitives/trie/src/trie.rs

@@ -2258,6 +2258,45 @@ mod tests {
         }
     }
 
+    #[test]
+    fn put_chunk_rejects_oversized_child_suffix() {
+        // A `TrieNodeChild.suffix` deserialized from the network used to cause an
+        // out-of-bounds slice panic inside `KeyNibbles::Add` when concatenated with the
+        // proof node's parent key. Both operands pass the `length <= 126` deserialize
+        // check individually, but together they exceed the 63-byte storage buffer.
+        // `put_chunk` must reject such a chunk gracefully instead of panicking.
+        use nimiq_primitives::trie::trie_node::TrieNodeChild;
+
+        let env = MdbxDatabase::new_volatile(Default::default()).unwrap();
+        let trie = MerkleRadixTrie::new_incomplete(&env, TestTrie);
+        let mut raw_txn = env.write_transaction();
+        let mut txn: WriteTransactionProxy = (&mut raw_txn).into();
+        assert!(!trie.is_complete(&txn));
+
+        let parent_key: KeyNibbles = "a".repeat(40).parse().unwrap();
+        let oversized_suffix: KeyNibbles = "b".repeat(90).parse().unwrap();
+
+        let mut proof_node: TrieProofNode = TrieNode::new_empty(parent_key.clone()).into();
+        // 40 + 90 = 130 nibbles > MAX_BYTES * 2 = 126 — would overflow `Add`.
+        proof_node.children[0] = Some(TrieNodeChild {
+            suffix: oversized_suffix,
+            hash: Blake2bHash::default(),
+        });
+
+        let proof_root = TrieNode::new_root();
+        let expected_hash = proof_root.hash_assert();
+        let proof = TrieProof::new(vec![proof_node, proof_root.into()], Default::default());
+
+        let malicious_chunk =
+            TrieChunk::new(None, vec![TrieItem::new(parent_key, vec![0x42])], proof);
+
+        // The exact error doesn't matter — what matters is that this returns instead
+        // of panicking inside `KeyNibbles::Add`.
+        assert!(trie
+            .put_chunk(&mut txn, KeyNibbles::ROOT, malicious_chunk, expected_hash)
+            .is_err());
+    }
+
     #[test]
     fn partial_tree_put_chunks_manual() {
         let key_1 = "413f22".parse().unwrap();