use keyhog_core::{Chunk, ChunkMetadata, RawMatch};
use regex_syntax::ast::{Ast, RepetitionKind, RepetitionRange};
use super::{absolute_line, absolute_offset, floor_char_boundary, CompiledScanner};
use crate::types::CompiledPattern;
pub(crate) const MAX_BOUNDARY_SEAM_BYTES: usize = crate::types::WINDOW_OVERLAP_BYTES;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum BoundaryContextBytes {
Bounded(usize),
FullAdjacentChunks,
}
pub(crate) fn derive_pattern_boundary_context<'a>(
patterns: impl IntoIterator<Item = &'a CompiledPattern>,
) -> BoundaryContextBytes {
let mut max_bound = 0usize;
for pattern in patterns {
let Some(bound) = regex_match_byte_upper_bound(pattern.regex.as_str()) else {
return BoundaryContextBytes::FullAdjacentChunks;
};
max_bound = max_bound.max(bound);
}
BoundaryContextBytes::Bounded(max_bound)
}
pub(crate) fn regex_match_byte_upper_bound(source: &str) -> Option<usize> {
let ast = match regex_syntax::ast::parse::Parser::new().parse(source) {
Ok(ast) => ast,
Err(_) => return None, };
ast_match_byte_upper_bound(&ast)
}
fn ast_match_byte_upper_bound(ast: &Ast) -> Option<usize> {
match ast {
Ast::Empty(_) | Ast::Flags(_) | Ast::Assertion(_) => Some(0),
Ast::Literal(literal) => Some(literal.c.len_utf8()),
Ast::Dot(_) | Ast::ClassUnicode(_) | Ast::ClassPerl(_) | Ast::ClassBracketed(_) => Some(4),
Ast::Group(group) => ast_match_byte_upper_bound(&group.ast),
Ast::Alternation(alternation) => {
let mut max_bound = 0usize;
for ast in &alternation.asts {
max_bound = max_bound.max(ast_match_byte_upper_bound(ast)?);
}
Some(max_bound)
}
Ast::Concat(concat) => {
let mut total = 0usize;
for ast in &concat.asts {
total = total.saturating_add(ast_match_byte_upper_bound(ast)?);
}
Some(total)
}
Ast::Repetition(repetition) => {
let inner = ast_match_byte_upper_bound(&repetition.ast)?;
let max_repetitions = match repetition.op.kind {
RepetitionKind::ZeroOrOne => 1,
RepetitionKind::ZeroOrMore | RepetitionKind::OneOrMore => return None,
RepetitionKind::Range(RepetitionRange::Exactly(n)) => n,
RepetitionKind::Range(RepetitionRange::Bounded(_, n)) => n,
RepetitionKind::Range(RepetitionRange::AtLeast(_)) => return None,
};
Some(inner.saturating_mul(max_repetitions as usize))
}
}
}
#[cfg(test)]
pub(crate) fn scan_chunk_boundaries(
scanner: &CompiledScanner,
chunks: &[Chunk],
per_chunk_results: &mut [Vec<RawMatch>],
) -> crate::error::Result<()> {
scan_chunk_boundaries_with_route(
scanner,
chunks,
per_chunk_results,
scanner.default_execution_route(),
)
}
pub(crate) fn scan_chunk_boundaries_with_route(
scanner: &CompiledScanner,
chunks: &[Chunk],
per_chunk_results: &mut [Vec<RawMatch>],
route: crate::ScanExecutionRoute,
) -> crate::error::Result<()> {
if chunks.len() < 2 {
return Ok(());
}
if chunks.len() != per_chunk_results.len() {
crate::telemetry::record_boundary_result_cardinality_mismatch();
return Ok(());
}
use std::collections::HashMap;
let mut groups: HashMap<(&str, &str), Vec<usize>> = HashMap::new();
for (i, c) in chunks.iter().enumerate() {
let Some(path) = c.metadata.path.as_deref() else {
continue;
};
groups
.entry((c.metadata.source_type.as_ref(), path))
.or_default()
.push(i);
}
for (_, mut indices) in groups {
if indices.len() < 2 {
continue;
}
indices.sort_by_key(|&i| chunks[i].metadata.base_offset);
for w in indices.windows(2) {
let (ai, bi) = (w[0], w[1]);
scan_one_pair(
scanner,
&chunks[ai],
&chunks[bi],
ai,
bi,
per_chunk_results,
route,
)?;
}
}
Ok(())
}
fn scan_one_pair(
scanner: &CompiledScanner,
a: &Chunk,
b: &Chunk,
ai: usize,
bi: usize,
per_chunk_results: &mut [Vec<RawMatch>],
route: crate::ScanExecutionRoute,
) -> crate::error::Result<()> {
if ai >= per_chunk_results.len() || bi >= per_chunk_results.len() {
crate::telemetry::record_boundary_result_cardinality_mismatch();
return Ok(());
}
let a_bytes = a.data.as_ref().as_bytes();
let b_bytes = b.data.as_ref().as_bytes();
let a_end = a.metadata.base_offset.saturating_add(a_bytes.len());
if a_end != b.metadata.base_offset {
return Ok(());
}
if a_bytes.is_empty() || b_bytes.is_empty() {
return Ok(());
}
let path = b.metadata.path.as_deref().or(a.metadata.path.as_deref());
let context = boundary_context_for_pair(scanner, path);
let context_bytes = match context {
BoundaryContextBytes::Bounded(bytes) => bytes,
BoundaryContextBytes::FullAdjacentChunks => MAX_BOUNDARY_SEAM_BYTES,
};
let tail_start = a_bytes.len().saturating_sub(context_bytes);
let tail_start = floor_char_boundary(a.data.as_ref(), tail_start);
let tail = &a.data.as_ref()[tail_start..];
let head_end = b_bytes.len().min(context_bytes);
let head_end = floor_char_boundary(b.data.as_ref(), head_end);
let head = &b.data.as_ref()[..head_end];
if tail.is_empty() || head.is_empty() {
return Ok(());
}
let Some(boundary_base_offset) = absolute_offset(a.metadata.base_offset, tail_start) else {
return Ok(());
};
let mut buf = String::with_capacity(tail.len() + head.len());
buf.push_str(tail);
let seam_local = buf.len();
buf.push_str(head);
let boundary_base_line = absolute_line(
a.metadata.base_line,
memchr::memchr_iter(b'\n', &a_bytes[..tail_start]).count(),
);
let boundary_chunk = Chunk {
data: buf.into(),
metadata: ChunkMetadata {
base_offset: boundary_base_offset,
base_line: boundary_base_line,
..b.metadata.clone()
},
};
let boundary_matches = scan_boundary_chunk_whole(scanner, &boundary_chunk, route)?;
let Some(seam_file_offset) = absolute_offset(boundary_base_offset, seam_local) else {
return Ok(());
};
for m in boundary_matches {
let start = m.location.offset;
let end = start.saturating_add(m.credential.as_ref().len());
if end <= seam_file_offset {
continue;
}
let already_seen = per_chunk_results[ai]
.iter()
.chain(per_chunk_results[bi].iter())
.any(|x| {
x.location.offset == m.location.offset
&& x.detector_id == m.detector_id
&& x.credential_hash == m.credential_hash
});
if already_seen {
continue;
}
per_chunk_results[bi].push(m);
}
Ok(())
}
fn boundary_context_for_pair(
scanner: &CompiledScanner,
path: Option<&str>,
) -> BoundaryContextBytes {
if matches!(
scanner.pattern_boundary_context,
BoundaryContextBytes::FullAdjacentChunks
) {
return BoundaryContextBytes::FullAdjacentChunks;
}
if scanner.config.entropy_enabled
&& crate::entropy::is_entropy_appropriate(path, scanner.config.entropy_in_source_files)
{
return BoundaryContextBytes::FullAdjacentChunks;
}
scanner.pattern_boundary_context
}
fn scan_boundary_chunk_whole(
scanner: &CompiledScanner,
chunk: &Chunk,
route: crate::ScanExecutionRoute,
) -> crate::error::Result<Vec<RawMatch>> {
let mut matches = scanner.scan_inner(
chunk,
crate::hw_probe::ScanBackend::CpuFallback,
None,
route,
)?;
scanner.post_process_matches(chunk, &mut matches, None, route)?;
Ok(matches)
}