guard_core_engine/patterns/
pickle.rs

1//! Pickle opcode-stream validation ported from
2//! `guard_core/handlers/_suspatterns_pickle.py` (spec 4.0.2).
3//!
4//! The reference walks a bounded 4096-byte window through the `CPython` pickle
5//! opcode dispatch (with class resolution, extension registry and persistent
6//! loading blocked). The walk answers exactly two questions: does the prefix
7//! before a candidate look like a valid opcode stream, and does the suffix
8//! after it reach a REDUCE/BUILD opcode without an error.
9
10const PICKLE_OPCODE_WORK_BUDGET_BYTES: usize = 4096;
11
12const REDUCE: u8 = 0x52; // 'R'
13const BUILD: u8 = 0x62; // 'b'
14const FRAME_OPCODE: u8 = 0x95;
15
16#[derive(Debug)]
17enum WalkError {
18    ShortRead,
19    Blocked,
20}
21
22struct WalkState<'a> {
23    window: &'a [u8],
24    pos: usize,
25    stack: Vec<u8>,
26    marks: Vec<usize>,
27    memo: std::collections::HashMap<u32, u8>,
28}
29
30fn pickle_read(state: &mut WalkState, size: usize) -> Result<Vec<u8>, WalkError> {
31    if state.pos + size > state.window.len() {
32        return Err(WalkError::ShortRead);
33    }
34    let bytes = state.window[state.pos..state.pos + size].to_vec();
35    state.pos += size;
36    Ok(bytes)
37}
38
39fn pickle_readline(state: &mut WalkState) -> Result<Vec<u8>, WalkError> {
40    let mut line = Vec::new();
41    loop {
42        if state.pos >= state.window.len() {
43            return Err(WalkError::ShortRead);
44        }
45        let byte = state.window[state.pos];
46        state.pos += 1;
47        line.push(byte);
48        if byte == 0x0a {
49            break;
50        }
51    }
52    Ok(line)
53}
54
55fn le_int(bytes: &[u8]) -> usize {
56    let mut value = 0usize;
57    for b in bytes.iter().rev() {
58        value = value * 0x100 + usize::from(*b);
59    }
60    value
61}
62
63fn push_mark(state: &mut WalkState) {
64    state.marks.push(state.stack.len());
65}
66
67fn pop_mark(state: &mut WalkState) -> Result<(), WalkError> {
68    let Some(mark) = state.marks.pop() else {
69        return Err(WalkError::Blocked);
70    };
71    // discard the marked segment like the reference's stack juggling
72    while state.stack.len() > mark {
73        state.stack.pop();
74    }
75    Ok(())
76}
77
78fn is_digits(text: &[u8]) -> bool {
79    !text.is_empty() && text.iter().all(u8::is_ascii_digit)
80}
81
82#[allow(
83    clippy::too_many_lines,
84    clippy::match_same_arms,
85    reason = "one arm per pickle opcode in reference-table order: arms with identical bodies are distinct opcodes that share behavior, and splitting or merging them would obscure the 1:1 opcode mapping"
86)]
87fn dispatch_opcode(state: &mut WalkState, key: u8) -> Result<(), WalkError> {
88    match key {
89        0x28 => {
90            // '(' MARK
91            push_mark(state);
92            Ok(())
93        }
94        0x30 => {
95            // '0' POP
96            state
97                .stack
98                .pop()
99                .map_or(Err(WalkError::Blocked), |_| Ok(()))
100        }
101        0x31 => pop_mark(state), // '1' POP_MARK
102        0x32 => {
103            // '2' DUP
104            let Some(last) = state.stack.last().copied() else {
105                return Err(WalkError::Blocked);
106            };
107            state.stack.push(last);
108            Ok(())
109        }
110        0x5d | 0x7d | 0x29 => {
111            // ']' EMPTY_LIST, '}' EMPTY_DICT, ')' EMPTY_TUPLE
112            state.stack.push(1);
113            Ok(())
114        }
115        0x6c | 0x74 | 0x64 => {
116            // 'l' LIST, 't' TUPLE, 'd' DICT
117            pop_mark(state)?;
118            state.stack.push(1);
119            Ok(())
120        }
121        0x61 => {
122            // 'a' APPEND
123            state
124                .stack
125                .pop()
126                .map_or(Err(WalkError::Blocked), |_| Ok(()))
127        }
128        0x65 => pop_mark(state), // 'e' APPENDS
129        0x73 => {
130            // 's' SETITEM
131            state
132                .stack
133                .pop()
134                .map_or(Err(WalkError::Blocked), |_| Ok(()))?;
135            state
136                .stack
137                .pop()
138                .map_or(Err(WalkError::Blocked), |_| Ok(()))
139        }
140        0x75 => pop_mark(state), // 'u' SETITEMS
141        0x4e | 0x89 | 0x88 => {
142            // 'N' NONE, NEWFALSE, NEWTRUE
143            state.stack.push(1);
144            Ok(())
145        }
146        0x49 => {
147            // 'I' INT
148            let data = pickle_readline(state)?;
149            let text = &data[..data.len().saturating_sub(1)];
150            if text == b"01" || text == b"00" {
151                return Ok(());
152            }
153            if text.first() == Some(&b'-') && is_digits(&text[1..]) {
154                return Ok(());
155            }
156            if is_digits(text) {
157                return Ok(());
158            }
159            Err(WalkError::Blocked)
160        }
161        0x4c => {
162            // 'L' LONG
163            let data = pickle_readline(state)?;
164            let mut text = &data[..data.len().saturating_sub(1)];
165            if text.last() == Some(&b'L') {
166                text = &text[..text.len() - 1];
167            }
168            if text.first() == Some(&b'-') && is_digits(&text[1..]) {
169                return Ok(());
170            }
171            if is_digits(text) {
172                return Ok(());
173            }
174            Err(WalkError::Blocked)
175        }
176        0x46 => {
177            // 'F' FLOAT
178            let data = pickle_readline(state)?;
179            let text = std::str::from_utf8(&data[..data.len().saturating_sub(1)])
180                .map_err(|_| WalkError::Blocked)?;
181            if text.parse::<f64>().is_ok() {
182                Ok(())
183            } else {
184                Err(WalkError::Blocked)
185            }
186        }
187        0x4a => {
188            pickle_read(state, 4)?;
189            Ok(())
190        }
191        0x4b => {
192            pickle_read(state, 1)?;
193            Ok(())
194        }
195        0x4d => {
196            pickle_read(state, 2)?;
197            Ok(())
198        }
199        0x47 => {
200            pickle_read(state, 8)?;
201            Ok(())
202        }
203        0x53 => {
204            // 'S' STRING
205            let data = pickle_readline(state)?;
206            let body = &data[..data.len().saturating_sub(1)];
207            if body.len() < 2
208                || body[0] != body[body.len() - 1]
209                || (body[0] != 0x22 && body[0] != 0x27)
210            {
211                return Err(WalkError::Blocked);
212            }
213            Ok(())
214        }
215        0x56 => {
216            pickle_readline(state)?;
217            Ok(())
218        }
219        0x58 => {
220            let length = le_int(&pickle_read(state, 4)?);
221            pickle_read(state, length)?;
222            Ok(())
223        }
224        0x8c => {
225            let length = pickle_read(state, 1)?[0] as usize;
226            pickle_read(state, length)?;
227            Ok(())
228        }
229        0x54 => {
230            let length = le_int(&pickle_read(state, 4)?);
231            pickle_read(state, length)?;
232            Ok(())
233        }
234        0x55 => {
235            let length = pickle_read(state, 1)?[0] as usize;
236            pickle_read(state, length)?;
237            Ok(())
238        }
239        0x42 => {
240            let length = le_int(&pickle_read(state, 4)?);
241            pickle_read(state, length)?;
242            Ok(())
243        }
244        0x8e => {
245            let length = le_int(&pickle_read(state, 8)?);
246            pickle_read(state, length)?;
247            Ok(())
248        }
249        0x43 => {
250            let length = pickle_read(state, 1)?[0] as usize;
251            pickle_read(state, length)?;
252            Ok(())
253        }
254        0x96 => {
255            let length = le_int(&pickle_read(state, 8)?);
256            pickle_read(state, length)?;
257            Ok(())
258        }
259        0x8a => {
260            let length = pickle_read(state, 1)?[0] as usize;
261            pickle_read(state, length)?;
262            Ok(())
263        }
264        0x8b => {
265            let length = le_int(&pickle_read(state, 4)?);
266            pickle_read(state, length)?;
267            Ok(())
268        }
269        0x80 => {
270            pickle_read(state, 1)?;
271            Ok(())
272        }
273        0x94 => {
274            // MEMOIZE
275            if state.stack.is_empty() {
276                return Err(WalkError::Blocked);
277            }
278            let synthetic = u32::try_from(state.memo.len()).unwrap_or(u32::MAX);
279            state.memo.insert(synthetic, 1);
280            Ok(())
281        }
282        0x71 => {
283            // 'q' BINPUT
284            let index = pickle_read(state, 1)?[0];
285            if state.stack.is_empty() {
286                return Err(WalkError::Blocked);
287            }
288            memo_set(state, u32::from(index));
289            Ok(())
290        }
291        0x72 => {
292            // 'r' LONG_BINPUT
293            let index = le_int(&pickle_read(state, 4)?);
294            if state.stack.is_empty() {
295                return Err(WalkError::Blocked);
296            }
297            memo_set(state, u32::try_from(index).unwrap_or(u32::MAX));
298            Ok(())
299        }
300        0x68 => {
301            // 'h' BINGET
302            let index = pickle_read(state, 1)?[0];
303            if !memo_has(state, u32::from(index)) {
304                return Err(WalkError::Blocked);
305            }
306            state.stack.push(1);
307            Ok(())
308        }
309        0x6a => {
310            // 'j' LONG_BINGET
311            let index = le_int(&pickle_read(state, 4)?);
312            if !memo_has(state, u32::try_from(index).unwrap_or(u32::MAX)) {
313                return Err(WalkError::Blocked);
314            }
315            state.stack.push(1);
316            Ok(())
317        }
318        0x67 => {
319            // 'g' GET
320            let data = pickle_readline(state)?;
321            let text = std::str::from_utf8(&data[..data.len().saturating_sub(1)])
322                .map_err(|_| WalkError::Blocked)?;
323            let index: usize = text.parse().map_err(|_| WalkError::Blocked)?;
324            if !memo_has(state, u32::try_from(index).unwrap_or(u32::MAX)) {
325                return Err(WalkError::Blocked);
326            }
327            state.stack.push(1);
328            Ok(())
329        }
330        // 'c' GLOBAL, 'i' INST, 'o' OBJ, NEWOBJ, NEWOBJ_EX, STACK_GLOBAL,
331        // PERSID, BINPERSID, EXT1/2/4: class resolution blocked
332        0x63 | 0x69 | 0x6f | 0x81 | 0x82 | 0x93 | 0x50 | 0x51 | 0x84 | 0x85 | 0x86 => {
333            Err(WalkError::Blocked)
334        }
335        _ => Err(WalkError::Blocked),
336    }
337}
338
339fn memo_set(state: &mut WalkState, index: u32) {
340    state.memo.insert(index, 1);
341}
342
343fn memo_has(state: &WalkState, index: u32) -> bool {
344    state.memo.get(&index).is_some_and(|v| *v == 1)
345}
346
347/// Returns `Some(true)` (verdict reached: stream valid / REDUCE or BUILD hit),
348/// `Some(false)` (unknown opcode, blocked resolution, malformed payload, or a
349/// complete window that exhausted without a verdict), or `None` (window
350/// exhausted before a verdict could form; treated like `true` by the
351/// reference's `is not False` checks).
352fn walk_opcodes(
353    state: &mut WalkState,
354    stop_at_reduce_or_build: bool,
355    is_complete: bool,
356) -> Option<bool> {
357    while state.pos < state.window.len() {
358        let key = match pickle_read(state, 1) {
359            Ok(bytes) => bytes[0],
360            Err(WalkError::ShortRead) => return (!is_complete).then_some(true),
361            Err(WalkError::Blocked) => return Some(false),
362        };
363        if stop_at_reduce_or_build && (key == REDUCE || key == BUILD) {
364            return Some(true);
365        }
366        if key == FRAME_OPCODE {
367            match pickle_read(state, 8) {
368                Ok(_) => continue,
369                Err(WalkError::ShortRead) => return (!is_complete).then_some(true),
370                Err(WalkError::Blocked) => return Some(false),
371            }
372        }
373        match dispatch_opcode(state, key) {
374            Ok(()) => {}
375            Err(WalkError::ShortRead) => return (!is_complete).then_some(true),
376            Err(WalkError::Blocked) => return Some(false),
377        }
378    }
379    if is_complete {
380        Some(!stop_at_reduce_or_build)
381    } else {
382        Some(true)
383    }
384}
385
386fn fresh_state(window: &[u8], seed_stack: bool) -> WalkState<'_> {
387    WalkState {
388        window,
389        pos: 0,
390        stack: if seed_stack { vec![1] } else { Vec::new() },
391        marks: Vec::new(),
392        memo: std::collections::HashMap::new(),
393    }
394}
395
396/// Python `_pickle_prefix_window_from_chars`; `None` when not byte-safe.
397fn window_from_chars(chars: &str) -> Option<Vec<u8>> {
398    let mut bytes = Vec::with_capacity(chars.len());
399    for c in chars.chars() {
400        let code = u32::from(c);
401        if code <= 0xff {
402            bytes.push(code as u8);
403        } else if (0xdc80..=0xdcff).contains(&code) {
404            bytes.push((code - 0xdc80 + 0x80) as u8);
405        } else {
406            return None;
407        }
408    }
409    Some(bytes)
410}
411
412#[must_use]
413pub fn pickle_prefix_is_opcode_stream(prefix: &str) -> bool {
414    if prefix.is_empty() || prefix.ends_with('\n') {
415        return true;
416    }
417    let is_complete = prefix.len() <= PICKLE_OPCODE_WORK_BUDGET_BYTES;
418    let Some(window) = window_from_chars(if is_complete {
419        prefix
420    } else {
421        &prefix[..PICKLE_OPCODE_WORK_BUDGET_BYTES]
422    }) else {
423        return false;
424    };
425    // a None verdict (incomplete window) is tolerated by the reference
426    walk_opcodes(&mut fresh_state(&window, false), false, is_complete) != Some(false)
427}
428
429#[must_use]
430pub fn pickle_suffix_reaches_reduce_or_build(suffix: &str) -> bool {
431    let is_complete = suffix.len() <= PICKLE_OPCODE_WORK_BUDGET_BYTES;
432    let Some(window) = window_from_chars(if is_complete {
433        suffix
434    } else {
435        &suffix[..PICKLE_OPCODE_WORK_BUDGET_BYTES]
436    }) else {
437        return false;
438    };
439    walk_opcodes(&mut fresh_state(&window, true), true, is_complete) != Some(false)
440}
441
442/// `_pickle_global_candidate_is_injection`: the prefix before the candidate
443/// must walk as a valid opcode stream and the suffix after the class-name
444/// group must reach REDUCE or BUILD.
445#[must_use]
446pub fn pickle_global_candidate_is_injection(
447    haystack: &str,
448    candidate: super::pyregex::Candidate,
449    group_one_end: usize,
450) -> bool {
451    if !pickle_prefix_is_opcode_stream(&haystack[..candidate.start]) {
452        return false;
453    }
454    pickle_suffix_reaches_reduce_or_build(&haystack[group_one_end..])
455}
456
457#[cfg(test)]
458mod tests {
459    use super::*;
460
461    #[test]
462    fn plain_text_prefix_is_stream_like() {
463        // empty or newline-terminated prefixes are trivially accepted
464        assert!(pickle_prefix_is_opcode_stream(""));
465        assert!(pickle_prefix_is_opcode_stream("hello\n"));
466    }
467
468    #[test]
469    fn proto_stream_prefix_walks() {
470        // "\x80\x04\x95" + 8 size bytes + "\x8c\x04main"
471        let mut bytes = vec![0x80u8, 0x04, 0x95];
472        bytes.extend_from_slice(&[0x00; 8]);
473        bytes.extend_from_slice(&[0x8c, 0x04, b'm', b'a', b'i', b'n']);
474        let text: String = bytes.iter().map(|b| char::from(*b)).collect();
475        assert!(pickle_prefix_is_opcode_stream(&text));
476    }
477
478    #[test]
479    fn reduce_opcode_stops_suffix_walk() {
480        let mut bytes = vec![0x80u8, 0x04, 0x95];
481        bytes.extend_from_slice(&[0x00; 8]);
482        bytes.push(b'R');
483        let text: String = bytes.iter().map(|b| char::from(*b)).collect();
484        assert!(pickle_suffix_reaches_reduce_or_build(&text));
485    }
486
487    #[test]
488    fn unknown_opcode_fails() {
489        assert!(!pickle_suffix_reaches_reduce_or_build("\u{1}nonsense"));
490    }
491}
guard_core_engine/patterns/pickle.rs

guard_core_engine/patterns/
pickle.rs
