Fix OOB heap read crash with malicious snapshot (#9268)

* Validate quantized u8 data size on load

* Validate other quantization types

* Add test

* Use fs_err

lib/quantization/src/encoded_storage.rs

@@ -62,6 +62,39 @@ pub trait EncodedStorageBuilder {
     fn push_vector_data(&mut self, other: &[u8]) -> Result<(), Self::Error>;
 }
 
+/// Validate that every encoded vector in `storage` has exactly `expected_size` bytes — the
+/// per-vector size the quantizer derives from its metadata.
+///
+/// The scoring hot paths assume each stored vector has this exact size: the storage stride and
+/// the quantizer metadata are both derived from the same vector parameters, so on consistent data
+/// they always match. Verifying it once here at load time keeps that invariant guaranteed in
+/// release builds, without paying for a bounds check on every score.
+///
+/// The storage uses a fixed stride for every vector, so inspecting the first encoded vector is
+/// enough to validate all of them. An empty storage has no vector data to score, so there is
+/// nothing to check.
+pub(crate) fn validate_storage_vector_size(
+    storage: &impl EncodedStorage,
+    expected_size: usize,
+) -> std::io::Result<()> {
+    if storage.vectors_count() == 0 {
+        return Ok(());
+    }
+
+    let actual_size = storage.get_vector_data(0).len();
+    if actual_size != expected_size {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            format!(
+                "Quantized vector storage is inconsistent with its metadata: encoded vector size \
+                 is {actual_size} bytes, but metadata expects {expected_size} bytes",
+            ),
+        ));
+    }
+
+    Ok(())
+}
+
 #[cfg(feature = "testing")]
 pub struct TestEncodedStorage {
     data: Vec<u8>,
@@ -232,3 +265,40 @@ impl EncodedStorageBuilder for TestEncodedStorageBuilder {
         Ok(())
     }
 }
+
+#[cfg(all(test, feature = "testing"))]
+mod tests {
+    use super::*;
+
+    fn storage_with_stride(stride: usize, count: usize) -> TestEncodedStorage {
+        let mut builder = TestEncodedStorageBuilder::new(None, stride);
+        let vector = vec![0u8; stride];
+        for _ in 0..count {
+            builder.push_vector_data(&vector).unwrap();
+        }
+        builder.build().unwrap()
+    }
+
+    #[test]
+    fn accepts_matching_size() {
+        let storage = storage_with_stride(260, 4);
+        validate_storage_vector_size(&storage, 260).unwrap();
+    }
+
+    #[test]
+    fn rejects_mismatched_size() {
+        let storage = storage_with_stride(260, 4);
+        // Both a smaller and a larger expected size must be rejected (exact match required).
+        for expected_size in [130, 520] {
+            let err = validate_storage_vector_size(&storage, expected_size).unwrap_err();
+            assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        }
+    }
+
+    #[test]
+    fn skips_empty_storage() {
+        let storage = storage_with_stride(260, 0);
+        // With no stored vectors there is nothing to check, so any size is accepted.
+        validate_storage_vector_size(&storage, 999).unwrap();
+    }
+}

lib/quantization/src/encoded_vectors_binary.rs

@@ -15,6 +15,7 @@ use fs_err as fs;
 use serde::{Deserialize, Serialize};
 use strum::EnumIter;
 
+use crate::encoded_storage::validate_storage_vector_size;
 use crate::encoded_vectors::validate_vector_parameters;
 use crate::vector_stats::{VectorElementStats, VectorStats};
 use crate::{
@@ -514,6 +515,12 @@ impl<TBitsStoreType: BitsStoreType, TStorage: EncodedStorage>
             encoded_vectors,
             bits_store_type: PhantomData,
         };
+
+        // Validate the storage's vector size against the metadata once here, so the size
+        // invariant the scoring hot path relies on (it XORs the stored vector against an
+        // equally-sized query) also holds in release builds without a per-score check.
+        validate_storage_vector_size(&result.encoded_vectors, result.quantized_vector_size())?;
+
         Ok(result)
     }
 

lib/quantization/src/encoded_vectors_pq.rs

@@ -19,7 +19,7 @@ use fs_err as fs;
 use parking_lot::Mutex;
 use serde::{Deserialize, Serialize};
 
-use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder};
+use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder, validate_storage_vector_size};
 use crate::encoded_vectors::{EncodedVectors, VectorParameters, validate_vector_parameters};
 use crate::kmeans::kmeans;
 use crate::{ConditionalVariable, EncodingError};
@@ -148,6 +148,12 @@ impl<TStorage: EncodedStorage> EncodedVectorsPQ<TStorage> {
             metadata,
             metadata_path: Some(meta_path.to_path_buf()),
         };
+
+        // Validate the storage's vector size against the metadata once here, so the size
+        // invariant the scoring hot path relies on (it walks the query LUT one entry per stored
+        // byte) also holds in release builds without a per-score check.
+        validate_storage_vector_size(&result.encoded_vectors, result.quantized_vector_size())?;
+
         Ok(result)
     }
 

lib/quantization/src/encoded_vectors_tq.rs

