use regex_automata::{
hybrid::dfa::{Cache, DFA},
util::start::Config as StartConfig,
Anchored, MatchKind,
};
#[derive(Debug)]
pub struct HoldBackDfa {
dfa: DFA,
empty: bool,
}
impl HoldBackDfa {
#[must_use]
pub fn new(patterns: &[String]) -> Self {
if patterns.is_empty() {
return Self::nothing();
}
let patterns: Vec<String> = patterns.iter().map(|p| ascii_word_boundaries(p)).collect();
let built = DFA::builder()
.configure(DFA::config().match_kind(MatchKind::All))
.build_many(&patterns);
match built {
Ok(dfa) => Self { dfa, empty: false },
Err(_) => Self::compile_individually(&patterns),
}
}
fn nothing() -> Self {
let dfa = DFA::builder()
.configure(DFA::config().match_kind(MatchKind::All))
.build_many::<&str>(&[])
.unwrap_or_else(|_| never_match_dfa());
Self { dfa, empty: true }
}
fn compile_individually(patterns: &[String]) -> Self {
let good: Vec<String> = patterns
.iter()
.filter(|p| {
DFA::builder()
.configure(DFA::config().match_kind(MatchKind::All))
.build(p)
.is_ok()
})
.cloned()
.collect();
if good.is_empty() {
return Self::nothing();
}
match DFA::builder()
.configure(DFA::config().match_kind(MatchKind::All))
.build_many(&good)
{
Ok(dfa) => Self { dfa, empty: false },
Err(_) => Self::nothing(),
}
}
#[must_use]
pub fn safe_flush_len(&self, buf: &[u8], at_eof: bool) -> usize {
if self.empty || at_eof {
return buf.len();
}
let mut cache = self.dfa.create_cache();
for s in 0..buf.len() {
if self.alive_at_eof(&mut cache, &buf[s..]) {
return s;
}
}
buf.len()
}
fn alive_at_eof(&self, cache: &mut Cache, tail: &[u8]) -> bool {
let Ok(start) = self
.dfa
.start_state(cache, &StartConfig::new().anchored(Anchored::Yes))
else {
return true;
};
let mut state = start;
for &b in tail {
state = match self.dfa.next_state(cache, state, b) {
Ok(s) => s,
Err(_) => return true,
};
if state.is_dead() {
return false;
}
}
!state.is_dead()
}
}
fn ascii_word_boundaries(pattern: &str) -> String {
let mut out = String::with_capacity(pattern.len());
let mut chars = pattern.chars().peekable();
while let Some(c) = chars.next() {
if c == '\\' {
match chars.next() {
Some(esc @ ('b' | 'B')) => {
out.push_str("(?-u:\\");
out.push(esc);
out.push(')');
}
Some(esc) => {
out.push('\\');
out.push(esc);
}
None => out.push('\\'),
}
continue;
}
out.push(c);
}
out
}
fn never_match_dfa() -> DFA {
match DFA::builder()
.configure(DFA::config().match_kind(MatchKind::All))
.build(r"[^\s\S]")
{
Ok(dfa) => dfa,
Err(_) => never_match_dfa(),
}
}
#[cfg(test)]
mod tests {
#![allow(
clippy::unwrap_used,
clippy::expect_used,
reason = "tests assert on known-good values"
)]
use super::*;
fn dfa(pats: &[&str]) -> HoldBackDfa {
HoldBackDfa::new(&pats.iter().map(|p| (*p).to_owned()).collect::<Vec<_>>())
}
#[test]
fn empty_pattern_set_flushes_everything() {
let d = dfa(&[]);
assert_eq!(d.safe_flush_len(b"anything at all", false), 15);
}
#[test]
fn holds_from_a_forming_match() {
let d = dfa(&[r"AKIA[0-9A-Z]{16}"]);
let buf = b"see AKIAABC";
let flush = d.safe_flush_len(buf, false);
assert_eq!(&buf[..flush], b"see ");
}
#[test]
fn flushes_when_match_dies() {
let d = dfa(&[r"AKIA[0-9A-Z]{16}"]);
let buf = b"AKIA!rest";
assert_eq!(d.safe_flush_len(buf, false), buf.len());
}
#[test]
fn eof_collapses_holdback() {
let d = dfa(&[r"AKIA[0-9A-Z]{16}"]);
let buf = b"see AKIAABC";
assert_eq!(d.safe_flush_len(buf, true), buf.len());
}
#[test]
fn unbounded_tail_holds_from_match_start() {
let d = dfa(&[r"sk-[A-Za-z0-9]{20,}"]);
let buf = b"key sk-aaaaaaaaaaaaaaaaaaaaaa";
let flush = d.safe_flush_len(buf, false);
assert_eq!(&buf[..flush], b"key ");
}
#[test]
fn clean_buffer_with_patterns_flushes_whole() {
let d = dfa(&[r"AKIA[0-9A-Z]{16}"]);
let buf = b"the quick brown fox";
assert_eq!(d.safe_flush_len(buf, false), buf.len());
}
#[test]
fn one_bad_pattern_does_not_disable_others() {
let d = dfa(&[r"AKIA[0-9A-Z]{16}", r"(unclosed"]);
let buf = b"AKIAABC";
assert_eq!(d.safe_flush_len(buf, false), 0);
}
#[test]
fn rewrites_unicode_word_boundaries_to_ascii() {
assert_eq!(
ascii_word_boundaries(r"\b\d{3}\b"),
r"(?-u:\b)\d{3}(?-u:\b)"
);
assert_eq!(ascii_word_boundaries(r"\d{3}\."), r"\d{3}\.");
assert_eq!(ascii_word_boundaries(r"\B"), r"(?-u:\B)");
}
#[test]
fn word_boundary_pattern_builds_and_holds() {
let d = dfa(&[r"\b(?:\d[ \-]?){13,19}\b"]);
let buf = b"card 4111 1111 1111 ";
let flush = d.safe_flush_len(buf, false);
assert_eq!(&buf[..flush], b"card ");
}
}