use crate::error::{Result, ScanError};
use crate::types::*;
use aho_corasick::{AhoCorasick, AhoCorasickBuilder};
use keyhog_core::{CompanionSpec, DetectorSpec, PatternSpec};
use regex::Regex;
use std::borrow::Cow;
static BUILD_GPU_LITERALS_INVOCATIONS: std::sync::atomic::AtomicUsize =
std::sync::atomic::AtomicUsize::new(0);
pub(crate) fn build_gpu_literals_invocations() -> usize {
BUILD_GPU_LITERALS_INVOCATIONS.load(std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn build_ac_pattern_set(literals: &[String]) -> Result<Option<AhoCorasick>> {
if literals.is_empty() {
return Ok(None);
}
Ok(Some(
AhoCorasickBuilder::new()
.ascii_case_insensitive(true)
.build(literals)?,
))
}
pub(crate) fn build_gpu_literals<'a>(
ac_literals: impl IntoIterator<Item = &'a [u8]>,
phase2_keywords: impl IntoIterator<Item = &'a [u8]>,
phase2_always_anchor_literals: impl IntoIterator<Item = &'a [u8]>,
confirmed_anchor_literals: impl IntoIterator<Item = &'a [u8]>,
generic_keyword_literals: impl IntoIterator<Item = &'a [u8]>,
) -> Option<std::sync::Arc<Vec<Vec<u8>>>> {
BUILD_GPU_LITERALS_INVOCATIONS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
build_gpu_literal_rows(
ac_literals
.into_iter()
.chain(phase2_keywords)
.chain(phase2_always_anchor_literals)
.chain(confirmed_anchor_literals)
.chain(generic_keyword_literals),
"GPU fused literal set",
)
}
static GPU_LITERAL_EMPTY_WARNED: std::sync::OnceLock<()> = std::sync::OnceLock::new();
fn build_gpu_literal_rows<'a>(
literals: impl Iterator<Item = &'a [u8]>,
label: &'static str,
) -> Option<std::sync::Arc<Vec<Vec<u8>>>> {
let mut rows = Vec::new();
for literal in literals {
if literal.is_empty() {
tracing::warn!("{label} contains an empty literal; disabling GPU literal scan");
if GPU_LITERAL_EMPTY_WARNED.set(()).is_ok() {
eprintln!(
"keyhog: a detector produced an empty literal in the {label}, so the GPU \
literal matcher was discarded and every scan will route through CPU/SIMD instead of the GPU \
literal path. Check your detector definitions for an empty AC literal (an empty `keywords`/\
prefix entry). Use --require-gpu when GPU acceleration is mandatory."
);
}
return None;
}
rows.push(literal.to_vec());
}
if rows.is_empty() {
None
} else {
tracing::info!(patterns = rows.len(), "{} prepared for VYRE", label);
Some(std::sync::Arc::new(rows))
}
}
pub(crate) fn build_same_prefix_patterns(literals: &[String]) -> crate::engine::CsrU32 {
let mut groups: std::collections::HashMap<&str, Vec<usize>> = std::collections::HashMap::new();
for (index, literal) in literals.iter().enumerate() {
groups.entry(literal.as_str()).or_default().push(index);
}
let mut pairs = Vec::new();
for indices in groups.values().filter(|indices| indices.len() > 1) {
for &row in indices {
pairs.extend(
indices
.iter()
.copied()
.filter(|&other| other != row)
.map(|other| (row, other)),
);
}
}
crate::engine::CsrU32::from_pairs(literals.len(), pairs)
}
pub(crate) fn build_prefix_propagation(literals: &[String]) -> crate::engine::CsrU32 {
crate::engine::CsrU32::from_pairs(
literals.len(),
crate::prefix_trie::build_propagation_pairs(literals),
)
}
const PHASE2_KEYWORD_BUCKET_COUNT: usize = 1 << 16;
#[inline]
fn phase2_keyword_prefix(bytes: &[u8]) -> u16 {
u16::from_be_bytes([bytes[0].to_ascii_lowercase(), bytes[1].to_ascii_lowercase()])
}
pub(crate) struct Phase2KeywordIndex {
bucket_offsets: Box<[u32]>,
bucket_keyword_ids: Box<[u32]>,
keywords: Box<[Box<[u8]>]>,
}
pub(crate) struct Phase2KeywordMatches<'a> {
index: &'a Phase2KeywordIndex,
haystack: &'a [u8],
next_position: usize,
}
impl Phase2KeywordIndex {
pub(crate) fn build(keywords: &[Cow<'_, str>]) -> Option<Self> {
if keywords.iter().any(|keyword| keyword.len() < 2) {
tracing::warn!(
"phase-2 keyword index received a sub-bigram literal; keyword-gate optimization disabled (recall preserved)"
);
return None;
}
let mut rows = Vec::with_capacity(keywords.len());
for (keyword_id, keyword) in keywords.iter().enumerate() {
let Ok(keyword_id) = u32::try_from(keyword_id) else {
tracing::warn!(
keywords = keywords.len(),
"phase-2 keyword index exceeds u32 rows; keyword-gate optimization disabled (recall preserved)"
);
return None;
};
rows.push((phase2_keyword_prefix(keyword.as_bytes()), keyword_id));
}
rows.sort_unstable();
let mut bucket_offsets = vec![0u32; PHASE2_KEYWORD_BUCKET_COUNT + 1];
let mut cursor = 0usize;
for (bucket, offset) in bucket_offsets
.iter_mut()
.take(PHASE2_KEYWORD_BUCKET_COUNT)
.enumerate()
{
*offset = u32::try_from(cursor).expect("keyword rows were bounded to u32");
while rows
.get(cursor)
.is_some_and(|(prefix, _)| usize::from(*prefix) == bucket)
{
cursor += 1;
}
}
bucket_offsets[PHASE2_KEYWORD_BUCKET_COUNT] =
u32::try_from(cursor).expect("keyword rows were bounded to u32");
Some(Self {
bucket_offsets: bucket_offsets.into_boxed_slice(),
bucket_keyword_ids: rows.into_iter().map(|(_, keyword_id)| keyword_id).collect(),
keywords: keywords
.iter()
.map(|keyword| Box::<[u8]>::from(keyword.as_bytes()))
.collect(),
})
}
#[inline]
pub(crate) fn find_iter<'a>(&'a self, haystack: &'a str) -> Phase2KeywordMatches<'a> {
Phase2KeywordMatches {
index: self,
haystack: haystack.as_bytes(),
next_position: 0,
}
}
}
impl Iterator for Phase2KeywordMatches<'_> {
type Item = usize;
fn next(&mut self) -> Option<Self::Item> {
let mut best: Option<(usize, usize)> = None;
let mut position = self.next_position;
while position + 1 < self.haystack.len() && best.is_none_or(|(end, _)| position + 2 <= end)
{
let prefix = usize::from(phase2_keyword_prefix(&self.haystack[position..]));
let candidate_start = self.index.bucket_offsets[prefix] as usize;
let candidate_end = self.index.bucket_offsets[prefix + 1] as usize;
for &keyword_id in &self.index.bucket_keyword_ids[candidate_start..candidate_end] {
let keyword_id = keyword_id as usize;
let keyword = &self.index.keywords[keyword_id];
let remaining = &self.haystack[position..];
if remaining.len() >= keyword.len()
&& remaining[..keyword.len()].eq_ignore_ascii_case(keyword)
{
let candidate = (position + keyword.len(), keyword_id);
if best.is_none_or(|current| candidate < current) {
best = Some(candidate);
}
}
}
position += 1;
}
match best {
Some((end, keyword_id)) => {
self.next_position = end;
Some(keyword_id)
}
None => {
self.next_position = self.haystack.len();
None
}
}
}
}
pub(crate) fn build_phase2_keyword_index<'a>(
phase2_patterns: &'a [(CompiledPattern, Vec<String>)],
) -> (
Option<Phase2KeywordIndex>,
crate::engine::CsrU32,
Vec<Cow<'a, str>>,
) {
let mut all_keywords = Vec::new();
let mut keyword_pattern_pairs = Vec::new();
let mut keyword_map: std::collections::HashMap<Cow<'a, str>, usize, ahash::RandomState> =
std::collections::HashMap::with_hasher(ahash::RandomState::new());
let mut add_candidate = |candidate: Cow<'a, str>, pattern_idx: usize| {
use std::collections::hash_map::Entry;
let idx = match keyword_map.entry(candidate) {
Entry::Occupied(entry) => *entry.get(),
Entry::Vacant(entry) => {
let idx = all_keywords.len();
all_keywords.push(entry.key().clone());
entry.insert(idx);
idx
}
};
keyword_pattern_pairs.push((idx, pattern_idx));
};
for (pattern_idx, (pattern, keywords)) in phase2_patterns.iter().enumerate() {
let allows_repeated_separator = pattern.allows_repeated_keyword_separator;
for keyword in keywords {
if keyword.len() >= 4 {
add_candidate(Cow::Borrowed(keyword.as_str()), pattern_idx);
}
if allows_repeated_separator {
if let Some(stem) = longest_compound_keyword_segment(keyword) {
if stem.len() >= 2 && stem != *keyword {
add_candidate(Cow::Owned(stem), pattern_idx);
}
}
}
}
}
if all_keywords.is_empty() {
return (
None,
crate::engine::CsrU32::from_pairs(0, std::iter::empty()),
Vec::new(),
);
}
let index = Phase2KeywordIndex::build(&all_keywords);
(
index,
crate::engine::CsrU32::from_pairs(all_keywords.len(), keyword_pattern_pairs),
all_keywords,
)
}
fn longest_compound_keyword_segment(keyword: &str) -> Option<String> {
keyword
.split(['_', '-', '.'])
.filter(|segment| {
segment.len() >= 2 && segment.bytes().all(|byte| byte.is_ascii_alphanumeric())
})
.max_by_key(|segment| segment.len())
.map(str::to_ascii_lowercase)
}
fn regex_allows_repeated_compound_keyword_separator(regex: &str) -> bool {
let Ok(hir) = regex_syntax::Parser::new().parse(regex) else {
return false;
};
hir_contains_repeated_separator(&hir)
}
fn hir_contains_repeated_separator(hir: ®ex_syntax::hir::Hir) -> bool {
use regex_syntax::hir::HirKind;
match hir.kind() {
HirKind::Repetition(repetition) => {
let repeats = repetition.max.is_none_or(|maximum| maximum > 1);
(repeats && hir_is_compound_keyword_separator(&repetition.sub))
|| hir_contains_repeated_separator(&repetition.sub)
}
HirKind::Capture(capture) => hir_contains_repeated_separator(&capture.sub),
HirKind::Concat(parts) | HirKind::Alternation(parts) => {
parts.iter().any(hir_contains_repeated_separator)
}
HirKind::Empty | HirKind::Literal(_) | HirKind::Class(_) | HirKind::Look(_) => false,
}
}
fn hir_is_compound_keyword_separator(hir: ®ex_syntax::hir::Hir) -> bool {
use regex_syntax::hir::{Class, HirKind};
match hir.kind() {
HirKind::Class(Class::Unicode(class)) => {
let mut has_join_punctuation = false;
for range in class.iter() {
let start = u32::from(range.start());
let end = u32::from(range.end());
if end - start > 32 {
return false;
}
for codepoint in start..=end {
let Some(character) = char::from_u32(codepoint) else {
return false;
};
if matches!(character, '_' | '-' | '.') {
has_join_punctuation = true;
} else if !character.is_whitespace() {
return false;
}
}
}
has_join_punctuation
}
HirKind::Class(Class::Bytes(class)) => {
let mut has_join_punctuation = false;
for range in class.iter() {
for byte in range.start()..=range.end() {
if matches!(byte, b'_' | b'-' | b'.') {
has_join_punctuation = true;
} else if !byte.is_ascii_whitespace() {
return false;
}
}
}
has_join_punctuation
}
_ => false,
}
}
pub(crate) fn log_quality_warnings(warnings: &[String]) {
for warning in warnings {
tracing::warn!(target: "keyhog::scanner::quality", "{}", warning);
}
}
pub(crate) fn compile_detector_companions(
detector: &DetectorSpec,
) -> Result<Vec<CompiledCompanion>> {
detector
.companions
.iter()
.map(|companion| compile_companion(companion, &detector.id))
.collect()
}
pub(crate) fn compile_pattern(
detector_index: usize,
pattern_index: usize,
spec: &PatternSpec,
detector_id: &str,
detector_keywords: &[String],
) -> Result<CompiledPattern> {
spec.validate_required_literals()
.map_err(|reason| ScanError::DetectorPatternPolicy {
detector_id: detector_id.to_string(),
index: pattern_index,
reason,
})?;
let validated =
shared_regex_compile(spec.regex.as_str()).map_err(|source| ScanError::RegexCompile {
detector_id: detector_id.to_string(),
index: pattern_index,
source,
})?;
if let Some(group) = spec.group {
let captures_len = validated.captures_len();
if group >= captures_len {
return Err(ScanError::CaptureGroupOutOfRange {
detector_id: detector_id.to_string(),
index: pattern_index,
group,
captures_len,
});
}
}
drop(validated);
Ok(CompiledPattern {
detector_index,
regex: LazyRegex::detector(spec.regex.as_str()),
group: spec.group,
client_safe: spec.client_safe,
weak_anchor: spec.weak_anchor,
structural_password_slot: spec.structural_password_slot,
match_proves_keyword_nearby: match_proves_keyword_nearby(
spec.regex.as_str(),
detector_keywords,
),
allows_repeated_keyword_separator: regex_allows_repeated_compound_keyword_separator(
spec.regex.as_str(),
),
homoglyph_variant: false,
})
}
pub(crate) fn match_proves_keyword_nearby(regex: &str, detector_keywords: &[String]) -> bool {
let prefixes = super::compiler_prefix::extract_literal_prefixes(regex);
!prefixes.is_empty()
&& prefixes.iter().all(|prefix| {
detector_keywords.iter().any(|keyword| {
!keyword.is_empty()
&& prefix
.as_bytes()
.get(..keyword.len())
.is_some_and(|head| head.eq_ignore_ascii_case(keyword.as_bytes()))
})
})
}
const REGEX_CACHE_SHARDS: usize = 64;
const REGEX_CACHE_CAPACITY: usize = 8192;
type RegexCacheShard = parking_lot::Mutex<lru::LruCache<String, std::sync::Weak<Regex>>>;
static REGEX_CACHE: std::sync::OnceLock<Box<[RegexCacheShard]>> = std::sync::OnceLock::new();
fn regex_cache() -> &'static [RegexCacheShard] {
REGEX_CACHE.get_or_init(|| {
let per_shard = (REGEX_CACHE_CAPACITY / REGEX_CACHE_SHARDS).max(1);
let nz = std::num::NonZeroUsize::new(per_shard).unwrap_or(std::num::NonZeroUsize::MIN); (0..REGEX_CACHE_SHARDS)
.map(|_| parking_lot::Mutex::new(lru::LruCache::new(nz)))
.collect::<Vec<_>>()
.into_boxed_slice()
})
}
fn regex_cache_shard(pattern: &str) -> &'static RegexCacheShard {
let idx = (crate::util_hash::hash_fast(pattern.as_bytes()) as usize) % REGEX_CACHE_SHARDS;
®ex_cache()[idx]
}
pub(crate) fn shared_regex_compile(
pattern: &str,
) -> std::result::Result<std::sync::Arc<Regex>, regex::Error> {
let regex = regex::RegexBuilder::new(pattern)
.case_insensitive(true)
.size_limit(REGEX_SIZE_LIMIT_BYTES)
.dfa_size_limit(regex_dfa_limit())
.crlf(true)
.build()?;
Ok(std::sync::Arc::new(regex))
}
pub(crate) fn shared_regex(
pattern: &str,
) -> std::result::Result<std::sync::Arc<Regex>, regex::Error> {
let shard = regex_cache_shard(pattern);
if let Some(hit) = shard.lock().get(pattern).and_then(std::sync::Weak::upgrade) {
return Ok(hit);
}
let arc = shared_regex_compile(pattern)?;
let mut lock = shard.lock();
if let Some(hit) = lock.get(pattern).and_then(std::sync::Weak::upgrade) {
return Ok(hit);
}
lock.put(pattern.to_string(), std::sync::Arc::downgrade(&arc));
Ok(arc)
}
pub(crate) fn shared_regex_cache_workload_probe(pattern: &str) -> (usize, usize) {
for shard in regex_cache() {
shard.lock().clear();
}
let first = shared_regex(pattern).expect("probe regex compiles");
let first_weak = std::sync::Arc::downgrade(&first);
let second = shared_regex(pattern).expect("live probe regex reuses cache");
let live_compiles = if std::sync::Arc::ptr_eq(&first, &second) {
1
} else {
2
};
drop((first, second));
let expired = first_weak.upgrade().is_none();
let completed_workload = shared_regex(pattern).expect("expired probe regex recompiles");
let completed_workload_compiles = live_compiles + usize::from(expired);
drop(completed_workload);
(live_compiles, completed_workload_compiles)
}
pub(crate) fn companion_regex(
pattern: &str,
) -> std::result::Result<std::sync::Arc<Regex>, regex::Error> {
regex::RegexBuilder::new(pattern)
.size_limit(REGEX_SIZE_LIMIT_BYTES)
.dfa_size_limit(regex_dfa_limit())
.crlf(true)
.build()
.map(std::sync::Arc::new)
}
pub(crate) fn compile_companion(
spec: &CompanionSpec,
detector_id: &str,
) -> Result<CompiledCompanion> {
let regex = companion_regex(&spec.regex).map_err(|e| ScanError::RegexCompile {
detector_id: detector_id.to_string(),
index: FIRST_CAPTURE_GROUP_INDEX,
source: e,
})?;
let capture_group = match spec.capture_group {
Some(group) if group < regex.captures_len() => Some(group),
Some(group) => {
return Err(ScanError::Config(format!(
"detector {detector_id:?} companion {:?} selects capture group {group}, \
but its regex exposes groups 0..{}",
spec.name,
regex.captures_len().saturating_sub(1)
)));
}
None => (regex.captures_len() > 1).then_some(FIRST_CAPTURE_GROUP_INDEX),
};
Ok(CompiledCompanion {
name: std::sync::Arc::from(spec.name.as_str()),
regex: LazyRegex::companion(spec.regex.as_str()),
capture_group,
within_lines: spec.within_lines,
within_bytes: spec.within_bytes,
direction: spec.direction,
scope: spec.scope,
requirement: spec.effective_requirement(),
value_relation: spec.value_relation,
})
}