Add expansion maximum value check to filter decode.

Signed-off-by: wuranxx <[email protected]>
valkey-io · Nov 14, 2024 · 2044cef · 2044cef
1 parent 0ce16d1
commit 2044cef
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/src/bloom/utils.rs b/src/bloom/utils.rs
@@ -1,5 +1,7 @@
 use crate::{
-    configs::{self, BLOOM_FP_RATE_MAX, BLOOM_FP_RATE_MIN},
+    configs::{
+        self, BLOOM_EXPANSION_MAX, BLOOM_EXPANSION_MIN, BLOOM_FP_RATE_MAX, BLOOM_FP_RATE_MIN,
+    },
     metrics,
 };
 use bloomfilter;
@@ -242,7 +244,7 @@ impl BloomFilterType {
                     match bincode::deserialize::<(u32, f64, Vec<BloomFilter>)>(&decoded_bytes[1..])
                     {
                         Ok(values) => {
-                            if values.0 == 0 {
+                            if !(BLOOM_EXPANSION_MIN..=BLOOM_EXPANSION_MAX).contains(&values.0) {
                                 return Err(BloomError::BadExpansion);
                             }
                             if !(values.1 > BLOOM_FP_RATE_MIN && values.1 < BLOOM_FP_RATE_MAX) {
@@ -813,7 +815,6 @@ mod tests {
         bf.expansion = 0;
 
         let encoder_result = bf.encode_bloom_filter();
-        assert!(encoder_result.is_ok());
 
         // 1. unsupport expansion
         let vec = encoder_result.unwrap();
@@ -823,6 +824,15 @@ mod tests {
             BloomFilterType::decode_bloom_filter(&vec, true).err(),
             Some(BloomError::BadExpansion)
         );
+
+        // 1.2 Exceeded the maximum expansion
+        bf.expansion = BLOOM_EXPANSION_MAX + 1;
+
+        let vec = bf.encode_bloom_filter().unwrap();
+        assert_eq!(
+            BloomFilterType::decode_bloom_filter(&vec, true).err(),
+            Some(BloomError::BadExpansion)
+        );
         // recover
         bf.expansion = origin_expansion;