guard_core_engine/redos/
probe_batches.rs

1//! Probe-set digests, timing validation, and batch measurement.
2//!
3//! Port of the reference `_redos_probe_batches.py`.
4
5use sha2::{Digest, Sha256};
6
7use super::child::ChildOutcome;
8
9/// Reference `_REACH_PROBE_BATCH_SIZE`.
10pub const REACH_PROBE_BATCH_SIZE: usize = 128;
11
12/// One timed probe set with its validated rows and host load factor.
13pub type TimedProbeSet = (Vec<String>, Option<Vec<Vec<f64>>>, f64);
14
15/// A validated reach-probe timing (the reference `ReachProbeTiming`).
16#[derive(Debug, Clone, PartialEq)]
17pub struct ReachProbeTiming {
18    /// Sorted samples per probe row.
19    pub samples_by_size: Vec<Vec<f64>>,
20    /// Host load factor the samples were measured under.
21    pub load_factor: f64,
22}
23
24/// Reference `_probe_set_digest`.
25#[must_use]
26pub fn probe_set_digest(probes: &[String]) -> [u8; 32] {
27    let mut digest = Sha256::new();
28    digest.update((probes.len() as u32).to_be_bytes());
29    for probe in probes {
30        let encoded = probe.as_bytes();
31        digest.update((encoded.len() as u64).to_be_bytes());
32        digest.update(encoded);
33    }
34    digest.finalize().into()
35}
36
37fn valid_sample(value: f64) -> bool {
38    value.is_finite() && value >= 0.0
39}
40
41fn valid_sample_row(row: &[f64]) -> bool {
42    !row.is_empty() && row.iter().all(|sample| valid_sample(*sample))
43}
44
45fn valid_load_factor(value: f64) -> bool {
46    valid_sample(value) && value > 0.0
47}
48
49/// Reference `_valid_timing_rows`: validates the row count and every
50/// sample, then splits rows and load factor.
51pub fn valid_timing_rows(
52    timing: Option<&ReachProbeTiming>,
53    expected_rows: usize,
54) -> Option<(Vec<Vec<f64>>, f64)> {
55    let timing = timing?;
56    if timing.samples_by_size.len() != expected_rows || !valid_load_factor(timing.load_factor) {
57        return None;
58    }
59    if !timing
60        .samples_by_size
61        .iter()
62        .all(|row| valid_sample_row(row))
63    {
64        return None;
65    }
66    Some((timing.samples_by_size.clone(), timing.load_factor))
67}
68
69/// Decode a reach-timing child payload into validated rows plus the
70/// reference scan (reference `_decode_reach_timing`).
71pub fn decode_reach_timing(outcome: Option<&ChildOutcome>) -> Option<(Vec<Vec<f64>>, f64)> {
72    match outcome? {
73        ChildOutcome::Timing { results, reference } => {
74            if results.is_empty() || !results.iter().all(|row| valid_sample_row(row)) {
75                return None;
76            }
77            if !valid_sample(*reference) {
78                return None;
79            }
80            let rows = results
81                .iter()
82                .map(|row| {
83                    let mut row = row.clone();
84                    row.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
85                    row
86                })
87                .collect();
88            Some((rows, *reference))
89        }
90        _ => None,
91    }
92}
93
94/// Reference `_batched_reach_probe_timings`.
95pub fn batched_reach_probe_timings(
96    pattern: &str,
97    mut probe_sets: Vec<Vec<String>>,
98    deadline: std::time::Instant,
99    flags: super::ast::Flags,
100    time_probes: super::cost_arbiter::TimeProbes<'_>,
101) -> Vec<TimedProbeSet> {
102    let mut out: Vec<TimedProbeSet> = Vec::new();
103    let mut batch: Vec<Vec<String>> = Vec::new();
104    for probes in probe_sets.drain(..) {
105        batch.push(probes);
106        if batch.len() < REACH_PROBE_BATCH_SIZE {
107            continue;
108        }
109        out.extend(measure_probe_batch(
110            pattern,
111            batch,
112            deadline,
113            flags,
114            time_probes,
115        ));
116        batch = Vec::new();
117    }
118    if !batch.is_empty() {
119        out.extend(measure_probe_batch(
120            pattern,
121            batch,
122            deadline,
123            flags,
124            time_probes,
125        ));
126    }
127    out
128}
129
130fn measure_probe_batch(
131    pattern: &str,
132    batch: Vec<Vec<String>>,
133    deadline: std::time::Instant,
134    flags: super::ast::Flags,
135    time_probes: super::cost_arbiter::TimeProbes<'_>,
136) -> Vec<TimedProbeSet> {
137    let flattened: Vec<String> = batch.iter().flatten().cloned().collect();
138    let timing = time_probes(pattern, flattened.clone(), deadline, flags);
139    let valid = valid_timing_rows(timing.as_ref(), flattened.len());
140    let Some((rows, load_factor)) = valid else {
141        return batch
142            .into_iter()
143            .map(|probes| (probes, None, 1.0))
144            .collect();
145    };
146    let mut out: Vec<TimedProbeSet> = Vec::new();
147    let mut offset = 0usize;
148    for probes in batch {
149        let next_offset = offset + probes.len();
150        out.push((
151            probes,
152            Some(rows[offset..next_offset].to_vec()),
153            load_factor,
154        ));
155        offset = next_offset;
156    }
157    out
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    fn probes(pairs: &[&str]) -> Vec<String> {
165        pairs.iter().map(|p| (*p).to_owned()).collect()
166    }
167
168    #[test]
169    fn digest_preserves_probe_set_boundaries() {
170        assert_ne!(
171            probe_set_digest(&probes(&["a", "bc"])),
172            probe_set_digest(&probes(&["ab", "c"]))
173        );
174        assert_eq!(
175            probe_set_digest(&probes(&["a", "bc"])),
176            probe_set_digest(&probes(&["a", "bc"]))
177        );
178    }
179
180    #[test]
181    fn valid_timing_rows_reject_shape_and_value_drift() {
182        let timing = ReachProbeTiming {
183            samples_by_size: vec![vec![0.1, 0.2]],
184            load_factor: 1.0,
185        };
186        assert_eq!(
187            valid_timing_rows(Some(&timing), 2),
188            None,
189            "row count must match"
190        );
191        let mut nan = timing.clone();
192        nan.samples_by_size = vec![vec![f64::NAN]];
193        assert_eq!(valid_timing_rows(Some(&nan), 1), None);
194        let mut negative = timing.clone();
195        negative.samples_by_size = vec![vec![-1.0]];
196        assert_eq!(valid_timing_rows(Some(&negative), 1), None);
197        let mut empty_row = timing.clone();
198        empty_row.samples_by_size = vec![vec![]];
199        assert_eq!(valid_timing_rows(Some(&empty_row), 1), None);
200        let mut load = timing;
201        load.load_factor = 0.0;
202        assert_eq!(valid_timing_rows(Some(&load), 1), None);
203        assert_eq!(valid_timing_rows(None, 1), None);
204    }
205
206    #[test]
207    fn valid_timing_rows_accepts_clean_rows() {
208        let timing = ReachProbeTiming {
209            samples_by_size: vec![vec![0.1, 0.2], vec![0.3]],
210            load_factor: 2.5,
211        };
212        let (rows, load) = valid_timing_rows(Some(&timing), 2).expect("valid");
213        assert_eq!(rows.len(), 2);
214        assert_eq!(load, 2.5);
215    }
216
217    #[test]
218    fn decode_reach_timing_sorts_each_row() {
219        let outcome = ChildOutcome::Timing {
220            results: vec![vec![0.3, 0.1, 0.2]],
221            reference: 0.002,
222        };
223        let (rows, reference) = decode_reach_timing(Some(&outcome)).expect("decoded");
224        assert_eq!(rows, vec![vec![0.1, 0.2, 0.3]]);
225        assert_eq!(reference, 0.002);
226    }
227
228    #[test]
229    fn decode_reach_timing_rejects_other_outcomes_and_bad_rows() {
230        assert_eq!(
231            decode_reach_timing(Some(&ChildOutcome::Failed("x".into()))),
232            None
233        );
234        assert_eq!(decode_reach_timing(None), None);
235        let empty = ChildOutcome::Timing {
236            results: vec![],
237            reference: 0.1,
238        };
239        assert_eq!(decode_reach_timing(Some(&empty)), None);
240        let bad_reference = ChildOutcome::Timing {
241            results: vec![vec![0.1]],
242            reference: -1.0,
243        };
244        assert_eq!(decode_reach_timing(Some(&bad_reference)), None);
245        let not_finite = ChildOutcome::Timing {
246            results: vec![vec![f64::INFINITY]],
247            reference: 0.1,
248        };
249        assert_eq!(decode_reach_timing(Some(&not_finite)), None);
250    }
251
252    #[test]
253    fn batched_timings_preserve_probe_order_and_all_rows() {
254        let probe_sets: Vec<Vec<String>> = (0..(REACH_PROBE_BATCH_SIZE + 1))
255            .map(|index| probes(&[&format!("p{index}")]))
256            .collect();
257        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60);
258        let calls: std::sync::Mutex<Vec<Vec<String>>> = std::sync::Mutex::new(Vec::new());
259        let time_probes = |pattern: &str,
260                           batch: Vec<String>,
261                           _deadline: std::time::Instant,
262                           _flags: super::super::ast::Flags|
263         -> Option<ReachProbeTiming> {
264            assert_eq!(pattern, "test");
265            calls.lock().expect("mutex").push(batch.clone());
266            Some(ReachProbeTiming {
267                samples_by_size: batch.iter().map(|_probe| vec![0.0]).collect(),
268                load_factor: 1.0,
269            })
270        };
271        let result = batched_reach_probe_timings(
272            "test",
273            probe_sets.clone(),
274            deadline,
275            super::super::ast::Flags::default(),
276            &time_probes,
277        );
278        let calls = calls.into_inner().expect("mutex");
279        assert_eq!(calls.len(), 2);
280        assert_eq!(calls[0].len(), REACH_PROBE_BATCH_SIZE);
281        assert_eq!(calls[1], probe_sets[REACH_PROBE_BATCH_SIZE]);
282        let sets: Vec<Vec<String>> = result
283            .iter()
284            .map(|(probes, _rows, _load)| probes.clone())
285            .collect();
286        assert_eq!(sets, probe_sets);
287        assert_eq!(result[0].1, Some(vec![vec![0.0]]));
288        assert_eq!(
289            result.last().expect("nonempty").1.clone().expect("rows"),
290            vec![vec![0.0]]
291        );
292    }
293
294    #[test]
295    fn an_exact_batch_multiple_leaves_no_final_flush() {
296        let probe_sets: Vec<Vec<String>> = (0..(2 * REACH_PROBE_BATCH_SIZE))
297            .map(|index| probes(&[&format!("p{index}")]))
298            .collect();
299        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60);
300        let calls: std::sync::Mutex<Vec<usize>> = std::sync::Mutex::new(Vec::new());
301        let time_probes = |_pattern: &str,
302                           batch: Vec<String>,
303                           _deadline: std::time::Instant,
304                           _flags: super::super::ast::Flags|
305         -> Option<ReachProbeTiming> {
306            calls.lock().expect("mutex").push(batch.len());
307            Some(ReachProbeTiming {
308                samples_by_size: batch.iter().map(|_probe| vec![0.0]).collect(),
309                load_factor: 1.0,
310            })
311        };
312        let result = batched_reach_probe_timings(
313            "test",
314            probe_sets,
315            deadline,
316            super::super::ast::Flags::default(),
317            &time_probes,
318        );
319        let calls = calls.into_inner().expect("mutex");
320        // exactly two full batches and nothing left over
321        assert_eq!(calls, vec![REACH_PROBE_BATCH_SIZE, REACH_PROBE_BATCH_SIZE]);
322        assert_eq!(result.len(), 2 * REACH_PROBE_BATCH_SIZE);
323    }
324
325    #[test]
326    fn batched_timings_fail_closed_for_malformed_child_results() {
327        let probe_sets = vec![probes(&["a", "b"])];
328        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60);
329        let time_probes = |_pattern: &str,
330                           _batch: Vec<String>,
331                           _deadline: std::time::Instant,
332                           _flags: super::super::ast::Flags|
333         -> Option<ReachProbeTiming> { None };
334        let result = batched_reach_probe_timings(
335            "test",
336            probe_sets,
337            deadline,
338            super::super::ast::Flags::default(),
339            &time_probes,
340        );
341        assert_eq!(result.len(), 1);
342        assert_eq!(result[0].1, None);
343        assert_eq!(result[0].2, 1.0);
344    }
345
346    #[test]
347    fn reach_probe_timing_holds_rows_and_load() {
348        let timing = ReachProbeTiming {
349            samples_by_size: vec![vec![1.0]],
350            load_factor: 2.0,
351        };
352        let clone = timing.clone();
353        assert_eq!(clone, timing);
354    }
355}
guard_core_engine/redos/probe_batches.rs

guard_core_engine/redos/
probe_batches.rs
