You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datasketches.apache.org by ti...@apache.org on 2026/07/29 06:55:47 UTC
(datasketches-rust) 02/02: fix(tdigest): reject invalid serialized state
This is an automated email from the ASF dual-hosted git repository.
tisonkun pushed a commit to branch codex/tdigest-p0-fixes
in repository https://gitbox.apache.org/repos/asf/datasketches-rust.git
commit 37865efbd3a349328acaf144e92bc76c7baabecf
Author: tison <wa...@gmail.com>
AuthorDate: Wed Jul 29 14:54:35 2026 +0800
fix(tdigest): reject invalid serialized state
Signed-off-by: tison <wa...@gmail.com>
---
datasketches/src/tdigest/sketch.rs | 40 ++++++++++++++---
datasketches/tests/serde_tests/tdigest.rs | 75 +++++++++++++++++++++++++++++++
2 files changed, 108 insertions(+), 7 deletions(-)
diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs
index 7b39c3a..136d145 100644
--- a/datasketches/src/tdigest/sketch.rs
+++ b/datasketches/src/tdigest/sketch.rs
@@ -577,6 +577,8 @@ impl TDigestMut {
};
check_non_nan(min, "min")?;
check_non_nan(max, "max")?;
+ check_finite(min, "min")?;
+ check_finite(max, "max")?;
let mut centroids = Vec::with_capacity(num_centroids);
let mut centroids_weight = 0u64;
for _ in 0..num_centroids {
@@ -594,9 +596,10 @@ impl TDigestMut {
check_non_nan(mean, "centroid mean")?;
check_finite(mean, "centroid")?;
let weight = check_nonzero(weight, "centroid weight")?;
- centroids_weight += weight.get();
+ centroids_weight = checked_weight_sum(centroids_weight, weight.get())?;
centroids.push(Centroid { mean, weight });
}
+ checked_weight_sum(centroids_weight, num_buffered as u64)?;
let mut buffer = Vec::with_capacity(num_buffered);
for _ in 0..num_buffered {
let value = if is_f32 {
@@ -643,6 +646,8 @@ impl TDigestMut {
let max = cursor.read_f64_be().map_err(make_error("max"))?;
check_non_nan(min, "min in compat double format")?;
check_non_nan(max, "max in compat double format")?;
+ check_finite(min, "min in compat double format")?;
+ check_finite(max, "max in compat double format")?;
let k = cursor.read_f64_be().map_err(make_error("k"))? as u16;
if k < 10 {
return Err(Error::deserial(format!(
@@ -654,12 +659,13 @@ impl TDigestMut {
let mut total_weight = 0u64;
let mut centroids = Vec::with_capacity(num_centroids);
for _ in 0..num_centroids {
- let weight = cursor.read_f64_be().map_err(make_error("weight"))? as u64;
+ let weight = cursor.read_f64_be().map_err(make_error("weight"))?;
let mean = cursor.read_f64_be().map_err(make_error("mean"))?;
- let weight = check_nonzero(weight, "centroid weight in compat double format")?;
+ let weight =
+ check_compat_weight(weight, "centroid weight in compat double format")?;
check_non_nan(mean, "centroid mean in compat double format")?;
check_finite(mean, "centroid mean in compat double format")?;
- total_weight += weight.get();
+ total_weight = checked_weight_sum(total_weight, weight.get())?;
centroids.push(Centroid { mean, weight });
}
Ok(TDigestMut::make(
@@ -682,6 +688,8 @@ impl TDigestMut {
let max = cursor.read_f64_be().map_err(make_error("max"))?;
check_non_nan(min, "min in compat float format")?;
check_non_nan(max, "max in compat float format")?;
+ check_finite(min, "min in compat float format")?;
+ check_finite(max, "max in compat float format")?;
let k = cursor.read_f32_be().map_err(make_error("k"))? as u16;
if k < 10 {
return Err(Error::deserial(format!(
@@ -696,12 +704,13 @@ impl TDigestMut {
let mut total_weight = 0u64;
let mut centroids = Vec::with_capacity(num_centroids);
for _ in 0..num_centroids {
- let weight = cursor.read_f32_be().map_err(make_error("weight"))? as u64;
+ let weight = cursor.read_f32_be().map_err(make_error("weight"))? as f64;
let mean = cursor.read_f32_be().map_err(make_error("mean"))? as f64;
- let weight = check_nonzero(weight, "centroid weight in compat float format")?;
+ let weight =
+ check_compat_weight(weight, "centroid weight in compat float format")?;
check_non_nan(mean, "centroid mean in compat float format")?;
check_finite(mean, "centroid mean in compat float format")?;
- total_weight += weight.get();
+ total_weight = checked_weight_sum(total_weight, weight.get())?;
centroids.push(Centroid { mean, weight });
}
Ok(TDigestMut::make(
@@ -1327,6 +1336,23 @@ fn check_nonzero(value: u64, tag: &'static str) -> Result<NonZeroU64, Error> {
.ok_or_else(|| Error::deserial(format!("malformed data: {tag} cannot be zero")))
}
+fn check_compat_weight(value: f64, tag: &'static str) -> Result<NonZeroU64, Error> {
+ check_non_nan(value, tag)?;
+ check_finite(value, tag)?;
+ if !(1.0..u64::MAX as f64).contains(&value) {
+ return Err(Error::deserial(format!(
+ "malformed data: {tag} must be representable as a positive u64"
+ )));
+ }
+ check_nonzero(value as u64, tag)
+}
+
+fn checked_weight_sum(total_weight: u64, weight: u64) -> Result<u64, Error> {
+ total_weight
+ .checked_add(weight)
+ .ok_or_else(|| Error::deserial("malformed data: total weight overflow"))
+}
+
/// Generates cluster sizes proportional to `q*(1-q)`.
///
/// The use of a normalizing function results in a strictly bounded number of clusters no matter
diff --git a/datasketches/tests/serde_tests/tdigest.rs b/datasketches/tests/serde_tests/tdigest.rs
index bfe4d40..8f216bb 100644
--- a/datasketches/tests/serde_tests/tdigest.rs
+++ b/datasketches/tests/serde_tests/tdigest.rs
@@ -18,6 +18,7 @@
use std::fs;
use std::path::PathBuf;
+use datasketches::error::ErrorKind;
use datasketches::tdigest::TDigestMut;
use googletest::assert_that;
use googletest::prelude::eq;
@@ -25,6 +26,18 @@ use googletest::prelude::near;
use crate::serialization_test_data;
+fn assert_invalid_data(bytes: &[u8], is_f32: bool) {
+ let error = TDigestMut::deserialize(bytes, is_f32).unwrap_err();
+ assert_eq!(error.kind(), ErrorKind::InvalidData);
+}
+
+fn serialize_multiple_values() -> Vec<u8> {
+ let mut tdigest = TDigestMut::default();
+ tdigest.update(0.0);
+ tdigest.update(1.0);
+ tdigest.serialize()
+}
+
fn test_sketch_file(path: PathBuf, n: u64, with_buffer: bool, is_f32: bool) {
let bytes = fs::read(&path).unwrap();
let td = TDigestMut::deserialize(&bytes, is_f32).unwrap();
@@ -186,3 +199,65 @@ fn test_many_values() {
assert_eq!(td.rank(500.0), deserialized_td.rank(500.0));
assert_eq!(td.quantile(0.5), deserialized_td.quantile(0.5));
}
+
+#[test]
+fn test_deserialize_rejects_non_finite_extrema() {
+ let bytes = serialize_multiple_values();
+ for (offset, value) in [(16, f64::INFINITY), (24, f64::NEG_INFINITY)] {
+ let mut malformed = bytes.clone();
+ malformed[offset..offset + size_of::<f64>()].copy_from_slice(&value.to_le_bytes());
+ assert_invalid_data(&malformed, false);
+ }
+
+ for filename in [
+ "tdigest_ref_k100_n10000_double.sk",
+ "tdigest_ref_k100_n10000_float.sk",
+ ] {
+ let path = serialization_test_data("reference_files", filename);
+ let bytes = fs::read(path).unwrap();
+ for (offset, value) in [(4, f64::INFINITY), (12, f64::NEG_INFINITY)] {
+ let mut malformed = bytes.clone();
+ malformed[offset..offset + size_of::<f64>()].copy_from_slice(&value.to_be_bytes());
+ assert_invalid_data(&malformed, false);
+ }
+ }
+}
+
+#[test]
+fn test_deserialize_rejects_invalid_compat_weights() {
+ let path = serialization_test_data("reference_files", "tdigest_ref_k100_n10000_double.sk");
+ let bytes = fs::read(path).unwrap();
+ for value in [f64::INFINITY, f64::MAX] {
+ let mut malformed = bytes.clone();
+ malformed[32..40].copy_from_slice(&value.to_be_bytes());
+ assert_invalid_data(&malformed, false);
+ }
+
+ let path = serialization_test_data("reference_files", "tdigest_ref_k100_n10000_float.sk");
+ let mut bytes = fs::read(path).unwrap();
+ bytes[30..34].copy_from_slice(&f32::INFINITY.to_be_bytes());
+ assert_invalid_data(&bytes, false);
+}
+
+#[test]
+fn test_deserialize_rejects_total_weight_overflow() {
+ let bytes = serialize_multiple_values();
+ let mut overflowing_centroids = bytes.clone();
+ overflowing_centroids[40..48].copy_from_slice(&u64::MAX.to_le_bytes());
+ overflowing_centroids[56..64].copy_from_slice(&1_u64.to_le_bytes());
+ assert_invalid_data(&overflowing_centroids, false);
+
+ let mut overflowing_buffer = bytes;
+ overflowing_buffer[8..12].copy_from_slice(&1_u32.to_le_bytes());
+ overflowing_buffer[12..16].copy_from_slice(&1_u32.to_le_bytes());
+ overflowing_buffer[40..48].copy_from_slice(&u64::MAX.to_le_bytes());
+ overflowing_buffer.truncate(56);
+ assert_invalid_data(&overflowing_buffer, false);
+
+ let path = serialization_test_data("reference_files", "tdigest_ref_k100_n10000_double.sk");
+ let mut overflowing_compat = fs::read(path).unwrap();
+ let near_max = u64::MAX as f64 - 2_048.0;
+ overflowing_compat[32..40].copy_from_slice(&near_max.to_be_bytes());
+ overflowing_compat[48..56].copy_from_slice(&4_096_f64.to_be_bytes());
+ assert_invalid_data(&overflowing_compat, false);
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datasketches.apache.org
For additional commands, e-mail: commits-help@datasketches.apache.org