diff --git a/crates/goose/src/security/classification_client.rs b/crates/goose/src/security/classification_client.rs index cb6bf54386..78f837cb1a 100644 --- a/crates/goose/src/security/classification_client.rs +++ b/crates/goose/src/security/classification_client.rs @@ -238,11 +238,27 @@ impl ClassificationClient { } pub async fn classify_chunked(&self, text: &str) -> ChunkedScan { - use crate::security::command_chunker::chunk_command; + use crate::security::command_chunker::{chunk_command, MAX_WINDOWS}; - let chunks = chunk_command(text); + const COMMAND_SCAN_CONCURRENCY: usize = 3; + + let mut chunks = chunk_command(text); let chunk_count = chunks.len(); + let mut unscanned = 0usize; + if chunk_count > MAX_WINDOWS { + unscanned = chunk_count - MAX_WINDOWS; + chunks.truncate(MAX_WINDOWS); + tracing::warn!( + monotonic_counter.goose.command_classifier_oversized = 1, + security.event_type = "command_classifier_chunking", + security.threat_type = "command_injection", + scanner.chunk_count = chunk_count, + scanner.window_cap = MAX_WINDOWS, + "command exceeds window cap; scanning capped windows and treating remainder as unscanned" + ); + } + if chunk_count == 1 { return match self.classify(text).await { Ok(conf) => ChunkedScan { @@ -275,7 +291,7 @@ impl ClassificationClient { let results: Vec> = stream::iter(chunks) .map(|chunk| async move { self.classify(&chunk).await }) - .buffer_unordered(3) + .buffer_unordered(COMMAND_SCAN_CONCURRENCY) .collect() .await; @@ -298,7 +314,7 @@ impl ClassificationClient { } } } - let failed = total - succeeded; + let failed = (total - succeeded) + unscanned; if failed > 0 { tracing::warn!( @@ -371,6 +387,22 @@ mod tests { .expect("client construction should succeed") } + #[tokio::test] + async fn classify_chunked_marks_oversized_commands_incomplete() { + use crate::security::command_chunker::exceeds_window_cap; + let client = unroutable_client(); + let huge = "a; ".repeat(4000); + assert!(exceeds_window_cap(&huge), "test input must exceed the cap"); + + let scan = client.classify_chunked(&huge).await; + + assert!( + scan.had_failures(), + "an oversized command must be reported as incomplete, not a clean pass" + ); + assert!(scan.failed >= 1); + } + #[tokio::test] async fn classify_chunked_reports_failures_when_all_windows_fail() { let client = unroutable_client(); diff --git a/crates/goose/src/security/command_chunker.rs b/crates/goose/src/security/command_chunker.rs index eab679567d..95c8295551 100644 --- a/crates/goose/src/security/command_chunker.rs +++ b/crates/goose/src/security/command_chunker.rs @@ -2,15 +2,25 @@ const MODEL_MAX_TOKENS: usize = 512; const SPECIAL_TOKEN_HEADROOM: usize = 12; +// Window size in BYTES, used as a worst-case proxy for tokens: the classifier's +// tokenizer never emits more than one token per byte, so a window of at most this +// many bytes cannot exceed MODEL_MAX_TOKENS. Changing the model/tokenizer to one +// that can produce >1 token/byte would break this bound. const MAX_WINDOW_CHARS: usize = MODEL_MAX_TOKENS - SPECIAL_TOKEN_HEADROOM; const OVERLAP_CHARS: usize = 256; +pub const MAX_WINDOWS: usize = 12; + pub fn chunk_command(text: &str) -> Vec { let overlap_ratio = OVERLAP_CHARS as f32 / MAX_WINDOW_CHARS as f32; chunk_with_params(text, MAX_WINDOW_CHARS, overlap_ratio) } +pub fn exceeds_window_cap(text: &str) -> bool { + chunk_command(text).len() > MAX_WINDOWS +} + #[allow(clippy::string_slice)] fn chunk_with_params(text: &str, max_chars: usize, overlap_ratio: f32) -> Vec { debug_assert!(max_chars > 0); @@ -22,19 +32,22 @@ fn chunk_with_params(text: &str, max_chars: usize, overlap_ratio: f32) -> Vec 0, "stride must be positive to make progress"); let mut chunks = Vec::new(); let mut start = 0; while start < text.len() { - let hard_end = (start + max_chars).min(text.len()); - let end = floor_char_boundary(text, hard_end); let real_start = floor_char_boundary(text, start); + let hard_end = (real_start + max_chars).min(text.len()); + let end = floor_char_boundary(text, hard_end); chunks.push(text[real_start..end].to_string()); if end >= text.len() { break; } - start = floor_char_boundary(text, start + stride).max(real_start + 1); + let next = floor_char_boundary(text, real_start + stride); + debug_assert!(next > real_start, "each window must advance past the last"); + start = next; } chunks } @@ -79,27 +92,26 @@ mod tests { #[test] fn full_text_is_covered() { - let text: String = (0..5000).map(|i| (b'a' + (i % 26) as u8) as char).collect(); + let text: String = (0..3000u32).map(|i| format!("{i:05}")).collect(); let chunks = chunk_with_params(&text, 300, 0.25); - let mut covered = vec![false; text.len()]; - let mut pos = 0; - let max_chars = 300usize; - let overlap = 75usize; - let stride = max_chars - overlap; - let mut start = 0; - while start < text.len() { - let end = (start + max_chars).min(text.len()); - for c in covered.iter_mut().take(end).skip(start) { + + let bytes = text.as_bytes(); + let mut covered = vec![false; bytes.len()]; + for chunk in &chunks { + let cb = chunk.as_bytes(); + let start = bytes + .windows(cb.len()) + .position(|w| w == cb) + .expect("each chunk is a substring of the input"); + for c in covered.iter_mut().skip(start).take(cb.len()) { *c = true; } - if end >= text.len() { - break; - } - start += stride; - pos += 1; } - assert!(covered.iter().all(|&c| c), "every byte must be covered"); - assert!(pos + 1 >= chunks.len().saturating_sub(1)); + + assert!( + covered.iter().all(|&c| c), + "every byte of the input must be covered by some window" + ); } #[test] @@ -143,16 +155,17 @@ mod tests { } #[test] - fn window_never_exceeds_token_budget() { + fn window_never_exceeds_char_budget() { let text: String = (0..10_000) .map(|i| (b'a' + (i % 26) as u8) as char) .collect(); let chunks = chunk_command(&text); for c in &chunks { assert!( - c.chars().count() <= MODEL_MAX_TOKENS, - "window has {} chars, exceeds token budget", - c.chars().count() + c.len() <= MAX_WINDOW_CHARS, + "window has {} bytes, exceeds worst-case token budget of {}", + c.len(), + MAX_WINDOW_CHARS ); } }