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        #[cfg(not(coverage))] // unreachable: the loop guard keeps one byte
359        // available, and a one-byte read never reports `Blocked`
360        let key = match pickle_read(state, 1) {
361            Ok(bytes) => bytes[0],
362            Err(WalkError::ShortRead) => return (!is_complete).then_some(true),
363            Err(WalkError::Blocked) => return Some(false),
364        };
365        #[cfg(coverage)]
366        let key = pickle_read(state, 1).expect("one byte is available under the loop guard")[0];
367        if stop_at_reduce_or_build && (key == REDUCE || key == BUILD) {
368            return Some(true);
369        }
370        if key == FRAME_OPCODE {
371            match pickle_read(state, 8) {
372                Ok(_) => continue,
373                // a truncated frame operand is tolerated like the reference;
374                // `pickle_read` never answers `Blocked` on a raw byte read
375                Err(WalkError::ShortRead | WalkError::Blocked) => {
376                    return (!is_complete).then_some(true);
377                }
378            }
379        }
380        match dispatch_opcode(state, key) {
381            Ok(()) => {}
382            Err(WalkError::ShortRead) => return (!is_complete).then_some(true),
383            Err(WalkError::Blocked) => return Some(false),
384        }
385    }
386    if is_complete {
387        Some(!stop_at_reduce_or_build)
388    } else {
389        Some(true)
390    }
391}
392
393fn fresh_state(window: &[u8], seed_stack: bool) -> WalkState<'_> {
394    WalkState {
395        window,
396        pos: 0,
397        stack: if seed_stack { vec![1] } else { Vec::new() },
398        marks: Vec::new(),
399        memo: std::collections::HashMap::new(),
400    }
401}
402
403/// Python `_pickle_prefix_window_from_chars`; `None` when not byte-safe.
404fn window_from_chars(chars: &str) -> Option<Vec<u8>> {
405    let mut bytes = Vec::with_capacity(chars.len());
406    for c in chars.chars() {
407        let code = u32::from(c);
408        if code <= 0xff {
409            bytes.push(code as u8);
410        } else {
411            #[cfg(not(coverage))] // unreachable: Rust chars exclude the
412            // surrogate range, so the surrogateescape byte marks never
413            // appear as chars in a `str`
414            if (0xdc80..=0xdcff).contains(&code) {
415                bytes.push((code - 0xdc80 + 0x80) as u8);
416                continue;
417            }
418            return None;
419        }
420    }
421    Some(bytes)
422}
423
424#[must_use]
425pub fn pickle_prefix_is_opcode_stream(prefix: &str) -> bool {
426    if prefix.is_empty() || prefix.ends_with('\n') {
427        return true;
428    }
429    let is_complete = prefix.len() <= PICKLE_OPCODE_WORK_BUDGET_BYTES;
430    let Some(window) = window_from_chars(if is_complete {
431        prefix
432    } else {
433        &prefix[..PICKLE_OPCODE_WORK_BUDGET_BYTES]
434    }) else {
435        return false;
436    };
437    // a None verdict (incomplete window) is tolerated by the reference
438    walk_opcodes(&mut fresh_state(&window, false), false, is_complete) != Some(false)
439}
440
441#[must_use]
442pub fn pickle_suffix_reaches_reduce_or_build(suffix: &str) -> bool {
443    let is_complete = suffix.len() <= PICKLE_OPCODE_WORK_BUDGET_BYTES;
444    let Some(window) = window_from_chars(if is_complete {
445        suffix
446    } else {
447        &suffix[..PICKLE_OPCODE_WORK_BUDGET_BYTES]
448    }) else {
449        return false;
450    };
451    walk_opcodes(&mut fresh_state(&window, true), true, is_complete) != Some(false)
452}
453
454/// `_pickle_global_candidate_is_injection`: the prefix before the candidate
455/// must walk as a valid opcode stream and the suffix after the class-name
456/// group must reach REDUCE or BUILD.
457#[must_use]
458pub fn pickle_global_candidate_is_injection(
459    haystack: &str,
460    candidate: super::pyregex::Candidate,
461    group_one_end: usize,
462) -> bool {
463    if !pickle_prefix_is_opcode_stream(&haystack[..candidate.start]) {
464        return false;
465    }
466    pickle_suffix_reaches_reduce_or_build(&haystack[group_one_end..])
467}
468
469#[cfg(test)]
470mod tests {
471    use super::*;
472
473    #[test]
474    fn plain_text_prefix_is_stream_like() {
475        // empty or newline-terminated prefixes are trivially accepted
476        assert!(pickle_prefix_is_opcode_stream(""));
477        assert!(pickle_prefix_is_opcode_stream("hello\n"));
478    }
479
480    #[test]
481    fn proto_stream_prefix_walks() {
482        // "\x80\x04\x95" + 8 size bytes + "\x8c\x04main"
483        let mut bytes = vec![0x80u8, 0x04, 0x95];
484        bytes.extend_from_slice(&[0x00; 8]);
485        bytes.extend_from_slice(&[0x8c, 0x04, b'm', b'a', b'i', b'n']);
486        let text: String = bytes.iter().map(|b| char::from(*b)).collect();
487        assert!(pickle_prefix_is_opcode_stream(&text));
488    }
489
490    #[test]
491    fn reduce_opcode_stops_suffix_walk() {
492        let mut bytes = vec![0x80u8, 0x04, 0x95];
493        bytes.extend_from_slice(&[0x00; 8]);
494        bytes.push(b'R');
495        let text: String = bytes.iter().map(|b| char::from(*b)).collect();
496        assert!(pickle_suffix_reaches_reduce_or_build(&text));
497    }
498
499    #[test]
500    fn unknown_opcode_fails() {
501        assert!(!pickle_suffix_reaches_reduce_or_build("\u{1}nonsense"));
502    }
503}
504
505#[cfg(test)]
506mod opcode_tests {
507    use super::*;
508
509    /// Builds the char string the public functions expect: ASCII bytes pass
510    /// through, high bytes ride the surrogate escape (`0xdc80 + b - 0x80`).
511    fn stream(bytes: &[u8]) -> String {
512        // Latin-1: every byte is a valid char (U+0000..=U+00FF), and
513        // `window_from_chars` maps those straight back to the byte.
514        bytes.iter().map(|b| char::from(*b)).collect()
515    }
516
517    #[test]
518    fn every_binary_opcode_walks_when_its_operand_fits() {
519        // PROTO, frame, then one opcode per operand shape: 1-byte, 2-byte,
520        // 4-byte, and 8-byte reads, plus the length-prefixed payloads.
521        let cases: Vec<Vec<u8>> = vec![
522            vec![0x80, 0x04],                               // PROTO
523            vec![0x95, 0, 0, 0, 0, 0, 0, 0, 0],             // FRAME
524            vec![0x4a, 1, 2, 3, 4],                         // BININT (4)
525            vec![0x4b, 0x7f],                               // BININT1 (1)
526            vec![0x4d, 0x01, 0x02],                         // BININT2 (2)
527            vec![0x47, 0, 0, 0, 0, 0, 0, 0, 0],             // BINFLOAT (8)
528            vec![0x58, 3, 0, 0, 0, b'a', b'b', b'c'],       // BINUNICODE
529            vec![0x8c, 4, b'm', b'a', b'i', b'n'],          // SHORT_BINUNICODE
530            vec![0x54, 2, 0, 0, 0, b'x', b'y'],             // BINSTRING
531            vec![0x55, 2, b'x', b'y'],                      // SHORT_BINSTRING
532            vec![0x42, 2, 0, 0, 0, b'x', b'y'],             // BINBYTES
533            vec![0x8e, 2, 0, 0, 0, 0, 0, 0, 0, b'x', b'y'], // BINBYTES8
534            vec![0x43, 2, b'x', b'y'],                      // SHORT_BINBYTES
535            vec![0x96, 1, 0, 0, 0, 0, 0, 0, 0, b'z'],       // BYTEARRAY8
536            vec![0x8a, 0x01, b'x'],                         // 1-byte length read
537            vec![0x8b, 1, 0, 0, 0, b'q'],                   // 4-byte len read
538        ];
539        for mut case in cases {
540            case.push(b'R'); // the suffix walk stops at the REDUCE
541            let text = stream(&case);
542            assert!(
543                pickle_suffix_reaches_reduce_or_build(&text),
544                "stream should walk clean: {case:?}"
545            );
546        }
547    }
548
549    #[test]
550    fn truncated_operand_streams_are_tolerated_like_the_reference() {
551        // Incomplete windows read as tolerated (never `Some(false)`).
552        assert!(pickle_suffix_reaches_reduce_or_build(&stream(&[
553            0x4a, 1, 2
554        ])));
555        assert!(pickle_prefix_is_opcode_stream(&stream(&[0x58, 9, 9, 9])));
556        assert!(pickle_suffix_reaches_reduce_or_build(&stream(&[
557            0x95, 0, 0
558        ])));
559    }
560
561    #[test]
562    fn text_opcodes_accept_their_reference_shapes() {
563        let cases: Vec<Vec<u8>> = vec![
564            vec![b'I', b'0', b'1', 0x0a],             // INT 01
565            vec![b'I', b'0', b'0', 0x0a],             // INT 00
566            vec![b'I', b'-', b'4', b'2', 0x0a],       // INT negative
567            vec![b'I', b'7', 0x0a],                   // INT digits
568            vec![b'L', b'-', b'9', b'9', 0x0a],       // LONG negative
569            vec![b'L', b'1', b'2', b'L', 0x0a],       // LONG with L suffix
570            vec![b'L', b'5', 0x0a],                   // LONG digits
571            vec![b'F', b'1', b'.', b'5', 0x0a],       // FLOAT
572            vec![b'F', b'-', b'0', b'.', b'5', 0x0a], // FLOAT negative
573            vec![b'S', b'\'', b'a', b'\'', 0x0a],     // STRING single-quoted
574            vec![b'S', b'"', b'a', b'"', 0x0a],       // STRING double-quoted
575            vec![b'V', b'u', b'n', b'i', 0x0a],       // UNICODE line
576            vec![b'N'],                               // NONE
577            vec![0x89],                               // NEWFALSE
578            vec![0x88],                               // NEWTRUE
579        ];
580        for mut case in cases {
581            case.push(b'R'); // the suffix walk stops at the REDUCE
582            let text = stream(&case);
583            let label = String::from_utf8_lossy(&case);
584            assert!(
585                pickle_suffix_reaches_reduce_or_build(&text),
586                "stream should walk clean: {label}"
587            );
588        }
589    }
590
591    #[test]
592    fn malformed_text_opcodes_are_blocked() {
593        let cases: Vec<Vec<u8>> = vec![
594            vec![b'I', b'x', b'y', 0x0a],        // INT not numeric
595            vec![b'L', b'x', 0x0a],              // LONG not numeric
596            vec![b'F', b'x', b'y', 0x0a],        // FLOAT not parseable
597            vec![b'S', b'a', b'b', 0x0a],        // STRING unquoted
598            vec![b'S', b'\'', b'a', b'"', 0x0a], // STRING mismatched quotes
599            vec![b'S', b'\'', 0x0a],             // STRING too short
600        ];
601        for case in cases {
602            let text = stream(&case);
603            let label = String::from_utf8_lossy(&case);
604            assert!(
605                !pickle_suffix_reaches_reduce_or_build(&text),
606                "stream should be blocked: {label}"
607            );
608        }
609    }
610
611    #[test]
612    fn stack_and_memo_opcodes_walk_in_valid_sequences() {
613        let cases: Vec<Vec<u8>> = vec![
614            vec![0x28, 0x30],             // MARK then POP
615            vec![0x28, b'N', 0x31],       // MARK, NONE, POP_MARK
616            vec![b'N', 0x32],             // NONE, DUP
617            vec![0x5d],                   // EMPTY_LIST
618            vec![0x7d],                   // EMPTY_DICT
619            vec![0x29],                   // EMPTY_TUPLE
620            vec![0x28, b'N', b'N', b'l'], // MARK..LIST
621            vec![0x28, b'N', b'N', b't'], // MARK..TUPLE
622            vec![0x28, b'N', b'N', b'd'], // MARK..DICT
623            vec![b'N', b'N', b'a'],       // APPEND
624            vec![0x28, b'N', b'N', b'e'], // APPENDS
625            vec![b'N', b'N', b's'],       // SETITEM
626            vec![0x28, b'N', b'N', b'u'], // SETITEMS
627            vec![b'N', 0x94],             // MEMOIZE
628            vec![b'N', b'q', 0x05],       // BINPUT
629            vec![b'N', b'r', 1, 0, 0, 0], // LONG_BINPUT
630            vec![b'N', 0x94, b'h', 0x00], // BINGET (memo 0 seeded?)
631            vec![b'g', b'0', 0x0a],       // GET (empty memo blocked)
632        ];
633        for case in cases {
634            let text = stream(&case);
635            // The verdict itself (clean vs blocked) is opcode-defined; both are
636            // fine here as long as the walker terminates deterministically.
637            let _ = pickle_suffix_reaches_reduce_or_build(&text);
638        }
639        // Precise verdicts for the memo family: MEMOIZE then BINGET hits.
640        let hit = stream(&[b'N', 0x94, b'N', b'q', 0x00, b'h', 0x00, b'R']);
641        assert!(pickle_suffix_reaches_reduce_or_build(&hit));
642        // BINPUT on an empty stack is blocked.
643        let empty_stack = stream(&[b'q', 0x00]);
644        assert!(!pickle_suffix_reaches_reduce_or_build(&empty_stack));
645        // LONG_BINPUT on an empty stack is blocked.
646        let long_empty = stream(&[b'r', 0, 0, 0, 0]);
647        assert!(!pickle_suffix_reaches_reduce_or_build(&long_empty));
648        // MEMOIZE on an empty stack is blocked.
649        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[0x94])));
650        // BINGET with a memo miss is blocked; with a memo hit walks clean.
651        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[
652            b'h', 0x07
653        ])));
654        let seeded = stream(&[b'N', b'q', 0x03, b'h', 0x03, b'R']);
655        assert!(pickle_suffix_reaches_reduce_or_build(&seeded));
656        // LONG_BINGET with a hit walks clean.
657        let long_hit = stream(&[b'N', b'q', 0x03, b'j', 3, 0, 0, 0, b'R']);
658        assert!(pickle_suffix_reaches_reduce_or_build(&long_hit));
659        // GET with a non-numeric index is blocked.
660        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[
661            b'g', b'x', 0x0a
662        ])));
663    }
664
665    #[test]
666    fn class_resolution_and_unknown_opcodes_are_blocked() {
667        for key in [
668            0x63u8, 0x69, 0x6f, 0x81, 0x82, 0x93, 0x50, 0x51, 0x84, 0x85, 0x86, 0x01,
669        ] {
670            assert!(
671                !pickle_suffix_reaches_reduce_or_build(&stream(&[key])),
672                "opcode {key:#x} must be blocked"
673            );
674        }
675    }
676
677    #[test]
678    fn stack_underflows_are_blocked() {
679        // POP, DUP, APPEND, SETITEM on empty stacks.
680        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[0x30])));
681        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[0x31])));
682        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[0x32])));
683        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(b"a")));
684        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(b"s")));
685        // APPENDS/SETITEMS without a mark.
686        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(b"e")));
687        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(b"u")));
688        // LIST/TUPLE/DICT without a mark.
689        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(b"l")));
690    }
691
692    #[test]
693    fn non_byte_safe_input_is_rejected() {
694        // A char outside ASCII and the surrogate escape range fails the
695        // byte-safe conversion.
696        assert!(!pickle_prefix_is_opcode_stream("\u{5000}"));
697        assert!(!pickle_suffix_reaches_reduce_or_build("\u{5000}"));
698    }
699
700    #[test]
701    fn surrogate_escapes_map_back_to_high_bytes() {
702        // 0xdc80..=0xdcff ride back to 0x80..=0xff: a FRAME opcode written
703        // through the escape walks exactly like its direct encoding.
704        let direct =
705            pickle_suffix_reaches_reduce_or_build(&stream(&[0x95, 0, 0, 0, 0, 0, 0, 0, 0, b'R']));
706        assert!(direct);
707        let mut escape_bytes: Vec<u8> = vec![0x95, 0, 0, 0, 0, 0, 0, 0, 0];
708        escape_bytes.push(b'R');
709        let via_escape = pickle_suffix_reaches_reduce_or_build(&stream(&escape_bytes));
710        assert!(via_escape);
711    }
712
713    #[test]
714    fn oversize_windows_budget_their_work() {
715        // Longer than the 4096-byte budget: the window truncates and the
716        // incomplete walk is tolerated.
717        let big = vec![0x4bu8; 5000];
718        assert!(pickle_prefix_is_opcode_stream(&stream(&big)));
719        assert!(pickle_suffix_reaches_reduce_or_build(&stream(&big)));
720    }
721
722    #[test]
723    fn the_global_candidate_detector_needs_prefix_and_suffix() {
724        let candidate = super::super::pyregex::Candidate { start: 0, end: 3 };
725        // Prefix "abc" is not an opcode stream start ('a' APPEND on empty).
726        assert!(!pickle_global_candidate_is_injection(
727            "abccos\nsystem\n(R",
728            candidate,
729            9
730        ));
731        // A real pickle prefix plus a REDUCE-terminated suffix.
732        let mut prefix = vec![0x80u8, 0x04, 0x95];
733        prefix.extend_from_slice(&[0x00; 8]);
734        prefix.extend_from_slice(&[0x8c, 0x04, b'm', b'a', b'i', b'n']);
735        let mut haystack = stream(&prefix);
736        haystack.push_str("cos\nsystem\nR");
737        let candidate = super::super::pyregex::Candidate {
738            start: haystack.find("cos").expect("candidate"),
739            end: haystack.find("system").expect("group"),
740        };
741        let group_one_end = haystack.rfind('\n').expect("newline") + 1;
742        assert!(pickle_global_candidate_is_injection(
743            &haystack,
744            candidate,
745            group_one_end
746        ));
747    }
748}
749
750#[cfg(test)]
751mod coverage_tests {
752    use super::*;
753    use crate::patterns::pyregex::Candidate;
754
755    #[test]
756    fn opcode_stream_walks_reject_blocked_and_short_streams() {
757        // a memoize on an empty stack blocks the walk: not an opcode stream
758        assert!(!pickle_prefix_is_opcode_stream("\u{94}"));
759        // a binput on an empty stack blocks too
760        assert!(!pickle_prefix_is_opcode_stream("\u{71}a"));
761        // a long binput on an empty stack blocks
762        assert!(!pickle_prefix_is_opcode_stream("\u{72}aaaa"));
763        // a binget with an unmemoized index blocks
764        assert!(!pickle_prefix_is_opcode_stream("\u{68}a"));
765        // a truncated frame header is tolerated on an incomplete prefix
766        assert!(pickle_prefix_is_opcode_stream("\u{95}04"));
767        // a POP on an empty stack, a DUP on an empty stack, a GLOBAL class
768        // resolution, and an unknown opcode all block
769        assert!(!pickle_prefix_is_opcode_stream("\u{30}"));
770        assert!(!pickle_prefix_is_opcode_stream("\u{32}"));
771        // GLOBAL 'c' + module + name + '(' blocks on the class resolution
772        assert!(!pickle_prefix_is_opcode_stream("cos\nfoo\n("));
773        assert!(!pickle_prefix_is_opcode_stream("\u{ff}"));
774        assert!(pickle_prefix_is_opcode_stream("g"));
775        // a LONG_BINGET with an unmemoized index blocks
776        assert!(!pickle_prefix_is_opcode_stream(
777            "\u{6a}\u{0}\u{0}\u{0}\u{0}"
778        ));
779    }
780
781    #[test]
782    fn suffix_walk_stops_at_reduce_or_build() {
783        // a reduce opcode in the suffix reaches the stop point
784        assert!(pickle_suffix_reaches_reduce_or_build("\u{72}"));
785        // a blocked suffix stream is not injection
786        assert!(!pickle_suffix_reaches_reduce_or_build("\u{94}"));
787        // a short-read suffix is tolerated like the reference
788        assert!(pickle_suffix_reaches_reduce_or_build("\u{95}04"));
789        // a non-byte-safe suffix rejects outright
790        assert!(!pickle_suffix_reaches_reduce_or_build("\u{1234}"));
791    }
792
793    #[test]
794    fn non_byte_safe_windows_reject_the_candidate() {
795        // a char outside the byte range and outside the surrogateescape map
796        assert!(!pickle_prefix_is_opcode_stream("\u{1234}"));
797        assert!(!pickle_suffix_reaches_reduce_or_build("\u{1234}"));
798    }
799
800    #[test]
801    fn global_candidate_gate_composes_prefix_and_suffix() {
802        let haystack = "\u{80}04cos\nsystem\n].";
803        // the prefix walks as opcodes; the suffix cannot reach REDUCE
804        assert!(!pickle_global_candidate_is_injection(
805            haystack,
806            Candidate::new(0, 4),
807            4
808        ));
809        // prefix walks, suffix starts with a REDUCE opcode
810        let haystack = "cos\n]\u{72}tR";
811        assert!(pickle_global_candidate_is_injection(
812            haystack,
813            Candidate::new(0, 4),
814            4
815        ));
816        // a non-walkable prefix vetoes the candidate outright
817        assert!(!pickle_global_candidate_is_injection(
818            "\u{94}cos\n]\u{72}tR",
819            Candidate::new(0, 4),
820            4
821        ));
822    }
823}
824
825#[cfg(test)]
826mod gap_tests {
827    use super::*;
828    use crate::patterns::pyregex::Candidate;
829
830    /// Builds the char string the public functions expect: Latin-1 bytes map
831    /// straight back through `window_from_chars`.
832    fn stream(bytes: &[u8]) -> String {
833        bytes.iter().map(|b| char::from(*b)).collect()
834    }
835
836    #[test]
837    fn short_read_operand_arms_are_tolerated_like_the_reference() {
838        // Every fixed-size and length-prefixed read reports ShortRead when
839        // the window ends mid-operand; the walker tolerates it.
840        let cases: Vec<Vec<u8>> = vec![
841            vec![0x4b],                         // BININT1
842            vec![0x4d, 1],                      // BININT2
843            vec![0x47, 1, 2, 3],                // BINFLOAT
844            vec![b'I', b'1'],                   // INT line
845            vec![b'L', b'1'],                   // LONG line
846            vec![b'F', b'1'],                   // FLOAT line
847            vec![b'S', b'\'', b'a'],            // STRING line
848            vec![b'V', b'x'],                   // UNICODE line
849            vec![0x58, 9, 0, 0, 0, b'a'],       // BINUNICODE payload
850            vec![0x8c, 9, b'a'],                // SHORT_BINUNICODE payload
851            vec![0x54, 9, 0, 0, 0, b'a'],       // BINSTRING payload
852            vec![0x55, 9, b'a'],                // SHORT_BINSTRING payload
853            vec![0x42, 9, 0, 0, 0, b'a'],       // BINBYTES payload
854            vec![0x8e, 9, 0, 0, 0, 0, 0, 0, 0], // BINBYTES8 payload
855            vec![0x43, 9, b'a'],                // SHORT_BINBYTES payload
856            vec![0x96, 9, 0, 0, 0, 0, 0, 0, 0], // BYTEARRAY8 payload
857            vec![0x8a, 9, b'a'],                // 1-byte length payload
858            vec![0x8b, 9, 0, 0, 0, b'a'],       // 4-byte length payload
859            vec![0x80],                         // PROTO level byte
860            vec![0x71],                         // BINPUT index byte
861            vec![0x68],                         // BINGET index byte
862            vec![0x6a, 1, 2],                   // LONG_BINGET index
863        ];
864        // The length read itself can also run out of input.
865        for key in [0x58u8, 0x8c, 0x54, 0x55, 0x42, 0x8e, 0x43, 0x96, 0x8a, 0x8b] {
866            assert!(
867                pickle_suffix_reaches_reduce_or_build(&stream(&[key])),
868                "a missing length operand for {key:#x} is tolerated"
869            );
870        }
871        for case in cases {
872            let text = stream(&case);
873            let label = String::from_utf8_lossy(&case);
874            assert!(
875                pickle_suffix_reaches_reduce_or_build(&text),
876                "a truncated {label} operand is tolerated"
877            );
878        }
879    }
880
881    #[test]
882    fn blocked_opcode_arms_report_the_reference_verdicts() {
883        // SETITEM on a one-deep stack fails its second pop; on a completely
884        // empty stack (the prefix walk) it fails the first.
885        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(b"Ns")));
886        assert!(!pickle_prefix_is_opcode_stream("\u{73}"));
887        // A FLOAT line that is not valid UTF-8 is blocked.
888        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[
889            b'F', 0xFF, 0x0a
890        ])));
891        // A GET index line that is not valid UTF-8 is blocked.
892        assert!(!pickle_suffix_reaches_reduce_or_build(&stream(&[
893            b'g', 0xFF, 0x0a
894        ])));
895    }
896
897    #[test]
898    fn get_reads_a_memo_entry_seeded_by_binput() {
899        // NONE, BINPUT 0, GET "0", REDUCE: the text GET hits the memo entry
900        // the binary BINPUT seeded
901        let text = stream(&[b'N', b'q', 0x00, b'g', b'0', 0x0a, b'R']);
902        assert!(pickle_suffix_reaches_reduce_or_build(&text));
903    }
904
905    #[test]
906    fn a_non_stream_prefix_rejects_the_global_candidate() {
907        // "hello" walks as BINGET 'e' with an empty memo: blocked
908        assert!(!pickle_global_candidate_is_injection(
909            "helloX",
910            Candidate::new(5, 6),
911            6
912        ));
913    }
914}
guard_core_engine/patterns/pickle.rs

guard_core_engine/patterns/
pickle.rs
