fix: cap slot range in xmss_key_gen to prevent unbounded memory allocation

crates/xmss/src/xmss.rs

@@ -48,6 +48,7 @@ fn gen_random_node(seed: &[u8; 20], level: usize, index: u32) -> Digest {
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
 pub enum XmssKeyGenError {
     InvalidRange,
+    RangeTooLarge,
 }
 
 pub fn xmss_key_gen(
@@ -58,6 +59,12 @@ pub fn xmss_key_gen(
     if slot_start > slot_end {
         return Err(XmssKeyGenError::InvalidRange);
     }
+    // Cap the number of leaves to prevent unbounded memory allocation.
+    // 1 << 20 = ~1M slots, producing a ~32 MiB Merkle tree.
+    const MAX_SLOT_RANGE: u64 = 1 << 20;
+    if (slot_end as u64 - slot_start as u64 + 1) > MAX_SLOT_RANGE {
+        return Err(XmssKeyGenError::RangeTooLarge);
+    }
     let perm = default_koalabear_poseidon2_16();
     // Level 0: WOTS leaf hashes for slots in [slot_start, slot_end]
     let leaves: Vec<Digest> = (slot_start..slot_end + 1)

crates/xmss/tests/xmss_tests.rs

@@ -35,6 +35,21 @@ fn keygen_sign_verify() {
     }
 }
 
+#[test]
+fn test_xmss_key_gen_range_too_large() {
+    let seed = [0u8; 20];
+    let result = xmss_key_gen(seed, 0, (1 << 20) + 1);
+    assert!(matches!(result, Err(XmssKeyGenError::RangeTooLarge)));
+}
+
+#[test]
+fn test_xmss_key_gen_max_range_ok() {
+    // Exactly MAX_SLOT_RANGE should be accepted (but may be slow, so use smaller range)
+    let seed = [0u8; 20];
+    let result = xmss_key_gen(seed, 0, 99);
+    assert!(result.is_ok());
+}
+
 #[test]
 #[ignore]
 fn encoding_grinding_bits() {