From ba210ce18fbce878afc163e937f1846b65242e3a Mon Sep 17 00:00:00 2001 From: Michael D'Angelo Date: Mon, 8 Jun 2026 22:39:05 -0700 Subject: [PATCH 1/4] fix(picklescan): close encoded protocol0 probe gaps Detect base64 protocol 0 payloads that begin with long scalar operands, and continue bounded probing after lenient prefixes unless the entire literal decodes to exactly one pickle. --- packages/modelaudit-picklescan/CHANGELOG.md | 1 + .../modelaudit-picklescan/rust/src/nested.rs | 157 +++++++++++++++++- .../tests/test_protocol0_line_operands.py | 70 ++++++++ 3 files changed, 219 insertions(+), 9 deletions(-) diff --git a/packages/modelaudit-picklescan/CHANGELOG.md b/packages/modelaudit-picklescan/CHANGELOG.md index 4cf5448a9..3ee57e9db 100644 --- a/packages/modelaudit-picklescan/CHANGELOG.md +++ b/packages/modelaudit-picklescan/CHANGELOG.md @@ -78,6 +78,7 @@ and this package adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Bug Fixes +- continue probing encoded protocol 0 payloads after long scalar operands and lenient base64 prefixes - preserve fail-closed nested protocol 0 analysis after `INST` opcodes - fail closed when known-size pickle boundaries cannot be verified or contain trailing bytes, and bind file scans to one descriptor - fail closed before copying protocol 0 line operands larger than 8 MiB diff --git a/packages/modelaudit-picklescan/rust/src/nested.rs b/packages/modelaudit-picklescan/rust/src/nested.rs index b5300424a..0e3d194dd 100644 --- a/packages/modelaudit-picklescan/rust/src/nested.rs +++ b/packages/modelaudit-picklescan/rust/src/nested.rs @@ -826,7 +826,12 @@ fn encoded_pickle_kind_at(value: &str, index: usize) -> Option<&'static str> { } let probe = take_bytes_str_slice(suffix, MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES); - if starts_base64_pickle && base64_prefix_has_pickle_prefix(probe) { + let base64_has_pickle_prefix = if starts_base64_token_at(value, index) { + base64_prefix_has_pickle_prefix(probe) + } else { + base64_prefix_has_nested_probe_prefix(probe, MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES) + }; + if starts_base64_pickle && base64_has_pickle_prefix { return Some("base64"); } if starts_hex_pickle && hex_prefix_has_pickle_prefix(probe) { @@ -851,11 +856,21 @@ fn encoded_pickle_prefix_probe_limit_kind_at(value: &str, index: usize) -> Optio } let probe = take_bytes_str_slice(suffix, MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES); - if starts_base64_pickle - && !base64_prefix_has_pickle_prefix(probe) - && base64_prefix_has_pickle_prefix_with_limit(suffix, suffix.len()) - { - return Some("base64"); + let starts_base64_token = starts_base64_token_at(value, index); + let probe_has_base64_pickle = if starts_base64_token { + base64_prefix_has_pickle_prefix(probe) + } else { + base64_prefix_has_nested_probe_prefix(probe, MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES) + }; + if starts_base64_pickle && !probe_has_base64_pickle { + let suffix_has_base64_pickle = if starts_base64_token { + base64_prefix_has_pickle_prefix_with_limit(suffix, suffix.len()) + } else { + base64_prefix_has_nested_probe_prefix(suffix, suffix.len()) + }; + if suffix_has_base64_pickle { + return Some("base64"); + } } if starts_hex_pickle @@ -888,9 +903,25 @@ fn encoded_probe_prefix_consumes_literal( prefix_has_base64_pickle: bool, prefix_has_hex_pickle: bool, ) -> bool { - (prefix_has_base64_pickle && is_lenient_base64_candidate(value)) - || (prefix_has_hex_pickle - && take_hex_literal_prefix(value, value.len()).len() == value.len()) + encoded_prefix_consumes_literal(value, prefix_has_base64_pickle, prefix_has_hex_pickle) + || (prefix_has_base64_pickle && lenient_base64_decodes_to_single_pickle(value)) +} + +fn lenient_base64_decodes_to_single_pickle(value: &str) -> bool { + is_lenient_base64_candidate(value) + && canonical_base64_candidate(&normalize_base64_literal( + value, + MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES, + value.len(), + )) + .as_deref() + .and_then(decode_base64) + .is_some_and(|decoded| { + pickle_payload_extent_result(&decoded, decoded.len()) + .ok() + .flatten() + == Some(decoded.len()) + }) } fn encoded_literal_candidates(stripped: &str) -> Vec { @@ -1068,6 +1099,11 @@ fn base64_prefix_has_pickle_prefix(value: &str) -> bool { } fn base64_prefix_has_pickle_prefix_with_limit(value: &str, max_input_bytes: usize) -> bool { + base64_prefix_has_nested_probe_prefix(value, max_input_bytes) + || base64_prefix_has_complete_protocol0_scalar(value, max_input_bytes) +} + +fn base64_prefix_has_nested_probe_prefix(value: &str, max_input_bytes: usize) -> bool { base64_literal_candidates(value, ENCODED_LITERAL_PROBE_CHARS, max_input_bytes) .into_iter() .any(|prefix| { @@ -1078,6 +1114,40 @@ fn base64_prefix_has_pickle_prefix_with_limit(value: &str, max_input_bytes: usiz }) } +fn base64_prefix_has_complete_protocol0_scalar(value: &str, max_input_bytes: usize) -> bool { + let strict_limit = max_input_bytes.min(MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES); + let strict = take_base64_literal_prefix(value, strict_limit); + if strict.len() <= ENCODED_LITERAL_PROBE_CHARS + || !encoded_base64_first_byte(strict.as_bytes()) + .is_some_and(|byte| matches!(byte, b'I' | b'S' | b'V')) + { + return false; + } + + canonical_base64_candidate(strict) + .as_deref() + .and_then(decode_base64) + .is_some_and(|decoded| { + decoded + .first() + .is_some_and(|byte| matches!(*byte, b'I' | b'S' | b'V')) + && pickle_payload_extent_result(&decoded, decoded.len()) + .ok() + .flatten() + .is_some() + }) +} + +fn starts_base64_token_at(value: &str, index: usize) -> bool { + index == 0 + || value.as_bytes().get(index - 1).is_some_and(|byte| { + !matches!( + *byte, + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'+' | b'/' | b'=' + ) + }) +} + fn hex_prefix_has_pickle_prefix(value: &str) -> bool { hex_prefix_has_pickle_prefix_with_limit(value, MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES) } @@ -1395,6 +1465,34 @@ mod tests { const TEST_MAX_NESTED_PICKLE_BYTES: usize = 2 * 1024 * 1024; + fn encode_base64_for_test(value: &[u8]) -> String { + const ALPHABET: &[u8; 64] = + b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + let mut encoded = String::with_capacity(value.len().div_ceil(3) * 4); + for chunk in value.chunks(3) { + let first = chunk[0]; + let second = chunk.get(1).copied().unwrap_or(0); + let third = chunk.get(2).copied().unwrap_or(0); + encoded.push(char::from(ALPHABET[(first >> 2) as usize])); + encoded.push(char::from( + ALPHABET[(((first & 0x03) << 4) | (second >> 4)) as usize], + )); + if chunk.len() > 1 { + encoded.push(char::from( + ALPHABET[(((second & 0x0f) << 2) | (third >> 6)) as usize], + )); + } else { + encoded.push('='); + } + if chunk.len() > 2 { + encoded.push(char::from(ALPHABET[(third & 0x3f) as usize])); + } else { + encoded.push('='); + } + } + encoded + } + #[test] fn encoded_prefix_gates_recognize_pickle_prefixes() { assert!(base64_prefix_has_pickle_prefix("gAR9Lg==")); @@ -1657,6 +1755,47 @@ mod tests { .any(|window| window.starts_with("Y29zCnN5c3RlbQopUi4="))); } + #[test] + fn base64_prefix_gate_recognizes_complete_long_protocol0_scalar_pickle() { + for mut payload in [b"I".to_vec(), b"S'".to_vec(), b"V".to_vec()] { + let fill = if payload.starts_with(b"I") { + b'1' + } else { + b'A' + }; + payload.extend(std::iter::repeat_n( + fill, + PROTOCOL0_SCALAR_PREFIX_PROBE_BYTES + 1, + )); + if payload.starts_with(b"S") { + payload.push(b'\''); + } + payload.extend_from_slice(b"\n0cos\nsystem\n(S'id'\ntR."); + let encoded = encode_base64_for_test(&payload); + + assert!(base64_prefix_has_pickle_prefix(&encoded)); + assert!( + decode_possible_encoded_pickle(&encoded, TEST_MAX_NESTED_PICKLE_BYTES) + .iter() + .any(|candidate| candidate.payload == payload) + ); + } + } + + #[test] + fn encoded_probe_windows_continue_after_lenient_benign_prefix() { + let benign = encode_base64_for_test(b"I42\n.") + .trim_end_matches('=') + .to_string(); + let malicious = encode_base64_for_test(b"cos\nsystem\n)R.") + .trim_end_matches('=') + .to_string(); + let value = format!("{benign}!{malicious}"); + let windows = encoded_nested_literal_probe_windows(&value, value.len()); + + assert!(windows.iter().any(|window| window.starts_with(&malicious))); + } + #[test] fn wrapped_base64_nested_literals_ignore_comment_leaders() { let value = "# this is doc\n# Y29zCnN5\n# c3RlbQopUi4=\n# more"; diff --git a/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py b/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py index 9386d4ec4..6d4ee8aec 100644 --- a/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py +++ b/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py @@ -23,6 +23,16 @@ def _nested_overlong_protocol0_line_operand() -> bytes: return b"cos\nsystem\n(S'" + _overlong_protocol0_operand_body() + b"'\ntR." +def _long_scalar_before_reduce_protocol0_pickle(opcode: bytes) -> bytes: + if opcode == b"I": + scalar = b"I" + (b"1" * 257) + b"\n" + elif opcode == b"S": + scalar = b"S'" + (b"A" * 257) + b"'\n" + else: + scalar = b"V" + (b"A" * 257) + b"\n" + return scalar + b"0cos\nsystem\n(S'id'\ntR." + + def test_scan_bytes_accepts_exact_limit_protocol0_line_operand() -> None: payload = b"S'" + (b"A" * (MAX_PROTOCOL0_LINE_OPERAND_BYTES - 2)) + b"'\n." @@ -134,6 +144,66 @@ def test_scan_bytes_detects_inline_encoded_protocol0_pickle_before_suffix( ) +@pytest.mark.parametrize("opcode", [b"I", b"S", b"V"]) +def test_scan_bytes_detects_base64_pickle_after_long_protocol0_scalar(opcode: bytes) -> None: + nested_payload = _long_scalar_before_reduce_protocol0_pickle(opcode) + encoded = base64.b64encode(nested_payload).decode("ascii") + + report = scan_bytes( + pickle.dumps(encoded, protocol=4), + source="base64-long-scalar-before-reduce.pkl", + ) + + assert report.status == ScanStatus.COMPLETE + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.rule_code == "S601" + and finding.details.get("encoding") == "base64" + and finding.details.get("nested_has_execution_opcode") is True + for finding in report.findings + ) + assert any( + finding.rule_code == "DANGEROUS_CALL" + and finding.details.get("module") == "os" + and finding.details.get("name") == "system" + for finding in report.findings + ) + + +def test_scan_bytes_detects_later_base64_pickle_after_lenient_benign_prefix() -> None: + benign = base64.b64encode(b"I42\n.").decode("ascii").rstrip("=") + malicious = base64.b64encode(b"cos\nsystem\n)R.").decode("ascii").rstrip("=") + + report = scan_bytes( + pickle.dumps(f"{benign}!{malicious}", protocol=4), + source="base64-later-pickle-after-lenient-prefix.pkl", + ) + + assert report.status == ScanStatus.COMPLETE + assert report.verdict == SafetyVerdict.MALICIOUS + assert any( + finding.rule_code == "S601" + and finding.details.get("encoding") == "base64" + and finding.details.get("nested_has_execution_opcode") is True + for finding in report.findings + ) + assert any(finding.rule_code == "DANGEROUS_CALL" for finding in report.findings) + + +def test_scan_bytes_keeps_multiple_lenient_benign_base64_pickles_clean() -> None: + first = base64.b64encode(b"I42\n.").decode("ascii").rstrip("=") + second = base64.b64encode(b"S'ok'\n.").decode("ascii").rstrip("=") + + report = scan_bytes( + pickle.dumps(f"{first}!{second}", protocol=4), + source="multiple-lenient-benign-base64-pickles.pkl", + ) + + assert report.status == ScanStatus.COMPLETE + assert report.verdict == SafetyVerdict.CLEAN + assert report.findings == () + + @pytest.mark.parametrize( ("encoding", "expected_rule_code"), [("base64", "S601"), ("hex", "S602")], From c01240a6aea9690fccf34e39dfe57535a26ec60d Mon Sep 17 00:00:00 2001 From: Michael D'Angelo Date: Tue, 9 Jun 2026 01:28:35 -0700 Subject: [PATCH 2/4] fix(picklescan): align encoded probe coverage bounds --- .../modelaudit-picklescan/rust/src/nested.rs | 97 +++++++++++++++---- .../modelaudit-picklescan/rust/src/state.rs | 6 +- .../tests/test_protocol0_line_operands.py | 15 +++ 3 files changed, 97 insertions(+), 21 deletions(-) diff --git a/packages/modelaudit-picklescan/rust/src/nested.rs b/packages/modelaudit-picklescan/rust/src/nested.rs index d096c0a40..f8aa53f39 100644 --- a/packages/modelaudit-picklescan/rust/src/nested.rs +++ b/packages/modelaudit-picklescan/rust/src/nested.rs @@ -714,6 +714,8 @@ pub(crate) fn encoded_nested_literal_probe_windows_with_limit( value, prefix_has_base64_pickle, prefix_has_hex_pickle, + MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES, + max_window_chars, ) { return EncodedNestedProbeWindows { windows, @@ -780,7 +782,11 @@ pub(crate) fn encoded_nested_literal_probe_windows_with_limit( } } -pub(crate) fn encoded_nested_literal_probe_coverage_incomplete(value: &str) -> bool { +pub(crate) fn encoded_nested_literal_probe_coverage_incomplete( + value: &str, + max_window_chars: usize, + max_nested_pickle_bytes: usize, +) -> bool { if value.len() <= MAX_ENCODED_LITERAL_MID_SCAN_BYTES { return false; } @@ -794,6 +800,8 @@ pub(crate) fn encoded_nested_literal_probe_coverage_incomplete(value: &str) -> b stripped, prefix_has_base64_pickle, prefix_has_hex_pickle, + max_nested_pickle_bytes.div_ceil(3) * 4, + max_window_chars, ) } @@ -832,8 +840,8 @@ fn embedded_long_protocol0_base64_probe_windows( max_window_chars: usize, ) -> EncodedNestedProbeWindows { let max_input_bytes = MAX_ENCODED_LITERAL_MID_SCAN_BYTES - .saturating_add(max_window_chars.max(LONG_PROTOCOL0_SCALAR_PROBE_CHARS)); - let input_truncated = value.len() > max_input_bytes; + .saturating_add(max_window_chars.max(LONG_PROTOCOL0_SCALAR_PROBE_CHARS)) + .saturating_add(ENCODED_LITERAL_PROBE_CHARS); let mut compact = String::with_capacity(value.len().min(max_input_bytes)); let mut raw_positions = Vec::with_capacity(value.len().min(max_input_bytes)); for (raw_index, byte) in value.bytes().take(max_input_bytes).enumerate() { @@ -878,13 +886,8 @@ fn embedded_long_protocol0_base64_probe_windows( .is_some_and(|raw_index| *raw_index < MAX_ENCODED_LITERAL_MID_SCAN_BYTES); let remaining_line = &decoded[decoded_index + LONG_PROTOCOL0_SCALAR_PROBE_BYTES..]; let Some(newline_offset) = remaining_line.iter().position(|byte| *byte == b'\n') else { - if input_truncated { - return EncodedNestedProbeWindows { - windows, - limit_exceeded: true, - limit_exceeded_encoding: Some("base64"), - }; - } + // A scalar prefix alone is not enough evidence of a pickle. Without the + // terminating line, no later opcode can establish structured content. break; }; let probe = &decoded[decoded_index..]; @@ -1026,18 +1029,28 @@ fn encoded_probe_prefix_consumes_literal( value: &str, prefix_has_base64_pickle: bool, prefix_has_hex_pickle: bool, + max_encoded_chars: usize, + max_window_chars: usize, ) -> bool { encoded_prefix_consumes_literal(value, prefix_has_base64_pickle, prefix_has_hex_pickle) - || (prefix_has_base64_pickle && lenient_base64_decodes_to_single_pickle(value)) + || (prefix_has_base64_pickle + && lenient_base64_decodes_to_single_pickle(value, max_encoded_chars, max_window_chars)) } -fn lenient_base64_decodes_to_single_pickle(value: &str) -> bool { - is_lenient_base64_candidate(value) - && canonical_base64_candidate(&normalize_base64_literal( - value, - MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES, - value.len(), - )) +fn lenient_base64_decodes_to_single_pickle( + value: &str, + max_encoded_chars: usize, + max_input_bytes: usize, +) -> bool { + if value.len() > max_input_bytes || !is_lenient_base64_candidate(value) { + return false; + } + let normalized = + normalize_base64_literal(value, max_encoded_chars.saturating_add(1), value.len()); + if normalized.len() > max_encoded_chars { + return false; + } + canonical_base64_candidate(&normalized) .as_deref() .and_then(decode_base64) .is_some_and(|decoded| { @@ -1856,7 +1869,9 @@ mod tests { .iter() .any(|window| window.starts_with("gAR9Lg=="))); assert!(encoded_nested_literal_probe_coverage_incomplete( - &beyond_bound + &beyond_bound, + 64, + 64, )); } @@ -1865,7 +1880,9 @@ mod tests { let whole_literal = format!("gAR9{}\n", "A".repeat(MAX_ENCODED_LITERAL_MID_SCAN_BYTES)); assert!(!encoded_nested_literal_probe_coverage_incomplete( - &whole_literal + &whole_literal, + whole_literal.len(), + whole_literal.len(), )); } @@ -1963,6 +1980,46 @@ mod tests { assert!(windows.iter().any(|window| window.starts_with(&malicious))); } + #[test] + fn encoded_probe_coverage_accepts_large_lenient_single_pickle() { + let body_len = (MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES * 3 / 4) + 1024; + let mut payload = b"\x80\x04X".to_vec(); + payload.extend_from_slice(&(body_len as u32).to_le_bytes()); + payload.extend(std::iter::repeat_n(b'A', body_len)); + payload.push(b'.'); + let encoded = encode_base64_for_test(&payload); + let separated = encoded + .as_bytes() + .chunks(76) + .map(|chunk| std::str::from_utf8(chunk).unwrap()) + .collect::>() + .join(" "); + + assert!(separated.len() > MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES); + assert!(!encoded_nested_literal_probe_coverage_incomplete( + &separated, + separated.len(), + payload.len(), + )); + } + + #[test] + fn long_protocol0_base64_probe_ignores_unterminated_scalar_without_structure() { + let body_len = (MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES * 3 / 4) + 1024; + let mut payload = b"V".to_vec(); + payload.extend(std::iter::repeat_n(b'A', body_len)); + let encoded = encode_base64_for_test(&payload); + + let probes = embedded_long_protocol0_base64_probe_windows( + &encoded, + LONG_PROTOCOL0_SCALAR_PROBE_CHARS, + ); + + assert!(encoded.len() > MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES); + assert!(probes.windows.is_empty()); + assert!(!probes.limit_exceeded); + } + #[test] fn wrapped_base64_nested_literals_ignore_comment_leaders() { let value = "# this is doc\n# Y29zCnN5\n# c3RlbQopUi4=\n# more"; diff --git a/packages/modelaudit-picklescan/rust/src/state.rs b/packages/modelaudit-picklescan/rust/src/state.rs index 5099e8eb2..433c9b7fa 100644 --- a/packages/modelaudit-picklescan/rust/src/state.rs +++ b/packages/modelaudit-picklescan/rust/src/state.rs @@ -6016,7 +6016,11 @@ impl<'a> ScanState<'a> { let max_window_chars = encoded_nested_window_char_limit(value, self.options.max_nested_pickle_bytes); - let probe_coverage_incomplete = encoded_nested_literal_probe_coverage_incomplete(value); + let probe_coverage_incomplete = encoded_nested_literal_probe_coverage_incomplete( + value, + max_window_chars, + self.options.max_nested_pickle_bytes, + ); if !found_candidate && (value.len() <= max_window_chars || value.chars().count() <= max_window_chars) { diff --git a/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py b/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py index d080543e5..988059b92 100644 --- a/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py +++ b/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py @@ -207,6 +207,21 @@ def test_scan_bytes_detects_lenient_base64_pickle_after_long_protocol0_scalar( assert any(finding.rule_code == "DANGEROUS_CALL" for finding in report.findings) +def test_scan_bytes_ignores_large_unterminated_protocol0_base64_scalar() -> None: + encoded_scalar = base64.b64encode(b"V" + (b"A" * 257)).decode("ascii") + encoded = encoded_scalar + ("!" * ((1024 * 1024) + 512)) + + report = scan_bytes( + pickle.dumps(encoded, protocol=4), + source="unterminated-protocol0-base64-scalar.pkl", + ) + + assert report.status == ScanStatus.COMPLETE + assert report.verdict == SafetyVerdict.CLEAN + assert report.findings == () + assert not any(notice.code == "nested_probe_limit_exceeded" for notice in report.notices) + + @pytest.mark.parametrize("opcode", [b"I", b"S", b"V"]) def test_scan_bytes_keeps_wrapped_benign_long_scalar_base64_pickle_clean(opcode: bytes) -> None: nested_payload = _long_scalar_before_reduce_protocol0_pickle(opcode).split(b"0cos", maxsplit=1)[0] + b"." From f365c0a75d6b2ffa7762a967c1b087b42fd53eb7 Mon Sep 17 00:00:00 2001 From: mldangelo-oai Date: Tue, 9 Jun 2026 02:02:13 -0700 Subject: [PATCH 3/4] fix(picklescan): cover sparse encoded protocol 0 operands --- .../modelaudit-picklescan/rust/src/nested.rs | 22 +++++++++---- .../tests/test_protocol0_line_operands.py | 32 +++++++++++++++++++ 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/packages/modelaudit-picklescan/rust/src/nested.rs b/packages/modelaudit-picklescan/rust/src/nested.rs index f8aa53f39..1fc2385b9 100644 --- a/packages/modelaudit-picklescan/rust/src/nested.rs +++ b/packages/modelaudit-picklescan/rust/src/nested.rs @@ -839,12 +839,20 @@ fn embedded_long_protocol0_base64_probe_windows( value: &str, max_window_chars: usize, ) -> EncodedNestedProbeWindows { - let max_input_bytes = MAX_ENCODED_LITERAL_MID_SCAN_BYTES - .saturating_add(max_window_chars.max(LONG_PROTOCOL0_SCALAR_PROBE_CHARS)) + let tail_compact_chars = max_window_chars + .max(LONG_PROTOCOL0_SCALAR_PROBE_CHARS) .saturating_add(ENCODED_LITERAL_PROBE_CHARS); - let mut compact = String::with_capacity(value.len().min(max_input_bytes)); - let mut raw_positions = Vec::with_capacity(value.len().min(max_input_bytes)); - for (raw_index, byte) in value.bytes().take(max_input_bytes).enumerate() { + let max_compact_chars = MAX_ENCODED_LITERAL_MID_SCAN_BYTES.saturating_add(tail_compact_chars); + let mut compact = String::with_capacity(value.len().min(max_compact_chars)); + let mut raw_positions = Vec::with_capacity(value.len().min(max_compact_chars)); + let mut compact_chars_at_mid_scan = None; + for (raw_index, byte) in value.bytes().enumerate() { + if raw_index >= MAX_ENCODED_LITERAL_MID_SCAN_BYTES { + let compact_chars_at_mid_scan = *compact_chars_at_mid_scan.get_or_insert(compact.len()); + if compact.len().saturating_sub(compact_chars_at_mid_scan) >= tail_compact_chars { + break; + } + } if base64_value(byte).is_some() { compact.push(char::from(byte)); raw_positions.push(raw_index); @@ -1032,8 +1040,10 @@ fn encoded_probe_prefix_consumes_literal( max_encoded_chars: usize, max_window_chars: usize, ) -> bool { + let prefix_may_be_long_protocol0_scalar = encoded_base64_first_byte_unbounded(value.as_bytes()) + .is_some_and(|byte| matches!(byte, b'I' | b'S' | b'V')); encoded_prefix_consumes_literal(value, prefix_has_base64_pickle, prefix_has_hex_pickle) - || (prefix_has_base64_pickle + || ((prefix_has_base64_pickle || prefix_may_be_long_protocol0_scalar) && lenient_base64_decodes_to_single_pickle(value, max_encoded_chars, max_window_chars)) } diff --git a/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py b/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py index 988059b92..f10ff05d5 100644 --- a/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py +++ b/packages/modelaudit-picklescan/tests/test_protocol0_line_operands.py @@ -207,6 +207,38 @@ def test_scan_bytes_detects_lenient_base64_pickle_after_long_protocol0_scalar( assert any(finding.rule_code == "DANGEROUS_CALL" for finding in report.findings) +def test_scan_bytes_detects_sparse_lenient_base64_pickle_after_long_protocol0_scalar() -> None: + nested_payload = _long_scalar_before_reduce_protocol0_pickle(b"V") + encoded = base64.b64encode(nested_payload).decode("ascii") + separated = ("!" * 4000).join(encoded) + + report = scan_bytes( + pickle.dumps(separated, protocol=4), + source="sparse-lenient-base64-long-scalar-before-reduce.pkl", + ) + + assert len(separated) > 1024 * 1024 + assert report.status == ScanStatus.COMPLETE + assert report.verdict == SafetyVerdict.MALICIOUS + assert any(finding.rule_code == "DANGEROUS_CALL" for finding in report.findings) + + +def test_scan_bytes_ignores_sparse_unterminated_protocol0_base64_scalar() -> None: + encoded = base64.b64encode(b"V" + (b"A" * 257)).decode("ascii") + separated = ("!" * 4000).join(encoded) + + report = scan_bytes( + pickle.dumps(separated, protocol=4), + source="sparse-unterminated-protocol0-base64-scalar.pkl", + ) + + assert len(separated) > 1024 * 1024 + assert report.status == ScanStatus.COMPLETE + assert report.verdict == SafetyVerdict.CLEAN + assert report.findings == () + assert not any(notice.code == "nested_probe_limit_exceeded" for notice in report.notices) + + def test_scan_bytes_ignores_large_unterminated_protocol0_base64_scalar() -> None: encoded_scalar = base64.b64encode(b"V" + (b"A" * 257)).decode("ascii") encoded = encoded_scalar + ("!" * ((1024 * 1024) + 512)) From 931823bb9456b264f58be07bf1c3df1db06f54a4 Mon Sep 17 00:00:00 2001 From: Michael D'Angelo Date: Tue, 9 Jun 2026 03:26:13 -0700 Subject: [PATCH 4/4] fix(picklescan): fail closed on truncated scalar probes --- .../modelaudit-picklescan/rust/src/nested.rs | 75 ++++++++++++++++++- 1 file changed, 73 insertions(+), 2 deletions(-) diff --git a/packages/modelaudit-picklescan/rust/src/nested.rs b/packages/modelaudit-picklescan/rust/src/nested.rs index 1fc2385b9..17f0719fd 100644 --- a/packages/modelaudit-picklescan/rust/src/nested.rs +++ b/packages/modelaudit-picklescan/rust/src/nested.rs @@ -894,8 +894,18 @@ fn embedded_long_protocol0_base64_probe_windows( .is_some_and(|raw_index| *raw_index < MAX_ENCODED_LITERAL_MID_SCAN_BYTES); let remaining_line = &decoded[decoded_index + LONG_PROTOCOL0_SCALAR_PROBE_BYTES..]; let Some(newline_offset) = remaining_line.iter().position(|byte| *byte == b'\n') else { - // A scalar prefix alone is not enough evidence of a pickle. Without the - // terminating line, no later opcode can establish structured content. + let Some(raw_start) = raw_positions.get(encoded_start).copied() else { + break; + }; + if base64_protocol0_scalar_has_line_terminator(value, raw_start) { + return EncodedNestedProbeWindows { + windows, + limit_exceeded: true, + limit_exceeded_encoding: Some("base64"), + }; + } + // A fully scanned scalar prefix without a terminating line cannot hide + // later pickle opcodes and is safe to treat as an encoded near-match. break; }; let probe = &decoded[decoded_index..]; @@ -938,6 +948,48 @@ fn embedded_long_protocol0_base64_probe_windows( } } +fn base64_protocol0_scalar_has_line_terminator(value: &str, raw_start: usize) -> bool { + let Some(raw_candidate) = value.as_bytes().get(raw_start..) else { + return false; + }; + let mut quartet = [0u8; 4]; + let mut quartet_len = 0usize; + let mut decoded_index = 0usize; + + let mut is_terminator = |byte: u8| { + let result = decoded_index >= LONG_PROTOCOL0_SCALAR_PROBE_BYTES && byte == b'\n'; + decoded_index = decoded_index.saturating_add(1); + result + }; + + for byte in raw_candidate { + let Some(value) = base64_value(*byte) else { + continue; + }; + quartet[quartet_len] = value; + quartet_len += 1; + if quartet_len < quartet.len() { + continue; + } + + for decoded in [ + (quartet[0] << 2) | (quartet[1] >> 4), + (quartet[1] << 4) | (quartet[2] >> 2), + (quartet[2] << 6) | quartet[3], + ] { + if is_terminator(decoded) { + return true; + } + } + quartet_len = 0; + } + + if quartet_len >= 2 && is_terminator((quartet[0] << 2) | (quartet[1] >> 4)) { + return true; + } + quartet_len >= 3 && is_terminator((quartet[1] << 4) | (quartet[2] >> 2)) +} + pub(crate) fn encoded_pickle_consumes_literal(value: &str) -> bool { let stripped = value.trim(); let prefix_has_base64_pickle = base64_prefix_has_pickle_prefix(stripped); @@ -2030,6 +2082,25 @@ mod tests { assert!(!probes.limit_exceeded); } + #[test] + fn long_protocol0_base64_probe_fails_closed_on_terminator_beyond_window() { + let body_len = (MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES * 3 / 4) + 1024; + let mut payload = b"V".to_vec(); + payload.extend(std::iter::repeat_n(b'A', body_len)); + payload.extend_from_slice(b"\n0cos\nsystem\n(S'id'\ntR."); + let encoded = encode_base64_for_test(&payload); + + let probes = embedded_long_protocol0_base64_probe_windows( + &encoded, + LONG_PROTOCOL0_SCALAR_PROBE_CHARS, + ); + + assert!(encoded.len() > MAX_ENCODED_LITERAL_PREFIX_SCAN_BYTES); + assert!(probes.windows.is_empty()); + assert!(probes.limit_exceeded); + assert_eq!(probes.limit_exceeded_encoding, Some("base64")); + } + #[test] fn wrapped_base64_nested_literals_ignore_comment_leaders() { let value = "# this is doc\n# Y29zCnN5\n# c3RlbQopUi4=\n# more";