@@ -12,7 +12,7 @@ use fs_err as fs;
 use serde::{Deserialize, Serialize};
 
 use crate::EncodingError;
-use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder};
+use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder, validate_storage_vector_size};
 use crate::encoded_vectors::{EncodedVectors, VectorParameters, validate_vector_parameters};
 use crate::quantile::find_quantile_interval_per_coordinate_with_preprocess;
 use crate::turboquant::math::std_normal_cdf;
@@ -289,6 +289,12 @@ impl<TStorage: EncodedStorage> EncodedVectorsTQ<TStorage> {
             quantizer,
         };
 
+        // Validate the storage's vector size against the metadata once here, so the size
+        // invariant the scoring hot path relies on (it splits each stored vector into packed
+        // dimensions plus a fixed-size trailer) also holds in release builds without a per-score
+        // check.
+        validate_storage_vector_size(&result.encoded_vectors, result.quantized_vector_size())?;
+
         Ok(result)
     }
 

lib/quantization/src/encoded_vectors_u8.rs

@@ -12,7 +12,7 @@ use fs_err as fs;
 use serde::{Deserialize, Serialize};
 
 use crate::EncodingError;
-use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder};
+use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder, validate_storage_vector_size};
 use crate::encoded_vectors::{
     DistanceType, EncodedVectors, VectorParameters, validate_vector_parameters,
 };
@@ -320,6 +320,12 @@ impl<TStorage: EncodedStorage> EncodedVectorsU8<TStorage> {
             metadata,
             metadata_path: Some(meta_path.to_path_buf()),
         };
+
+        // Validate the storage's vector size against the metadata once here, so the size
+        // invariant the hot path in `score_bytes` relies on (it reads `actual_dim` bytes past a
+        // fixed offset) also holds in release builds without a per-score check.
+        validate_storage_vector_size(&result.encoded_vectors, result.quantized_vector_size())?;
+
         Ok(result)
     }
 
