use crate::middleware::holdback::HoldBackDfa;
use crate::middleware::MiddlewareError;
use crate::scanner::substitute::Substituter;
use crate::scanner::{Blocked, ScannerStack};
pub struct StreamOutput {
dfa: HoldBackDfa,
carry: Vec<u8>,
whole_stream: bool,
}
impl StreamOutput {
#[must_use]
pub fn new(stack: &ScannerStack) -> Self {
let whole_stream = stack.output_needs_whole_stream();
let dfa = HoldBackDfa::new(&stack.output_stream_patterns());
Self {
dfa,
carry: Vec::new(),
whole_stream,
}
}
pub fn push(
&mut self,
stack: &mut ScannerStack,
chunk: &str,
) -> Result<String, MiddlewareError> {
self.carry.extend_from_slice(chunk.as_bytes());
if self.whole_stream {
return Ok(String::new());
}
self.flush_safe_prefix(stack, false)
}
pub fn finish(mut self, stack: &mut ScannerStack) -> Result<String, MiddlewareError> {
self.flush_safe_prefix(stack, true)
}
fn flush_safe_prefix(
&mut self,
stack: &mut ScannerStack,
at_eof: bool,
) -> Result<String, MiddlewareError> {
let split = self.safe_split(stack, at_eof);
if split == 0 {
return Ok(String::new());
}
let prefix = self.take_prefix(split);
let scanned = stack.run_output(&prefix).map_err(|b| blocked(&b))?;
Ok(scanned)
}
fn take_prefix(&mut self, split: usize) -> String {
let rest = self.carry.split_off(split);
let prefix = std::mem::replace(&mut self.carry, rest);
String::from_utf8(prefix)
.unwrap_or_else(|e| String::from_utf8_lossy(e.as_bytes()).into_owned())
}
fn safe_split(&mut self, stack: &ScannerStack, at_eof: bool) -> usize {
let buf = &self.carry;
if at_eof {
return buf.len();
}
let mut split = self.dfa.safe_flush_len(buf, false);
split = split.min(token_boundary_floor(buf));
split = split.min(sentinel_floor(stack, buf));
char_boundary_floor(buf, split)
}
}
fn token_boundary_floor(buf: &[u8]) -> usize {
match buf.iter().rposition(u8::is_ascii_whitespace) {
Some(pos) => pos + 1,
None => 0,
}
}
fn sentinel_floor(stack: &ScannerStack, buf: &[u8]) -> usize {
let pairs: Vec<(String, String)> = stack
.ctx()
.vault
.entries()
.map(|(s, o)| (s.to_owned(), o.to_owned()))
.collect();
if pairs.is_empty() {
return buf.len();
}
let sub = Substituter::from_pairs(pairs.iter().map(|(s, o)| (s.as_str(), o.as_str())));
let longest = sub.longest_sentinel();
let window_start = buf.len().saturating_sub(longest);
for start in window_start..buf.len() {
if sub.is_sentinel_prefix(&buf[start..]) {
return start;
}
}
buf.len()
}
fn char_boundary_floor(buf: &[u8], mut split: usize) -> usize {
while split > 0 && split < buf.len() && buf[split] & 0xC0 == 0x80 {
split -= 1;
}
split
}
fn blocked(b: &Blocked) -> MiddlewareError {
MiddlewareError::Blocked {
detail: format!("{} (risk {})", b.scanner, b.risk),
}
}
#[cfg(test)]
mod tests {
#![allow(
clippy::unwrap_used,
clippy::expect_used,
reason = "tests assert on known-good values"
)]
use super::*;
#[test]
fn token_floor_holds_partial_trailing_token() {
assert_eq!(token_boundary_floor(b"see AKIA"), 4);
}
#[test]
fn token_floor_with_no_whitespace_holds_everything() {
assert_eq!(token_boundary_floor(b"AKIAABC"), 0);
}
#[test]
fn token_floor_trailing_space_flushes_whole() {
assert_eq!(token_boundary_floor(b"done "), 5);
}
#[test]
fn char_boundary_snaps_below_a_multibyte_char() {
let buf = "a€".as_bytes();
assert_eq!(char_boundary_floor(buf, 2), 1);
assert_eq!(char_boundary_floor(buf, 3), 1);
assert_eq!(char_boundary_floor(buf, 1), 1);
assert_eq!(char_boundary_floor(buf, 4), 4);
}
}