Skip to main content

guard_core_engine/redos/
repeat_prefix_state.rs

1//! Prefix-walk state machine shared by the repeat-prefix machinery.
2//!
3//! Port of the reference `_redos_repeat_prefix_state.py`.
4
5use super::intervals::IntervalSet;
6use super::parse_slots::node_intervals;
7use super::timeout::BuilderTimeout;
8
9/// Reference `_REPEAT_PREFIX_BUDGET`.
10pub const REPEAT_PREFIX_BUDGET: usize = 24000;
11/// Reference `_REPEAT_PREFIX_STATE_LIMIT`.
12pub const REPEAT_PREFIX_STATE_LIMIT: usize = 1024;
13/// Reference `_PREFIX_WALK_STATE_LIMIT`.
14pub const PREFIX_WALK_STATE_LIMIT: usize = 16384;
15/// Reference `_PREFIX_WALK_TEXT_BUDGET`.
16pub const PREFIX_WALK_TEXT_BUDGET: usize = 2_000_000;
17
18/// The reference `_RepeatPrefixState`.
19#[derive(Debug, Clone, PartialEq, Eq, Hash)]
20pub struct RepeatPrefixState {
21    pub text: String,
22    pub pending: String,
23    pub excluded: IntervalSet,
24    pub forbidden: Vec<String>,
25    pub captures: Vec<(u32, String)>,
26    pub last_atom: Option<IntervalSet>,
27}
28
29impl Default for RepeatPrefixState {
30    fn default() -> Self {
31        Self {
32            text: String::new(),
33            pending: String::new(),
34            excluded: IntervalSet::empty(),
35            forbidden: Vec::new(),
36            captures: Vec::new(),
37            last_atom: None,
38        }
39    }
40}
41
42/// Reference `_state_text_size`.
43#[must_use]
44pub fn state_text_size(state: &RepeatPrefixState) -> usize {
45    state.text.len()
46        + state.pending.len()
47        + state.forbidden.iter().map(String::len).sum::<usize>()
48        + state
49            .captures
50            .iter()
51            .map(|(_group, value)| value.len())
52            .sum::<usize>()
53}
54
55fn merge_pending(left: &str, right: &str) -> Option<String> {
56    let common = left.len().min(right.len());
57    if left.as_bytes()[..common] != right.as_bytes()[..common] {
58        return None;
59    }
60    Some(
61        if left.len() >= right.len() {
62            left
63        } else {
64            right
65        }
66        .to_owned(),
67    )
68}
69
70fn remaining_forbidden(words: &[String], piece: &str) -> Option<Vec<String>> {
71    for word in words {
72        if piece.starts_with(word.as_str()) {
73            return None;
74        }
75    }
76    Some(
77        words
78            .iter()
79            .filter(|word| word.starts_with(piece))
80            .map(|word| word[piece.len()..].to_owned())
81            .collect(),
82    )
83}
84
85/// Reference `_consume_piece`.
86pub fn consume_piece(
87    state: &RepeatPrefixState,
88    piece: &str,
89) -> Result<Option<RepeatPrefixState>, BuilderTimeout> {
90    if piece.is_empty() {
91        return Ok(Some(state.clone()));
92    }
93    // unreachable: `piece` is non-empty (checked two lines above), so the
94    // first char always exists
95    #[cfg(not(coverage))]
96    let Some(first) = piece.chars().next() else {
97        return Ok(Some(state.clone()));
98    };
99    #[cfg(coverage)]
100    let first = piece.chars().next().expect("piece is non-empty");
101    if state.excluded.contains(u32::from(first)) {
102        return Ok(None);
103    }
104    let pending: Vec<char> = state.pending.chars().collect();
105    let piece_chars: Vec<char> = piece.chars().collect();
106    let common = pending.len().min(piece_chars.len());
107    if pending[..common] != piece_chars[..common] {
108        return Ok(None);
109    }
110    let Some(forbidden) = remaining_forbidden(&state.forbidden, piece) else {
111        return Ok(None);
112    };
113    if state.text.len() + piece.len() > REPEAT_PREFIX_BUDGET {
114        return Err(BuilderTimeout(
115            "Pattern validation repeat-prefix length budget exceeded".into(),
116        ));
117    }
118    let last = piece_chars.last().copied().unwrap_or(first);
119    Ok(Some(RepeatPrefixState {
120        text: format!("{}{piece}", state.text),
121        pending: pending[common..].iter().collect(),
122        excluded: IntervalSet::empty(),
123        forbidden,
124        captures: state.captures.clone(),
125        last_atom: Some(IntervalSet::single(u32::from(last))),
126    }))
127}
128
129/// Reference `_consume_atom`.
130pub fn consume_atom(
131    state: &RepeatPrefixState,
132    intervals: &IntervalSet,
133) -> Result<Option<RepeatPrefixState>, BuilderTimeout> {
134    let mut available = intervals.difference(&state.excluded);
135    for word in &state.forbidden {
136        let Some(first) = word.chars().next() else {
137            continue;
138        };
139        if word.chars().count() == 1 {
140            available = available.difference(&IntervalSet::single(u32::from(first)));
141        }
142    }
143    let member = match state.pending.chars().next() {
144        Some(c) => u32::from(c),
145        None => match available.first_member() {
146            Some(member) => member,
147            None => return Ok(None),
148        },
149    };
150    if !available.contains(member) {
151        return Ok(None);
152    }
153    let member_text = char::from_u32(member).map(String::from).unwrap_or_default();
154    let Some(consumed) = consume_piece(state, &member_text)? else {
155        return Ok(None);
156    };
157    Ok(Some(RepeatPrefixState {
158        last_atom: Some(available),
159        ..consumed
160    }))
161}
162
163/// Reference `_capture_state`.
164#[must_use]
165pub fn capture_state(
166    state: &RepeatPrefixState,
167    group: Option<u32>,
168    start: usize,
169) -> RepeatPrefixState {
170    let Some(group) = group else {
171        return state.clone();
172    };
173    let mut captures: Vec<(u32, String)> = state.captures.clone();
174    captures.retain(|(existing, _)| *existing != group);
175    captures.push((group, state.text[start.min(state.text.len())..].to_owned()));
176    captures.sort_by_key(|(existing, _)| *existing);
177    RepeatPrefixState {
178        captures,
179        ..state.clone()
180    }
181}
182
183/// Reference `_positive_assertion`.
184#[must_use]
185pub fn positive_assertion(
186    behind: bool,
187    witnesses: &[RepeatPrefixState],
188    state: &RepeatPrefixState,
189) -> Vec<RepeatPrefixState> {
190    if behind {
191        return witnesses
192            .iter()
193            .map(|witness| {
194                let text = if witness.text.is_empty() {
195                    state.text.clone()
196                } else {
197                    let cut = state.text.len().saturating_sub(witness.text.len());
198                    format!("{}{}", &state.text[..cut], witness.text)
199                };
200                let mut captures: Vec<(u32, String)> = state.captures.clone();
201                for (group, value) in &witness.captures {
202                    captures.retain(|(existing, _)| existing != group);
203                    captures.push((*group, value.clone()));
204                }
205                captures.sort_by_key(|(group, _)| *group);
206                RepeatPrefixState {
207                    text,
208                    captures,
209                    ..state.clone()
210                }
211            })
212            .collect();
213    }
214    witnesses
215        .iter()
216        .filter_map(|witness| {
217            merge_pending(&state.pending, &witness.text).map(|pending| {
218                let mut captures: Vec<(u32, String)> = state.captures.clone();
219                for (group, value) in &witness.captures {
220                    captures.retain(|(existing, _)| existing != group);
221                    captures.push((*group, value.clone()));
222                }
223                captures.sort_by_key(|(group, _)| *group);
224                RepeatPrefixState {
225                    pending,
226                    captures,
227                    ..state.clone()
228                }
229            })
230        })
231        .collect()
232}
233
234/// Reference `_negative_assertion`.
235#[must_use]
236pub fn negative_assertion(
237    body: &[super::ast::Op],
238    flags: super::ast::Flags,
239    witnesses: &[RepeatPrefixState],
240    state: &RepeatPrefixState,
241) -> Vec<RepeatPrefixState> {
242    if body.is_empty() {
243        return Vec::new();
244    }
245    if body.len() == 1 && body[0].is_pairing() {
246        return vec![RepeatPrefixState {
247            excluded: state.excluded.union(&node_intervals(&body[0], flags)),
248            ..state.clone()
249        }];
250    }
251    let words: Vec<String> = witnesses
252        .iter()
253        .filter(|witness| !witness.text.is_empty())
254        .map(|witness| witness.text.clone())
255        .collect();
256    let mut forbidden = state.forbidden.clone();
257    forbidden.extend(words);
258    vec![RepeatPrefixState {
259        forbidden,
260        ..state.clone()
261    }]
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use crate::redos::ast::{Flags, Op};
268
269    fn state(text: &str) -> RepeatPrefixState {
270        RepeatPrefixState {
271            text: text.to_owned(),
272            ..RepeatPrefixState::default()
273        }
274    }
275
276    #[test]
277    fn consume_piece_appends_text_and_tracks_the_last_atom() {
278        let consumed = consume_piece(&state("ab"), "cd").expect("no budget error");
279        let consumed = consumed.expect("piece consumed");
280        assert_eq!(consumed.text, "abcd");
281        assert_eq!(
282            consumed.last_atom,
283            Some(IntervalSet::single(u32::from('d')))
284        );
285    }
286
287    #[test]
288    fn consume_piece_empty_returns_the_state() {
289        let original = state("ab");
290        let consumed = consume_piece(&original, "").expect("no budget error");
291        assert_eq!(consumed, Some(original));
292    }
293
294    #[test]
295    fn consume_piece_rejects_an_excluded_first_character() {
296        let mut excluded = state("ab");
297        excluded.excluded = IntervalSet::single(u32::from('x'));
298        assert_eq!(
299            consume_piece(&excluded, "xy").expect("no budget error"),
300            None
301        );
302    }
303
304    #[test]
305    fn consume_piece_rejects_a_pending_mismatch() {
306        let mut pending = state("ab");
307        pending.pending = "xy".into();
308        assert_eq!(
309            consume_piece(&pending, "ax").expect("no budget error"),
310            None
311        );
312    }
313
314    #[test]
315    fn consume_piece_consumes_a_matching_pending_prefix() {
316        let mut pending = state("ab");
317        pending.pending = "cd".into();
318        let consumed = consume_piece(&pending, "c").expect("no budget error");
319        let consumed = consumed.expect("prefix consumed");
320        assert_eq!(consumed.text, "abc");
321        assert_eq!(consumed.pending, "d");
322    }
323
324    #[test]
325    fn consume_piece_forbidden_prefix_kills_the_state() {
326        // The kill condition is `piece.startswith(word)`: the consumed
327        // piece covers the whole forbidden word.
328        let mut forbidden = state("ab");
329        forbidden.forbidden = vec!["xyz".into()];
330        assert_eq!(
331            consume_piece(&forbidden, "xyz").expect("no budget error"),
332            None
333        );
334    }
335
336    #[test]
337    fn consume_piece_forbidden_word_survivors_drop_the_consumed_prefix() {
338        let mut forbidden = state("");
339        forbidden.forbidden = vec!["abcd".into()];
340        let consumed = consume_piece(&forbidden, "ab").expect("no budget error");
341        let consumed = consumed.expect("prefix consumed");
342        assert_eq!(consumed.forbidden, vec!["cd".to_owned()]);
343    }
344
345    #[test]
346    fn consume_piece_over_the_length_budget_is_an_error() {
347        let mut long = state("");
348        long.text = "a".repeat(REPEAT_PREFIX_BUDGET);
349        let error = consume_piece(&long, "b").expect_err("budget exceeded");
350        assert_eq!(
351            error.0,
352            "Pattern validation repeat-prefix length budget exceeded"
353        );
354    }
355
356    #[test]
357    fn consume_atom_picks_the_first_available_member() {
358        let intervals = IntervalSet::new(&[(u32::from('a'), u32::from('c'))]);
359        let consumed = consume_atom(&state(""), &intervals)
360            .expect("no budget error")
361            .expect("atom consumed");
362        assert_eq!(consumed.text, "a");
363        assert_eq!(consumed.last_atom, Some(intervals));
364    }
365
366    #[test]
367    fn consume_atom_honors_pending_before_availability() {
368        let intervals = IntervalSet::new(&[(u32::from('a'), u32::from('z'))]);
369        let mut pending = state("");
370        pending.pending = "q".into();
371        let consumed = consume_atom(&pending, &intervals)
372            .expect("no budget error")
373            .expect("atom consumed");
374        assert_eq!(consumed.text, "q");
375    }
376
377    #[test]
378    fn consume_atom_returns_none_when_everything_is_excluded() {
379        let intervals = IntervalSet::single(u32::from('a'));
380        let mut excluded = state("");
381        excluded.excluded = IntervalSet::single(u32::from('a'));
382        assert_eq!(
383            consume_atom(&excluded, &intervals).expect("no budget error"),
384            None
385        );
386    }
387
388    #[test]
389    fn consume_atom_filters_single_char_forbidden_words() {
390        let intervals = IntervalSet::new(&[(u32::from('a'), u32::from('b'))]);
391        let mut forbidden = state("");
392        forbidden.forbidden = vec!["a".into()];
393        let consumed = consume_atom(&forbidden, &intervals)
394            .expect("no budget error")
395            .expect("atom consumed");
396        assert_eq!(consumed.text, "b");
397    }
398
399    #[test]
400    fn capture_state_ignores_anonymous_groups() {
401        let original = state("ab");
402        assert_eq!(capture_state(&original, None, 0), original);
403    }
404
405    #[test]
406    fn capture_state_records_the_group_body() {
407        let captured = capture_state(&state("abcd"), Some(2), 1);
408        assert_eq!(captured.captures, vec![(2, "bcd".to_owned())]);
409    }
410
411    #[test]
412    fn capture_state_replaces_an_existing_group_number() {
413        let mut original = state("xyz");
414        original.captures = vec![(1, "old".into())];
415        let captured = capture_state(&original, Some(1), 1);
416        assert_eq!(captured.captures, vec![(1, "yz".to_owned())]);
417    }
418
419    #[test]
420    fn positive_assertion_lookahead_merges_pending() {
421        let witness = state("bc");
422        let mut base = state("a");
423        base.pending = "b".into();
424        let result = positive_assertion(false, &[witness], &base);
425        assert_eq!(result.len(), 1);
426        assert_eq!(result[0].pending, "bc");
427    }
428
429    #[test]
430    fn positive_assertion_lookahead_drops_a_disjoint_pending() {
431        let witness = state("bc");
432        let mut base = state("a");
433        base.pending = "z".into();
434        let result = positive_assertion(false, &[witness], &base);
435        assert!(result.is_empty());
436    }
437
438    #[test]
439    fn positive_assertion_lookahead_drops_incompatible_witnesses() {
440        let witness = state("bc");
441        let mut base = state("a");
442        base.pending = "x".into();
443        let result = positive_assertion(false, &[witness], &base);
444        assert!(result.is_empty());
445    }
446
447    #[test]
448    fn positive_assertion_lookbehind_swaps_the_witness_text() {
449        // The reference keeps the prefix before the witness span.
450        let witness = state("bc");
451        let result = positive_assertion(true, &[witness], &state("abc"));
452        assert_eq!(result.len(), 1);
453        assert_eq!(result[0].text, "abc");
454    }
455
456    #[test]
457    fn positive_assertion_merges_witness_captures() {
458        let mut witness = state("b");
459        witness.captures = vec![(3, "b".into())];
460        let result = positive_assertion(false, &[witness], &state("a"));
461        assert_eq!(result[0].captures, vec![(3, "b".to_owned())]);
462    }
463
464    #[test]
465    fn negative_assertion_empty_body_dies() {
466        let result = negative_assertion(&[], Flags::default(), &[], &state("a"));
467        assert!(result.is_empty());
468    }
469
470    #[test]
471    fn negative_assertion_pairing_body_excludes_the_intervals() {
472        let mut base = state("a");
473        base.excluded = IntervalSet::single(1);
474        let result =
475            negative_assertion(&[Op::Literal(u32::from('x'))], Flags::default(), &[], &base);
476        assert_eq!(result.len(), 1);
477        assert!(result[0].excluded.contains(u32::from('x')));
478    }
479
480    #[test]
481    fn negative_assertion_non_pairing_body_forbids_witness_words() {
482        let witness = state("xy");
483        let mut base = state("a");
484        base.forbidden = vec!["old".into()];
485        let result = negative_assertion(
486            &[Op::GroupRef(1)],
487            Flags::default(),
488            std::slice::from_ref(&witness),
489            &base,
490        );
491        assert_eq!(result[0].forbidden, vec!["old".to_owned(), "xy".to_owned()]);
492    }
493
494    #[test]
495    fn state_text_size_sums_every_text_field() {
496        let mut sized = state("abc");
497        sized.pending = "de".into();
498        sized.forbidden = vec!["fgh".into(), "i".into()];
499        sized.captures = vec![(1, "jk".into())];
500        assert_eq!(state_text_size(&sized), 3 + 2 + 4 + 2);
501    }
502
503    #[test]
504    fn merge_pending_keeps_the_longer_affine_tail() {
505        assert_eq!(merge_pending("abx", "ab"), Some("abx".to_owned()));
506        assert_eq!(merge_pending("ab", "abx"), Some("abx".to_owned()));
507        assert_eq!(merge_pending("ab", "ab"), Some("ab".to_owned()));
508        assert_eq!(merge_pending("ax", "ab"), None);
509    }
510
511    #[test]
512    fn consume_atom_skips_empty_forbidden_words() {
513        // A forbidden word with no first char cannot narrow the available
514        // set; the walk still consumes its first member.
515        let state = RepeatPrefixState {
516            forbidden: vec![String::new()],
517            ..RepeatPrefixState::default()
518        };
519        // The empty word cannot narrow the available set, but it does make
520        // `consume_piece` reject every piece, so the atom is refused.
521        assert_eq!(
522            consume_atom(&state, &IntervalSet::from_range(0x61, 0x62)).expect("walk"),
523            None
524        );
525    }
526
527    #[test]
528    fn consume_atom_rejects_a_pending_char_outside_the_atom() {
529        let state = RepeatPrefixState {
530            pending: "x".to_owned(),
531            ..RepeatPrefixState::default()
532        };
533        assert_eq!(
534            consume_atom(&state, &IntervalSet::single(0x61)).expect("walk"),
535            None
536        );
537    }
538
539    #[test]
540    fn positive_assertion_rewrites_text_and_captures() {
541        let state = RepeatPrefixState {
542            text: "abcd".to_owned(),
543            captures: vec![(1, "old".to_owned()), (2, "keep".to_owned())],
544            ..RepeatPrefixState::default()
545        };
546        let witness = RepeatPrefixState {
547            text: String::new(),
548            captures: vec![(1, "new".to_owned())],
549            ..RepeatPrefixState::default()
550        };
551        // An empty witness keeps the state text; the capture merge replaces
552        // group 1 and preserves group 2 in order.
553        let rewritten = &positive_assertion(true, &[witness], &state)[0];
554        assert_eq!(rewritten.text, "abcd");
555        assert_eq!(
556            rewritten.captures,
557            vec![(1, "new".to_owned()), (2, "keep".to_owned())]
558        );
559        // A non-empty witness splices its text at the cut point.
560        let witness = RepeatPrefixState {
561            text: "xy".to_owned(),
562            ..RepeatPrefixState::default()
563        };
564        let rewritten = &positive_assertion(true, &[witness], &state)[0];
565        assert_eq!(rewritten.text, "abxy");
566    }
567}