#[cfg(feature = "ml")]
use std::collections::HashMap;
use std::collections::{BinaryHeap, HashSet};
use std::sync::Arc;
use keyhog_core::SensitiveString;
#[cfg(feature = "ml")]
use zeroize::Zeroize;
#[cfg(feature = "ml")]
pub(crate) fn ml_context_for_candidate(
text: &str,
line: usize,
file_path: Option<&str>,
context_radius_lines: usize,
) -> String {
let text_context = crate::pipeline::local_context_window(text, line, context_radius_lines);
match file_path {
Some(path) => format!("file:{path}\n{text_context}"),
None => text_context.to_string(),
}
}
#[cfg(feature = "ml")]
pub(crate) fn ml_features_for_candidate(
text: &str,
line: usize,
file_path: Option<&str>,
credential: &str,
context_radius_lines: usize,
config: &crate::types::ScannerConfig,
detector_service: &str,
detector_features: crate::ml_scorer::ml_features::CompiledDetectorMlFeatures,
channel: crate::ml_scorer::MlCandidateChannel,
) -> [f32; crate::ml_scorer::NUM_FEATURES] {
if credential.is_empty() {
return [0.0; crate::ml_scorer::NUM_FEATURES];
}
let text_context = crate::pipeline::local_context_window(text, line, context_radius_lines);
let compute = |context: &str| {
crate::ml_scorer::ml_features::compute_features_for_compiled_detector_with_config(
credential,
context,
&config.known_prefixes,
&config.secret_keywords,
&config.test_keywords,
&config.placeholder_keywords,
detector_service,
detector_features,
channel,
)
};
let Some(path) = file_path else {
return compute(text_context);
};
thread_local! {
static CONTEXT: std::cell::RefCell<String> = const { std::cell::RefCell::new(String::new()) };
}
CONTEXT.with(|cell| {
let mut context = cell.borrow_mut();
context.clear();
context.push_str("file:");
context.push_str(path);
context.push('\n');
context.push_str(text_context);
let features = compute(&context);
context.zeroize();
features
})
}
#[cfg(feature = "ml")]
#[derive(Debug, Clone)]
pub(crate) struct MlPendingMatch {
pub(crate) raw_match: keyhog_core::RawMatch,
pub(crate) heuristic_conf: f64,
pub(crate) code_context: crate::context::CodeContext,
pub(crate) context_multiplier: f64,
pub(crate) context_suppression_threshold: Option<f64>,
pub(crate) post_match: keyhog_core::DetectorPostMatchConfidenceSpec,
pub(crate) ml_features: [f32; crate::ml_scorer::NUM_FEATURES],
channel: crate::ml_scorer::MlCandidateChannel,
pub(crate) ml_weight: f64,
pub(crate) min_confidence_floor: f64,
pub(crate) is_named_detector: bool,
pub(crate) is_generic_detector: bool,
pub(crate) allow_canonical_hex_key: bool,
pub(crate) allow_encoded_text_lift: bool,
pub(crate) checksum: crate::checksum::ChecksumConfidenceDecision,
pub(crate) ml_mode: crate::detector_ml_policy::ActiveMlMode,
}
#[cfg(feature = "ml")]
impl MlPendingMatch {
pub(crate) fn detector_candidate(
raw_match: keyhog_core::RawMatch,
heuristic_conf: f64,
code_context: crate::context::CodeContext,
context_multiplier: f64,
context_suppression_threshold: Option<f64>,
post_match: keyhog_core::DetectorPostMatchConfidenceSpec,
ml_features: [f32; crate::ml_scorer::NUM_FEATURES],
ml_weight: f64,
min_confidence_floor: f64,
is_named_detector: bool,
is_generic_detector: bool,
allow_canonical_hex_key: bool,
allow_encoded_text_lift: bool,
checksum: crate::checksum::ChecksumConfidenceDecision,
ml_mode: crate::detector_ml_policy::ActiveMlMode,
) -> Self {
Self {
raw_match,
heuristic_conf,
code_context,
context_multiplier,
context_suppression_threshold,
post_match,
ml_features,
channel: crate::ml_scorer::MlCandidateChannel::Pattern,
ml_weight,
min_confidence_floor,
is_named_detector,
is_generic_detector,
allow_canonical_hex_key,
allow_encoded_text_lift,
checksum,
ml_mode,
}
}
#[cfg(feature = "entropy")]
pub(crate) fn entropy_candidate(
raw_match: keyhog_core::RawMatch,
heuristic_conf: f64,
context_multiplier: f64,
context_suppression_threshold: Option<f64>,
post_match: keyhog_core::DetectorPostMatchConfidenceSpec,
ml_features: [f32; crate::ml_scorer::NUM_FEATURES],
ml_weight: f64,
min_confidence_floor: f64,
allow_canonical_hex_key: bool,
checksum: crate::checksum::ChecksumConfidenceDecision,
ml_mode: crate::detector_ml_policy::ActiveMlMode,
) -> Self {
Self {
raw_match,
heuristic_conf,
code_context: crate::context::CodeContext::Unknown,
context_multiplier,
context_suppression_threshold,
post_match,
ml_features,
channel: crate::ml_scorer::MlCandidateChannel::Entropy,
ml_weight,
min_confidence_floor,
is_named_detector: false,
is_generic_detector: true,
allow_canonical_hex_key,
allow_encoded_text_lift: false,
checksum,
ml_mode,
}
}
}
#[cfg(feature = "ml")]
impl MlPendingMatch {
fn has_same_execution_as(&self, other: &Self) -> bool {
self.channel == other.channel
&& self.code_context == other.code_context
&& self.context_multiplier.to_bits() == other.context_multiplier.to_bits()
&& self.context_suppression_threshold.map(f64::to_bits)
== other.context_suppression_threshold.map(f64::to_bits)
&& post_match_execution_eq(self.post_match, other.post_match)
&& self.ml_features == other.ml_features
&& self.ml_weight.to_bits() == other.ml_weight.to_bits()
&& self.min_confidence_floor.to_bits() == other.min_confidence_floor.to_bits()
&& self.is_named_detector == other.is_named_detector
&& self.is_generic_detector == other.is_generic_detector
&& self.allow_canonical_hex_key == other.allow_canonical_hex_key
&& self.allow_encoded_text_lift == other.allow_encoded_text_lift
&& self.checksum == other.checksum
&& self.ml_mode == other.ml_mode
}
}
#[cfg(feature = "ml")]
fn post_match_execution_eq(
left: keyhog_core::DetectorPostMatchConfidenceSpec,
right: keyhog_core::DetectorPostMatchConfidenceSpec,
) -> bool {
left.placeholder_multiplier.to_bits() == right.placeholder_multiplier.to_bits()
&& left.minimum_byte_diversity.to_bits() == right.minimum_byte_diversity.to_bits()
&& left.low_diversity_multiplier.to_bits() == right.low_diversity_multiplier.to_bits()
&& left.maximum_repeat_ratio.to_bits() == right.maximum_repeat_ratio.to_bits()
&& left.degenerate_run_min_length == right.degenerate_run_min_length
&& left.degenerate_repeat_multiplier.to_bits()
== right.degenerate_repeat_multiplier.to_bits()
&& left.data_envelope_multiplier.map(f64::to_bits)
== right.data_envelope_multiplier.map(f64::to_bits)
&& left.fixture_path_multiplier.to_bits() == right.fixture_path_multiplier.to_bits()
&& left.ml_context_reapply_below.to_bits() == right.ml_context_reapply_below.to_bits()
}
#[cfg(feature = "ml")]
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct PendingMatchIdentity {
detector_id: Arc<str>,
credential_hash: keyhog_core::CredentialHash,
offset: usize,
channel: crate::ml_scorer::MlCandidateChannel,
}
#[cfg(feature = "ml")]
impl From<&MlPendingMatch> for PendingMatchIdentity {
fn from(pending: &MlPendingMatch) -> Self {
Self {
detector_id: pending.raw_match.detector_id.clone(),
credential_hash: pending.raw_match.credential_hash,
offset: pending.raw_match.location.offset,
channel: pending.channel,
}
}
}
#[cfg(any(feature = "entropy", test))]
pub(crate) struct RawMatchPriority<'a> {
pub(crate) confidence: Option<f64>,
pub(crate) severity: keyhog_core::Severity,
pub(crate) detector_id: &'a str,
pub(crate) credential: &'a str,
pub(crate) offset: usize,
pub(crate) line: Option<usize>,
}
#[cfg(any(feature = "entropy", test))]
impl RawMatchPriority<'_> {
fn cmp_raw_match(&self, other: &keyhog_core::RawMatch) -> std::cmp::Ordering {
let self_conf = self.confidence.unwrap_or(0.0); let other_conf = other.confidence.unwrap_or(0.0);
match other_conf.total_cmp(&self_conf) {
std::cmp::Ordering::Equal => {}
ord => return ord,
}
match other.severity.cmp(&self.severity) {
std::cmp::Ordering::Equal => {}
ord => return ord,
}
match self.detector_id.cmp(other.detector_id.as_ref()) {
std::cmp::Ordering::Equal => {}
ord => return ord,
}
match self.credential.cmp(other.credential.as_ref()) {
std::cmp::Ordering::Equal => {}
ord => return ord,
}
match self.offset.cmp(&other.location.offset) {
std::cmp::Ordering::Equal => self.line.cmp(&other.location.line),
ord => ord,
}
}
}
fn raw_match_identity_cmp(
a: &keyhog_core::RawMatch,
b: &keyhog_core::RawMatch,
) -> std::cmp::Ordering {
MatchIdentity::from(a).cmp(&MatchIdentity::from(b))
}
fn same_raw_match_identity(a: &keyhog_core::RawMatch, b: &keyhog_core::RawMatch) -> bool {
MatchIdentity::from(a) == MatchIdentity::from(b)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct MatchIdentity<'a> {
detector_id: &'a str,
credential: &'a str,
offset: usize,
}
impl<'a> From<&'a keyhog_core::RawMatch> for MatchIdentity<'a> {
fn from(raw_match: &'a keyhog_core::RawMatch) -> Self {
Self {
detector_id: raw_match.detector_id.as_ref(),
credential: raw_match.credential.as_ref(),
offset: raw_match.location.offset,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct OwnedMatchIdentity {
detector_id: Arc<str>,
credential: SensitiveString,
offset: usize,
}
impl From<&keyhog_core::RawMatch> for OwnedMatchIdentity {
fn from(raw_match: &keyhog_core::RawMatch) -> Self {
Self {
detector_id: raw_match.detector_id.clone(),
credential: raw_match.credential.clone(),
offset: raw_match.location.offset,
}
}
}
impl OwnedMatchIdentity {
fn matches_raw(&self, m: &keyhog_core::RawMatch) -> bool {
self.offset == m.location.offset
&& self.detector_id.as_ref() == m.detector_id.as_ref()
&& self.credential.as_ref() == m.credential.as_ref()
}
}
impl OwnedMatchIdentity {
#[cfg(any(feature = "entropy", test))]
fn from_priority(priority: &RawMatchPriority<'_>) -> Self {
Self {
detector_id: Arc::from(priority.detector_id),
credential: SensitiveString::from(priority.credential),
offset: priority.offset,
}
}
}
#[derive(Default)]
pub(crate) struct ScanState {
pub(crate) matches: BinaryHeap<keyhog_core::RawMatch>,
pub(crate) credential_interner: HashSet<SensitiveString>,
pub(crate) metadata_interner: HashSet<Arc<str>>,
claimed_match_identities: HashSet<OwnedMatchIdentity>,
pub(crate) static_intern: Option<Arc<crate::static_intern::StaticInterner>>,
#[cfg(feature = "ml")]
pub(crate) ml_pending: Vec<MlPendingMatch>,
#[cfg(feature = "ml")]
ml_pending_index: HashMap<PendingMatchIdentity, usize>,
}
impl ScanState {
pub(crate) fn intern_credential(&mut self, s: &str) -> SensitiveString {
if let Some(existing) = self.credential_interner.get(s) {
existing.clone()
} else {
let shared = SensitiveString::from(s);
self.credential_interner.insert(shared.clone());
shared
}
}
pub(crate) fn intern_metadata(&mut self, s: &str) -> Arc<str> {
if let Some(intern) = self.static_intern.as_ref() {
if let Some(arc) = intern.lookup(s) {
return arc;
}
}
if let Some(existing) = self.metadata_interner.get(s) {
return existing.clone();
}
let shared: Arc<str> = Arc::from(s);
self.metadata_interner.insert(shared.clone());
shared
}
pub(crate) fn with_static_intern(intern: Arc<crate::static_intern::StaticInterner>) -> Self {
Self {
static_intern: Some(intern),
..Self::default()
}
}
#[cfg(feature = "ml")]
pub(crate) fn push_detector_ml_pending(
&mut self,
raw_match: keyhog_core::RawMatch,
heuristic_conf: f64,
code_context: crate::context::CodeContext,
context_multiplier: f64,
context_suppression_threshold: Option<f64>,
post_match: keyhog_core::DetectorPostMatchConfidenceSpec,
ml_features: [f32; crate::ml_scorer::NUM_FEATURES],
ml_weight: f64,
min_confidence_floor: f64,
is_named_detector: bool,
is_generic_detector: bool,
allow_canonical_hex_key: bool,
allow_encoded_text_lift: bool,
checksum: crate::checksum::ChecksumConfidenceDecision,
ml_mode: crate::detector_ml_policy::ActiveMlMode,
) -> bool {
self.push_ml_pending(MlPendingMatch::detector_candidate(
raw_match,
heuristic_conf,
code_context,
context_multiplier,
context_suppression_threshold,
post_match,
ml_features,
ml_weight,
min_confidence_floor,
is_named_detector,
is_generic_detector,
allow_canonical_hex_key,
allow_encoded_text_lift,
checksum,
ml_mode,
))
}
#[cfg(all(feature = "ml", feature = "entropy"))]
pub(crate) fn push_entropy_ml_pending(
&mut self,
raw_match: keyhog_core::RawMatch,
heuristic_conf: f64,
context_multiplier: f64,
context_suppression_threshold: Option<f64>,
post_match: keyhog_core::DetectorPostMatchConfidenceSpec,
ml_features: [f32; crate::ml_scorer::NUM_FEATURES],
ml_weight: f64,
min_confidence_floor: f64,
allow_canonical_hex_key: bool,
checksum: crate::checksum::ChecksumConfidenceDecision,
ml_mode: crate::detector_ml_policy::ActiveMlMode,
) -> bool {
self.push_ml_pending(MlPendingMatch::entropy_candidate(
raw_match,
heuristic_conf,
context_multiplier,
context_suppression_threshold,
post_match,
ml_features,
ml_weight,
min_confidence_floor,
allow_canonical_hex_key,
checksum,
ml_mode,
))
}
#[cfg(feature = "ml")]
fn push_ml_pending(&mut self, candidate: MlPendingMatch) -> bool {
let identity = PendingMatchIdentity::from(&candidate);
if let Some(&index) = self.ml_pending_index.get(&identity) {
let existing = &mut self.ml_pending[index];
if same_raw_match_identity(&candidate.raw_match, &existing.raw_match)
&& candidate.has_same_execution_as(existing)
{
if candidate.raw_match < existing.raw_match {
*existing = candidate;
return true;
}
return false;
}
}
let index = self.ml_pending.len();
self.ml_pending.push(candidate);
self.ml_pending_index.insert(identity, index);
true
}
#[cfg(feature = "ml")]
pub(crate) fn take_ml_pending(&mut self) -> Vec<MlPendingMatch> {
self.ml_pending_index.clear();
std::mem::take(&mut self.ml_pending)
}
#[cfg(all(feature = "ml", feature = "entropy"))]
pub(crate) fn for_each_pre_entropy_pending_ml_line<F>(&self, mut visit: F)
where
F: FnMut(Option<usize>),
{
for pending in &self.ml_pending {
visit(pending.raw_match.location.line);
}
}
pub(crate) fn for_each_produced_match<F>(&self, mut visit: F)
where
F: FnMut(&keyhog_core::RawMatch),
{
for found in &self.matches {
visit(found);
}
#[cfg(feature = "ml")]
for pending in &self.ml_pending {
visit(&pending.raw_match);
}
}
pub(crate) fn push_match(&mut self, m: keyhog_core::RawMatch, limit: usize) -> bool {
let identity = OwnedMatchIdentity::from(&m);
if self.claimed_match_identities.contains(&identity) {
return self.replace_claimed_match_if_better(&identity, m);
}
if self.matches.len() < limit {
self.claimed_match_identities.insert(identity);
self.matches.push(m);
return true;
}
if let Some(mut worst) = self.matches.peek_mut() {
if m < *worst {
let displaced = OwnedMatchIdentity::from(&*worst);
*worst = m;
drop(worst);
self.claimed_match_identities.remove(&displaced);
self.claimed_match_identities.insert(identity);
return true;
}
}
false
}
fn replace_claimed_match_if_better(
&mut self,
identity: &OwnedMatchIdentity,
candidate: keyhog_core::RawMatch,
) -> bool {
let should_replace = self
.matches
.iter()
.any(|existing| identity.matches_raw(existing) && candidate < *existing);
if !should_replace {
return false;
}
let mut data = std::mem::take(&mut self.matches).into_vec();
let idx = data
.iter()
.position(|existing| identity.matches_raw(existing))
.expect("identity in claimed_match_identities implies heap entry");
data[idx] = candidate;
self.matches = BinaryHeap::from(data);
true
}
#[cfg(any(feature = "entropy", test))]
fn claimed_priority_would_replace(
&self,
identity: &OwnedMatchIdentity,
priority: &RawMatchPriority<'_>,
) -> bool {
self.matches
.iter()
.find(|existing| identity.matches_raw(existing))
.is_none_or(|existing| !priority.cmp_raw_match(existing).is_gt())
}
#[cfg(any(feature = "entropy", test))]
pub(crate) fn push_match_lazy<F>(
&mut self,
priority: RawMatchPriority<'_>,
limit: usize,
build: F,
) where
F: FnOnce(&mut Self) -> keyhog_core::RawMatch,
{
if limit == 0 {
return;
}
if self.matches.len() >= limit
&& self
.matches
.peek()
.is_some_and(|worst| priority.cmp_raw_match(worst).is_gt())
{
return;
}
let identity = OwnedMatchIdentity::from_priority(&priority);
if self.claimed_match_identities.contains(&identity) {
if !self.claimed_priority_would_replace(&identity, &priority) {
return;
}
let m = build(self);
self.replace_claimed_match_if_better(&identity, m);
return;
}
if self.matches.len() < limit {
let m = build(self);
self.claimed_match_identities.insert(identity);
self.matches.push(m);
return;
}
let m = build(self);
if let Some(mut worst) = self.matches.peek_mut() {
if m < *worst {
let displaced = OwnedMatchIdentity::from(&*worst);
*worst = m;
drop(worst);
self.claimed_match_identities.remove(&displaced);
self.claimed_match_identities.insert(identity);
}
}
}
pub(crate) fn into_matches(self) -> Vec<keyhog_core::RawMatch> {
let mut matches: Vec<_> = self.matches.into_iter().collect();
if matches.len() <= 1 {
return matches;
}
matches.sort_unstable_by(|a, b| raw_match_identity_cmp(a, b).then_with(|| a.cmp(b)));
matches.dedup_by(|a, b| same_raw_match_identity(a, b));
matches.sort_unstable();
matches
}
}