From 6a35ca36f8ac7301a2fa20578ba7a3202d3058ac Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 26 May 2026 06:55:03 -0600 Subject: [PATCH] harden(ingest): validate inbound CMS/CountSketch wire dims before matrix reconstruct MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Defensive, non-breaking hardening on the modified-OTLP frequency-sketch ingest path. Before reconstructing a CMS / CountSketch matrix from the wire-declared `(rows, cols)`, validate the dims so a malformed / hostile payload fails gracefully instead of risking a degenerate matrix or a huge allocation. New shared `validate_sketch_dims` helper (in count_min_sketch_accumulator, reused by count_sketch_accumulator) rejects: - `rows < 1` or `cols < 1` (degenerate, no cells); - narrow-hash-budget violation `rows * ceil(log2(cols)) > 64` — the point at which the cross-language Packed64 wire hasher's per-row column slices overflow / alias the 64-bit word and the matrix-cell layout degrades (mirrors sketchlib's `MatrixFastHash::assert_compatible`); checking the column-index bits alone keeps the realistic configs 5x2048 / 5x4096 / 5x2000 valid; - obviously-oversized dims `rows * cols > 8M cells` to cap allocation. On rejection the decoder returns `Err` (never panics); the existing OTLP ingest caller already skips that data point and increments `decoded_failed`. That call site now logs dim-validation rejections at WARN (ordinary decode fallbacks stay at DEBUG to avoid spam). Valid configs decode unchanged. Unit tests added to both accumulators: a malformed-dim data point is skipped (Err, no panic), a valid one still decodes, plus zero/budget/cap/ extreme-value (saturating) cases. Co-Authored-By: Claude Opus 4.7 (1M context) --- data_plane/src/drivers/ingest/otel.rs | 42 +++-- .../operators/count_min_sketch_accumulator.rs | 152 +++++++++++++++++- .../operators/count_sketch_accumulator.rs | 75 ++++++++- 3 files changed, 250 insertions(+), 19 deletions(-) diff --git a/data_plane/src/drivers/ingest/otel.rs b/data_plane/src/drivers/ingest/otel.rs index 12b19faf..cf54ef28 100644 --- a/data_plane/src/drivers/ingest/otel.rs +++ b/data_plane/src/drivers/ingest/otel.rs @@ -1340,17 +1340,37 @@ async fn route_modified_otlp_sketches_to_precompute( } Err(e) => { decoded_failed += 1; - debug!( - "OTLP modified-proto sketch decode failed \ - (metric={}, kind={:?}, encoding={}, \ - bytes={}): {} — falling through to §5.2 \ - fallback", - metric.name, - dp.kind, - dp.encoding, - dp.sketch.len(), - e - ); + let msg = e.to_string(); + // Defensive dim-validation rejections + // (malformed / degenerate CMS / CountSketch + // dims — see `validate_sketch_dims`) signal a + // misbehaving producer, so surface them at + // WARN; ordinary decode fallbacks stay at + // DEBUG to avoid log spam. + if msg.contains("rejecting") { + warn!( + "OTLP modified-proto sketch dropped on dim \ + validation (metric={}, kind={:?}, \ + encoding={}, bytes={}): {}", + metric.name, + dp.kind, + dp.encoding, + dp.sketch.len(), + msg + ); + } else { + debug!( + "OTLP modified-proto sketch decode failed \ + (metric={}, kind={:?}, encoding={}, \ + bytes={}): {} — falling through to §5.2 \ + fallback", + metric.name, + dp.kind, + dp.encoding, + dp.sketch.len(), + msg + ); + } continue; } } diff --git a/data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs b/data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs index 85c74591..692f29b6 100644 --- a/data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs +++ b/data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs @@ -117,9 +117,12 @@ impl CountMinSketchAccumulator { }; let rows = state.rows as usize; let cols = state.cols as usize; - if rows == 0 || cols == 0 { - return Err(format!("CountMinState has zero dims (rows={rows}, cols={cols})").into()); - } + // Defensive dim validation BEFORE reconstructing the matrix: + // reject degenerate / narrow-hash-budget-violating / absurdly + // oversized dims so a malformed payload fails gracefully (the + // ingest caller skips the data point) instead of building a + // degenerate or huge matrix. + validate_sketch_dims("CountMinState", rows, cols)?; let expected_len = rows * cols; let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { format!( @@ -299,6 +302,74 @@ impl CountMinSketchAccumulator { } } +/// Defensive upper bound on the number of matrix cells (`rows * cols`) +/// we'll reconstruct from an inbound wire-declared CMS / CountSketch +/// dimension pair. A malformed / hostile payload could declare absurd +/// dims (e.g. `rows = cols = u32::MAX`) and trick the decoder into a +/// huge `Vec` allocation before the `counts_*.len() != rows*cols` +/// check ever runs. Realistic sketches are at most a few hundred rows +/// by tens-of-thousands of columns, so 8M cells (~64 MiB of f64) is a +/// generous ceiling that no legitimate producer reaches. +pub(crate) const MAX_SKETCH_CELLS: usize = 8 * 1024 * 1024; + +/// Validate an inbound, wire-declared `(rows, cols)` pair for a +/// matrix-backed frequency sketch (CMS / CountSketch) BEFORE any matrix +/// is reconstructed from it. Returns `Ok(())` for dimensions a +/// legitimate producer could have emitted, and an `Err` (never a panic) +/// for malformed / degenerate ones so the ingest path can skip the data +/// point and fall through to its existing decode-failure accounting. +/// +/// Rejections: +/// 1. `rows < 1` or `cols < 1` — a zero-dim matrix has no cells. +/// 2. Narrow-hash-budget violation. The cross-language wire hasher +/// (`sketchlib`'s `MatrixHashType::Packed64`) derives every row's +/// column index from disjoint bit-fields of a single 64-bit hash +/// word: row `r` reads `mask_bits = ceil(log2(cols))` bits at offset +/// `r * mask_bits`. Once `rows * mask_bits > 64` the per-row column +/// slices overflow / alias the 64-bit word and the matrix-cell +/// layout is no longer the one the producer hashed into — the sketch +/// is internally degenerate. This mirrors sketchlib's own +/// `MatrixFastHash::assert_compatible` budget (`rows * (mask_bits + +/// 1) <= 64`); we check the column-index bits alone so realistic +/// configs (5x2048, 5x4096, 5x2000) — for which the sign bits share +/// the top of the word without affecting the cell layout — still +/// pass. +/// 3. Obviously-oversized dims: `rows * cols > MAX_SKETCH_CELLS`, +/// guarding against a huge allocation from a malformed payload. +/// +/// `what` names the wire struct for the error message (e.g. +/// `"CountMinState"`). +pub(crate) fn validate_sketch_dims(what: &str, rows: usize, cols: usize) -> Result<(), String> { + if rows < 1 || cols < 1 { + return Err(format!( + "{what} has degenerate dims (rows={rows}, cols={cols}); rejecting" + )); + } + // mask_bits = ceil(log2(cols)); cols >= 1 here. ilog2 is floor(log2). + let mask_bits = if cols.is_power_of_two() { + cols.ilog2() as usize + } else { + cols.ilog2() as usize + 1 + }; + if rows.saturating_mul(mask_bits) > 64 { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) exceed the 64-bit \ + packed-hash column budget (rows * ceil(log2(cols)) = {} > 64); \ + the sketch's matrix-cell layout is degenerate, rejecting", + rows.saturating_mul(mask_bits) + )); + } + if rows.saturating_mul(cols) > MAX_SKETCH_CELLS { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) declare {} cells, \ + exceeding the {MAX_SKETCH_CELLS}-cell ingest cap; rejecting to \ + avoid a huge allocation from a malformed payload", + rows.saturating_mul(cols) + )); + } + Ok(()) +} + impl SerializableToSink for CountMinSketchAccumulator { fn serialize_to_json(&self) -> Value { serde_json::json!({ @@ -916,6 +987,81 @@ mod tests { assert_eq!(v, 12.0); } + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / degenerate / narrow-hash-budget-violating CMS dims must + // be rejected gracefully (Err, never a panic); valid configs the + // backend actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + /// Build a bare `CountMinState` proto carrying the given dims and a + /// row-major INT64 counts vector sized to `rows*cols` so that, IF the + /// dims pass validation, the reshape also succeeds. Used to prove a + /// malformed-dim payload is rejected at the dim gate, not later. + fn cms_state_bytes(rows: u32, cols: u32) -> Vec { + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + let n = (rows as usize).saturating_mul(cols as usize); + let state = CountMinState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int: vec![0i64; n], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + state.encode_to_vec() + } + + #[test] + fn test_validate_sketch_dims_accepts_valid_configs() { + // The realistic configs the backend uses must pass unchanged. + for (r, c) in [(5usize, 2048usize), (5, 4096), (5, 2000), (4, 1000), (2, 3)] { + assert!( + validate_sketch_dims("CountMinState", r, c).is_ok(), + "valid config {r}x{c} was wrongly rejected" + ); + } + } + + #[test] + fn test_validate_sketch_dims_rejects_malformed() { + // Zero dims. + assert!(validate_sketch_dims("CountMinState", 0, 2048).is_err()); + assert!(validate_sketch_dims("CountMinState", 5, 0).is_err()); + // Narrow-hash-budget violation: 5 * ceil(log2(8192))=5*13=65 > 64. + let err = validate_sketch_dims("CountMinState", 5, 8192).unwrap_err(); + assert!(err.contains("budget"), "expected budget error, got: {err}"); + // Absurdly oversized: 1 x 16,777,216 = 16M cells > 8M cap. (1 row + // keeps the hash budget tiny — 1*24=24 — so the cap check, not the + // budget check, is what fires here.) + let err = validate_sketch_dims("CountMinState", 1, 16_777_216).unwrap_err(); + assert!(err.contains("cap"), "expected cell-cap error, got: {err}"); + // No panic on extreme dims (saturating_mul guards the products). + assert!(validate_sketch_dims("CountMinState", usize::MAX, usize::MAX).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + // A data point declaring narrow-hash-budget-violating dims must be + // skipped (Err returned, NOT a panic). The ingest caller turns + // this Err into a dropped data point + WARN log. + let bytes = cms_state_bytes(5, 8192); + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let ok_bytes = cms_state_bytes(5, 4096); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CMS should still decode"); + assert_eq!(acc.inner.rows(), 5); + assert_eq!(acc.inner.cols(), 4096); + } + #[test] fn test_query_statistic_rate_rejects_invalid_range_ms() { let cms = CountMinSketchAccumulator::new(2, 2); diff --git a/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs b/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs index c94bb9f3..cca852c0 100644 --- a/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs +++ b/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs @@ -91,11 +91,17 @@ impl CountSketchAccumulator { }; let rows = state.rows as usize; let cols = state.cols as usize; - if rows == 0 || cols == 0 { - return Err( - format!("CountSketchState has zero dims (rows={rows}, cols={cols})").into(), - ); - } + // Defensive dim validation BEFORE reconstructing the matrix: + // reject degenerate / narrow-hash-budget-violating / absurdly + // oversized dims so a malformed payload fails gracefully (the + // ingest caller skips the data point) instead of building a + // degenerate or huge matrix. Shares the CMS validator since the + // CountSketch matrix uses the same packed-hash column layout. + crate::precompute_engine::operators::count_min_sketch_accumulator::validate_sketch_dims( + "CountSketchState", + rows, + cols, + )?; let expected_len = rows * cols; let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { format!( @@ -566,4 +572,63 @@ mod tests { let mut acc = CountSketchAccumulator::new(2, 3); assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); } + + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / narrow-hash-budget-violating CountSketch dims must be + // rejected gracefully (Err, never a panic); valid configs the backend + // actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + use asap_sketchlib::proto::sketchlib::CounterType; + // 5 * ceil(log2(8192))=5*13=65 > 64 — narrow-hash-budget violation. + // counts sized to rows*cols so rejection is on dims, not length. + let n = 5usize * 8192usize; + let bytes = encode_state( + 5, + 8192, + CounterType::Int64 as i32, + vec![0i64; n], + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let n_ok = 5usize * 4096usize; + let ok_bytes = encode_state( + 5, + 4096, + CounterType::Int64 as i32, + vec![0i64; n_ok], + Vec::new(), + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CountSketch should still decode"); + assert_eq!(acc.inner.rows, 5); + assert_eq!(acc.inner.cols, 4096); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_oversized_dims() { + use asap_sketchlib::proto::sketchlib::CounterType; + // Declare 1 x 16,777,216 = 16M cells (> 8M cap) but send an empty + // counts vector: validation must reject on the dim cap BEFORE the + // decoder tries to allocate/reshape a 16M-entry matrix. (1 row keeps + // the hash budget tiny so the cap check, not the budget check, fires.) + let bytes = encode_state( + 1, + 16_777_216, + CounterType::Int64 as i32, + Vec::new(), + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "oversized dims should be rejected"); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("cap"), "expected cell-cap error, got: {msg}"); + } }