wafrift_evolution/differential/
binary_search.rs

1use super::probe::{
2    Probe, ProbeTarget, baseline_probe, command_path_probes, command_separator_probes,
3    sql_keyword_probes, sql_tautology_probes, xss_event_probes, xss_function_probes,
4    xss_tag_probes,
5};
6
7/// High-level probe family used to narrow the search space quickly.
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum ProbeFamily {
10    Sql,
11    Xss,
12    Command,
13}
14
15/// Generate a minimal first-pass probe set for quick analysis.
16#[must_use]
17pub fn generate_quick_probes() -> Vec<Probe> {
18    let mut probes = Vec::new();
19    probes.push(baseline_probe("test_benign", "baseline"));
20    probes.extend(generate_family_probes(ProbeFamily::Sql));
21    probes.extend(generate_family_probes(ProbeFamily::Xss));
22    probes.extend(generate_family_probes(ProbeFamily::Command));
23    probes
24}
25
26/// Generate a focused probe batch for one family during staged analysis.
27#[must_use]
28pub fn generate_family_probes(family: ProbeFamily) -> Vec<Probe> {
29    match family {
30        ProbeFamily::Sql => vec![
31            Probe {
32                payload: "' OR 1=1--".into(),
33                tests: ProbeTarget::SqlTautology("1=1".into()),
34                description: "classic SQLi".into(),
35                expected_blocked: true,
36            },
37            Probe {
38                payload: "test SELECT test".into(),
39                tests: ProbeTarget::SqlKeyword("SELECT".into()),
40                description: "SQL keyword".into(),
41                expected_blocked: true,
42            },
43            Probe {
44                payload: "test UNION test".into(),
45                tests: ProbeTarget::SqlKeyword("UNION".into()),
46                description: "SQL UNION".into(),
47                expected_blocked: true,
48            },
49            Probe {
50                payload: "'".into(),
51                tests: ProbeTarget::SqlQuote,
52                description: "single quote".into(),
53                expected_blocked: true,
54            },
55        ],
56        ProbeFamily::Xss => {
57            let mut tags = xss_tag_probes();
58            let mut events = xss_event_probes();
59            let mut funcs = xss_function_probes();
60            vec![
61                tags.remove(0),
62                tags.remove(0),
63                tags.remove(0),
64                events.remove(0),
65                funcs.remove(0),
66                funcs.remove(5.min(funcs.len().saturating_sub(1))),
67            ]
68        }
69        ProbeFamily::Command => {
70            let mut seps = command_separator_probes();
71            let mut paths = command_path_probes();
72            vec![seps.remove(0), seps.remove(0), paths.remove(0)]
73        }
74    }
75}
76
77/// Generate focused follow-up probes for families that blocked in the quick pass.
78#[must_use]
79pub fn generate_follow_up_probes(families: &[ProbeFamily]) -> Vec<Probe> {
80    let mut probes = Vec::new();
81    for family in families {
82        probes.extend(match family {
83            ProbeFamily::Sql => {
84                let mut sql = sql_keyword_probes();
85                sql.extend(sql_tautology_probes());
86                sql
87            }
88            ProbeFamily::Xss => {
89                let mut xss = xss_tag_probes();
90                xss.extend(xss_event_probes());
91                xss.extend(xss_function_probes());
92                xss
93            }
94            ProbeFamily::Command => {
95                let mut command = command_separator_probes();
96                command.extend(command_path_probes());
97                command
98            }
99        });
100    }
101    probes
102}
103
104/// Result of a binary search narrowing operation.
105#[derive(Debug, Clone)]
106pub struct NarrowingResult {
107    pub trigger: String,
108    pub start: usize,
109    pub end: usize,
110    pub probes_sent: usize,
111    pub description: String,
112}
113
114/// Binary search to find the minimum substring that triggers a WAF block.
115pub fn narrow_to_trigger(payload: &str, is_blocked: &dyn Fn(&str) -> bool) -> NarrowingResult {
116    let chars: Vec<char> = payload.chars().collect();
117    let len = chars.len();
118    let mut probes_sent = 0usize;
119    if len == 0 {
120        return NarrowingResult {
121            trigger: String::new(),
122            start: 0,
123            end: 0,
124            probes_sent,
125            description: "Empty payload cannot be narrowed".to_string(),
126        };
127    }
128
129    probes_sent += 1;
130    if !is_blocked(payload) {
131        return NarrowingResult {
132            trigger: payload.to_string(),
133            start: 0,
134            end: len,
135            probes_sent,
136            description: "Payload did not trigger a block during narrowing".to_string(),
137        };
138    }
139
140    let mut start = 0usize;
141    let mut end = len;
142
143    loop {
144        let removable_prefix =
145            max_removable_prefix(&chars, start, end, is_blocked, &mut probes_sent);
146        if removable_prefix > 0 {
147            start += removable_prefix;
148        }
149
150        let removable_suffix =
151            max_removable_suffix(&chars, start, end, is_blocked, &mut probes_sent);
152        if removable_suffix > 0 {
153            end -= removable_suffix;
154        }
155
156        if removable_prefix == 0 && removable_suffix == 0 {
157            break;
158        }
159    }
160
161    let trigger: String = chars[start..end].iter().collect();
162    probes_sent += 1;
163    let still_blocked = is_blocked(&trigger);
164
165    NarrowingResult {
166        trigger: trigger.clone(),
167        start,
168        end,
169        probes_sent,
170        description: if still_blocked {
171            format!(
172                "WAF trigger narrowed to '{}' ({} chars, positions {}-{} of {} char payload)",
173                if trigger.len() > 50 {
174                    &trigger[..50]
175                } else {
176                    &trigger
177                },
178                end - start,
179                start,
180                end,
181                len
182            )
183        } else {
184            "Could not narrow trigger (payload may use context-dependent matching)".to_string()
185        },
186    }
187}
188
189fn max_removable_prefix(
190    chars: &[char],
191    start: usize,
192    end: usize,
193    is_blocked: &dyn Fn(&str) -> bool,
194    probes_sent: &mut usize,
195) -> usize {
196    if end.saturating_sub(start) <= 1 {
197        return 0;
198    }
199    let mut lo = 0usize;
200    let mut hi = end - start - 1;
201    while lo < hi {
202        let mid = (lo + hi).div_ceil(2);
203        let candidate: String = chars[start + mid..end].iter().collect();
204        *probes_sent += 1;
205        if is_blocked(&candidate) {
206            lo = mid;
207        } else {
208            hi = mid - 1;
209        }
210    }
211    lo
212}
213
214fn max_removable_suffix(
215    chars: &[char],
216    start: usize,
217    end: usize,
218    is_blocked: &dyn Fn(&str) -> bool,
219    probes_sent: &mut usize,
220) -> usize {
221    if end.saturating_sub(start) <= 1 {
222        return 0;
223    }
224    let mut lo = 0usize;
225    let mut hi = end - start - 1;
226    while lo < hi {
227        let mid = (lo + hi).div_ceil(2);
228        let candidate: String = chars[start..end - mid].iter().collect();
229        *probes_sent += 1;
230        if is_blocked(&candidate) {
231            lo = mid;
232        } else {
233            hi = mid - 1;
234        }
235    }
236    lo
237}
238
239/// Find multiple independent triggers in a single payload.
240pub fn find_all_triggers(payload: &str, is_blocked: &dyn Fn(&str) -> bool) -> Vec<NarrowingResult> {
241    let mut triggers = Vec::new();
242    let mut remaining = payload.to_string();
243
244    for _ in 0..5 {
245        if !is_blocked(&remaining) {
246            break;
247        }
248        let result = narrow_to_trigger(&remaining, is_blocked);
249        if result.trigger.is_empty() || result.end <= result.start {
250            break;
251        }
252        let masked: String = remaining
253            .chars()
254            .enumerate()
255            .map(|(i, c)| {
256                if i >= result.start && i < result.end {
257                    'X'
258                } else {
259                    c
260                }
261            })
262            .collect();
263        triggers.push(result);
264        remaining = masked;
265    }
266
267    triggers
268}
269
270#[cfg(test)]
271mod tests {
272    use super::{
273        ProbeFamily, find_all_triggers, generate_family_probes, generate_follow_up_probes,
274        generate_quick_probes, narrow_to_trigger,
275    };
276
277    #[test]
278    fn quick_probes_smaller_set() {
279        let quick = generate_quick_probes();
280        assert!(quick.len() >= 10);
281    }
282
283    #[test]
284    fn family_probes_are_focused() {
285        assert_eq!(generate_family_probes(ProbeFamily::Sql).len(), 4);
286        assert_eq!(generate_family_probes(ProbeFamily::Command).len(), 3);
287    }
288
289    #[test]
290    fn follow_up_probes_expand_requested_families() {
291        let sql_only = generate_follow_up_probes(&[ProbeFamily::Sql]);
292        let both = generate_follow_up_probes(&[ProbeFamily::Sql, ProbeFamily::Xss]);
293        assert!(both.len() > sql_only.len());
294    }
295
296    #[test]
297    fn narrow_to_trigger_finds_minimal_substring() {
298        let payload = "prefixUNIONsuffix";
299        let result = narrow_to_trigger(payload, &|candidate| candidate.contains("UNION"));
300        assert_eq!(result.trigger, "UNION");
301        assert_eq!(result.start, 6);
302        assert_eq!(result.end, 11);
303    }
304
305    #[test]
306    fn find_all_triggers_masks_and_finds_multiple_regions() {
307        let payload = "aaaUNIONbbbSELECTccc";
308        let results = find_all_triggers(payload, &|candidate| {
309            candidate.contains("UNION") || candidate.contains("SELECT")
310        });
311        let triggers: Vec<_> = results
312            .iter()
313            .map(|result| result.trigger.as_str())
314            .collect();
315        assert!(triggers.contains(&"UNION"));
316        assert!(triggers.contains(&"SELECT"));
317    }
318}
wafrift_evolution/differential/binary_search.rs

wafrift_evolution/differential/
binary_search.rs