@@ -747,6 +753,9 @@ impl<TStorage: EncodedStorage> EncodedVectors for EncodedVectorsU8<TStorage> {
             .cpu_counter()
             .incr_delta(self.metadata.vector_parameters().dim);
 
+        // For storage-derived vectors this invariant is validated once at load time (see
+        // `load`), so we only assert it here in debug builds to also cover externally provided
+        // byte slices without adding overhead to this hot path in release builds.
         debug_assert!(bytes.len() >= ADDITIONAL_CONSTANT_SIZE + self.metadata.actual_dim());
 
         #[cfg(target_arch = "x86_64")]

lib/quantization/tests/integration/load_validation.rs

@@ -0,0 +1,158 @@
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::AtomicBool;
+
+    use quantization::encoded_storage::{TestEncodedStorage, TestEncodedStorageBuilder};
+    use quantization::encoded_vectors::{DistanceType, VectorParameters};
+    use quantization::encoded_vectors_u8;
+    use quantization::encoded_vectors_u8::{EncodedVectorsU8, ScalarQuantizationMethod};
+    use tempfile::Builder;
+
+    /// The hot-path `score_bytes` asserts (in debug builds) that each encoded vector is at least
+    /// as large as the metadata expects. `load` performs a stricter, exact-size check for
+    /// storage-derived vectors once at load time, so the invariant also holds in release builds.
+    /// This verifies that a storage whose vector stride does not exactly match the metadata is
+    /// rejected at load time instead of reaching the hot path and reading out of bounds.
+    #[test]
+    fn load_requires_storage_vector_size_to_match_metadata() {
+        let dir = Builder::new().prefix("storage_dir").tempdir().unwrap();
+
+        let vectors_count = 4;
+        let vector_dim = 256;
+        let vector_parameters = VectorParameters {
+            dim: vector_dim,
+            deprecated_count: None,
+            distance_type: DistanceType::Dot,
+            invert: false,
+        };
+        let vector_data: Vec<Vec<f32>> = (0..vectors_count)
+            .map(|i| vec![i as f32; vector_dim])
+            .collect();
+
+        let data_path = dir.path().join("data.bin");
+        let meta_path = dir.path().join("meta.json");
+        let quantized_vector_size =
+            encoded_vectors_u8::get_quantized_vector_size(&vector_parameters);
+        EncodedVectorsU8::encode(
+            vector_data.iter(),
+            TestEncodedStorageBuilder::new(Some(data_path.as_path()), quantized_vector_size),
+            &vector_parameters,
+            vectors_count,
+            None,
+            ScalarQuantizationMethod::Int8,
+            Some(meta_path.as_path()),
+            &AtomicBool::new(false),
+        )
+        .unwrap();
+
+        let try_load = |stride: usize| {
+            EncodedVectorsU8::<TestEncodedStorage>::load(
+                TestEncodedStorage::from_file(data_path.as_path(), stride).unwrap(),
+                meta_path.as_path(),
+            )
+        };
+
+        // Loading with the exact stride succeeds.
+        try_load(quantized_vector_size).unwrap();
+
+        // Loading with a stride that does not exactly match the metadata must be rejected,
+        // whether it is smaller or larger than expected. Both strides still divide the data
+        // file, so they pass the storage's own divisibility check and only our load-time
+        // consistency check catches the mismatch.
+        for stride in [quantized_vector_size / 2, quantized_vector_size * 2] {
+            match try_load(stride) {
+                Ok(_) => panic!(
+                    "loading a storage with vector stride {stride} should fail, \
+                     metadata expects {quantized_vector_size}"
+                ),
+                Err(err) => assert_eq!(err.kind(), std::io::ErrorKind::InvalidData),
+            }
+        }
+    }
+
+    /// Regression test for BBP-827 (vulnerability 1): heap out-of-bounds read via crafted
+    /// quantization metadata in a malicious snapshot.
+    ///
+    /// A snapshot's `quantized.meta.json` is attacker-controlled. The scalar-quantization SIMD
+    /// scoring functions use `metadata.actual_dim` as the iteration count for raw pointer
+    /// arithmetic over each stored vector (e.g. `impl_score_dot(q_ptr, v_ptr, actual_dim)`). If
+    /// `actual_dim` is inflated far beyond the real vector size, every score reads heap memory
+    /// past the end of the (small) stored vector — leaking adjacent heap data into scores, or
+    /// crashing outright for large enough values.
+    ///
+    /// This crafts that malicious snapshot: a storage honestly encoded for dimension 128 whose
+    /// `quantized.meta.json` is then patched to claim a far larger `actual_dim`, exactly as an
+    /// attacker would inside the snapshot archive. It asserts that `load` rejects the inconsistent
+    /// storage, so search queries can never reach the out-of-bounds read.
+    #[test]
+    fn bbp_827_load_rejects_inflated_actual_dim() {
+        let dir = Builder::new().prefix("storage_dir").tempdir().unwrap();
+
+        let vectors_count = 4;
+        let vector_dim = 128;
+        let vector_parameters = VectorParameters {
+            dim: vector_dim,
+            deprecated_count: None,
+            distance_type: DistanceType::Cosine,
+            invert: false,
+        };
+        let vector_data: Vec<Vec<f32>> = (0..vectors_count)
+            .map(|i| vec![i as f32; vector_dim])
+            .collect();
+
+        let data_path = dir.path().join("data.bin");
+        let meta_path = dir.path().join("quantized.meta.json");
+        let quantized_vector_size =
+            encoded_vectors_u8::get_quantized_vector_size(&vector_parameters);
+        EncodedVectorsU8::encode(
+            vector_data.iter(),
+            TestEncodedStorageBuilder::new(Some(data_path.as_path()), quantized_vector_size),
+            &vector_parameters,
+            vectors_count,
+            None,
+            ScalarQuantizationMethod::Int8,
+            Some(meta_path.as_path()),
+            &AtomicBool::new(false),
+        )
+        .unwrap();
+
+        let load = || {
+            EncodedVectorsU8::<TestEncodedStorage>::load(
+                TestEncodedStorage::from_file(data_path.as_path(), quantized_vector_size).unwrap(),
+                meta_path.as_path(),
+            )
+        };
+
+        // The honest snapshot loads fine.
+        load().unwrap();
+
+        // Sanity check: the encoded metadata records the real dimension under `actual_dim`. This
+        // guards against the field being renamed, which would silently make the patch below (and
+        // thus this regression test) a no-op.
+        let original = fs_err::read_to_string(&meta_path).unwrap();
+        let metadata: serde_json::Value = serde_json::from_str(&original).unwrap();
+        assert_eq!(metadata["actual_dim"].as_u64(), Some(vector_dim as u64));
+
+        let patch_actual_dim = |actual_dim: u64| {
+            let contents = fs_err::read_to_string(&meta_path).unwrap();
+            let mut metadata: serde_json::Value = serde_json::from_str(&contents).unwrap();
+            metadata["actual_dim"] = serde_json::Value::from(actual_dim);
+            fs_err::write(&meta_path, serde_json::to_string(&metadata).unwrap()).unwrap();
+        };
+
+        // Each stored vector is only `quantized_vector_size` bytes, but the patched metadata
+        // claims `actual_dim` per vector — so scoring would iterate `actual_dim` times over that
+        // small buffer, reading well past its allocation. `8192` is the report's score-inflation
+        // value; `10_000_000` is its crash value. Both must be rejected at load time.
+        for inflated_actual_dim in [8192, 10_000_000] {
+            patch_actual_dim(inflated_actual_dim);
+            match load() {
+                Ok(_) => panic!(
+                    "malicious snapshot with actual_dim={inflated_actual_dim} was accepted — \
+                     heap OOB read is reachable (BBP-827 regression)"
+                ),
+                Err(err) => assert_eq!(err.kind(), std::io::ErrorKind::InvalidData),
+            }
+        }
+    }
+}

lib/quantization/tests/integration/main.rs

@@ -1,6 +1,8 @@
 #[cfg(test)]
 pub mod empty_storage;
 #[cfg(test)]
+pub mod load_validation;
+#[cfg(test)]
 pub mod metrics;
 #[cfg(test)]
 pub mod stop_condition;