1pub(crate) mod auth;
12pub(crate) mod aws;
13pub(crate) mod credential;
14mod multi_step;
15pub(crate) mod request;
16pub(crate) mod response;
17
18use std::collections::HashMap;
19use std::sync::atomic::{AtomicUsize, Ordering};
20use std::sync::Arc;
21use std::time::Duration;
22
23use dashmap::DashMap;
24use futures_util::FutureExt;
25use keyhog_core::{
26 CredentialHash, MatchLocation, SensitiveString, Severity, VerificationResult, VerifiedFinding,
27};
28use reqwest::Client;
29use tokio::sync::{Notify, Semaphore};
30use tokio::task::{Id as TaskId, JoinError, JoinSet};
31
32use crate::cache;
33use crate::{into_finding, DedupedMatch, VerificationEngine, VerifyConfig, VerifyError};
34
35pub(crate) use auth::script_auth_result;
36pub(crate) use aws::{
37 build_aws_probe, classify_aws_sts_failure, classify_aws_sts_http_200,
38 parse_aws_sts_success_metadata, valid_aws_format, validate_aws_region,
39};
40pub(crate) use credential::{
41 rate_limit_feedback_sequence_for_test, retry_delay_bounds_for_attempt,
42 retry_loop_preserves_metadata_on_exhaustion_for_test,
43 retry_loop_records_rate_limit_feedback_for_test, verify_with_retry, VerificationAttempt,
44};
45pub(crate) use multi_step::rate_limit_service_name as multi_step_rate_limit_service_name;
46pub(crate) use request::{
47 apply_header_body_templates, build_request_for_step, clear_pinned_client_cache_for_test,
48 missing_companion_error, pinned_client_cache_len_for_host_for_test,
49 pinned_client_cache_len_for_test, pinned_client_for_test, resolved_client_for_url,
50 ssrf_check_url_with_resolved_addrs_for_test, validate_header_body_templates,
51 validate_template_companions, RequestBuildResult,
52};
53pub(crate) use response::{
54 body_indicates_error, evaluate_success, execute_and_read_response, extract_metadata,
55 extract_provider_evidence,
56};
57
58pub(crate) fn retryable_http_status(status: u16) -> bool {
63 status == 429 || (500..=504).contains(&status)
64}
65
66pub(crate) fn success_spec_is_explicit(spec: &keyhog_core::SuccessSpec) -> bool {
71 match spec.policy {
72 Some(keyhog_core::SuccessPolicy::StatusAuthoritative) => true,
73 Some(keyhog_core::SuccessPolicy::BodyPositive) => {
74 spec.body_contains
75 .as_deref()
76 .is_some_and(|needle| !needle.is_empty())
77 || spec.json_path.is_some()
78 }
79 Some(keyhog_core::SuccessPolicy::StatusWithErrorBackstop) | None => false,
80 }
81}
82
83pub(crate) fn resolve_live_verdict(is_live: bool, success_is_explicit: bool, body: &str) -> bool {
88 is_live && (success_is_explicit || !body_indicates_error(body))
89}
90
91static INFLIGHT_CAP_BYPASSES: AtomicUsize = AtomicUsize::new(0);
97static INFLIGHT_CAP_WARNED: std::sync::OnceLock<()> = std::sync::OnceLock::new();
98
99pub(crate) fn note_inflight_cap_bypass(max_inflight_keys: usize) -> usize {
100 let count = INFLIGHT_CAP_BYPASSES.fetch_add(1, Ordering::Relaxed) + 1;
101 if INFLIGHT_CAP_WARNED.set(()).is_ok() {
102 tracing::warn!(
103 max_inflight_keys,
104 "verifier inflight-dedup cap reached: verifying complete request identities \
105 WITHOUT the single-in-flight guard, so concurrent duplicate probes can hit the \
106 live API (rate-limit bans). Raise max_inflight_keys to restore dedup."
107 );
108 }
109 count
110}
111
112#[derive(Clone)]
113struct VerifyTaskShared {
114 global_semaphore: Arc<Semaphore>,
115 service_semaphores: Arc<HashMap<Arc<str>, Arc<Semaphore>>>,
116 max_concurrent_per_service: usize,
121 client: Client,
122 detectors: Arc<HashMap<Arc<str>, keyhog_core::DetectorSpec>>,
123 timeout: Duration,
124 cache: Arc<cache::VerificationCache>,
125 inflight: Arc<DashMap<cache::VerificationIdentity, Arc<Notify>>>,
126 inflight_count: Arc<AtomicUsize>,
127 max_inflight_keys: usize,
128 danger_allow_private_ips: bool,
129 danger_allow_http: bool,
130 insecure_tls: bool,
137 allow_script_verify: bool,
138 proxy_in_use: bool,
145 oob_session: Option<Arc<crate::oob::OobSession>>,
146}
147
148struct InflightGuard {
149 key: cache::VerificationIdentity,
150 inflight: Arc<DashMap<cache::VerificationIdentity, Arc<Notify>>>,
151 inflight_count: Arc<AtomicUsize>,
152 notify: Arc<Notify>,
153}
154
155impl Drop for InflightGuard {
156 fn drop(&mut self) {
157 self.inflight.remove(&self.key);
162 self.inflight_count.fetch_sub(1, Ordering::Release);
163 self.notify.notify_waiters();
164 }
165}
166
167fn try_reserve_inflight_slot(inflight_count: &AtomicUsize, max_inflight_keys: usize) -> bool {
168 let mut current = inflight_count.load(Ordering::Acquire);
169 loop {
170 if current >= max_inflight_keys {
171 return false;
172 }
173 match inflight_count.compare_exchange_weak(
174 current,
175 current + 1,
176 Ordering::AcqRel,
177 Ordering::Acquire,
178 ) {
179 Ok(_) => return true,
180 Err(observed) => current = observed,
181 }
182 }
183}
184
185async fn verify_group_task_safe(
186 shared: Arc<VerifyTaskShared>,
187 group: DedupedMatch,
188) -> VerifiedFinding {
189 let group_for_error = group.clone();
190 match std::panic::AssertUnwindSafe(verify_group_task(shared, group))
191 .catch_unwind()
192 .await
193 {
194 Ok(finding) => finding,
195 Err(e) => {
196 let reason = if let Some(s) = e.downcast_ref::<&str>() {
199 format!("verification task panicked: {s}")
200 } else if let Some(s) = e.downcast_ref::<String>() {
201 format!("verification task panicked: {s}")
202 } else {
203 "verification task panicked".to_string() };
205 tracing::error!(reason);
206 into_finding(
207 group_for_error,
208 VerificationResult::Error(reason),
209 HashMap::new(),
210 )
211 }
212 }
213}
214
215fn spawn_tracked_verify_task(
216 join_set: &mut JoinSet<VerifiedFinding>,
217 task_groups: &mut HashMap<TaskId, DedupedMatch>,
218 shared: Arc<VerifyTaskShared>,
219 group: DedupedMatch,
220) {
221 let group_for_error = group.clone();
222 let abort_handle = join_set.spawn(keyhog_profile::instrument_future(
223 keyhog_profile::Stage::LiveVerification,
224 verify_group_task_safe(shared, group),
225 ));
226 task_groups.insert(abort_handle.id(), group_for_error);
227}
228
229fn finding_for_join_error(
230 join_error: JoinError,
231 task_groups: &mut HashMap<TaskId, DedupedMatch>,
232) -> Option<VerifiedFinding> {
233 let task_id = join_error.id();
234 tracing::error!(
235 %join_error,
236 %task_id,
237 "a verification task failed to join; preserving the credential group as a verification error"
238 );
239 match task_groups.remove(&task_id) {
240 Some(group) => Some(into_finding(
241 group,
242 VerificationResult::Error(format!("verification task failed to join: {join_error}")),
243 HashMap::new(),
244 )),
245 None => {
246 tracing::error!(
247 %task_id,
248 "a verification task failed to join but had no tracked credential group"
249 );
250 None
251 }
252 }
253}
254
255#[doc(hidden)]
256pub async fn tracked_join_error_preservation_for_test() -> Option<VerifiedFinding> {
257 let mut join_set = JoinSet::new();
258 let mut task_groups = HashMap::new();
259 let group = DedupedMatch {
260 detector_id: Arc::from("test-detector"),
261 detector_name: Arc::from("Test Detector"),
262 service: Arc::from("test-service"),
263 severity: Severity::High,
264 credential: SensitiveString::from("test-secret-for-join-error"),
265 credential_hash: CredentialHash::ZERO,
266 companions: HashMap::new(),
267 primary_location: MatchLocation {
268 source: Arc::from("test"),
269 file_path: Some(Arc::from("fixture.txt")),
270 line: Some(1),
271 offset: 0,
272 commit: None,
273 author: None,
274 date: None,
275 },
276 additional_locations: Vec::new(),
277 entropy: None,
278 confidence: Some(0.9),
279 };
280 let abort_handle = join_set.spawn(async { std::future::pending::<VerifiedFinding>().await });
281 task_groups.insert(abort_handle.id(), group);
282 abort_handle.abort();
283 match join_set.join_next_with_id().await {
284 Some(Err(join_error)) => finding_for_join_error(join_error, &mut task_groups),
285 _ => None,
286 }
287}
288
289async fn verify_group_task(shared: Arc<VerifyTaskShared>, group: DedupedMatch) -> VerifiedFinding {
290 let global = &shared.global_semaphore;
291 let service_sem = shared
292 .service_semaphores
293 .get(&*group.service)
294 .cloned()
295 .unwrap_or_else(|| Arc::new(Semaphore::new(shared.max_concurrent_per_service))); let client = &shared.client;
297 let detector = shared.detectors.get(&*group.detector_id).cloned();
298 let timeout = shared.timeout;
299
300 let cache = &shared.cache;
301 let inflight = &shared.inflight;
302 let inflight_count = &shared.inflight_count;
303 let max_inflight_keys = shared.max_inflight_keys;
304 let Ok(_global_permit) = keyhog_profile::instrument_future(
305 keyhog_profile::Stage::LiveVerification,
306 global.acquire(),
307 )
308 .await
309 else {
310 return into_finding(
311 group,
312 VerificationResult::Error("semaphore closed".into()),
313 HashMap::new(),
314 );
315 };
316 let Ok(_service_permit) = keyhog_profile::instrument_future(
317 keyhog_profile::Stage::LiveVerification,
318 service_sem.acquire(),
319 )
320 .await
321 else {
322 return into_finding(
323 group,
324 VerificationResult::Error("service semaphore closed".into()),
325 HashMap::new(),
326 );
327 };
328
329 let verification_identity =
330 cache::verification_identity(&group.credential, &group.detector_id, &group.companions);
331 if let Some((cached_result, cached_meta)) =
332 cache.get_with_companions(&group.credential, &group.detector_id, &group.companions)
333 {
334 return into_finding(group, cached_result, cached_meta);
335 }
336
337 let _inflight_guard = loop {
338 let notify_to_await: Arc<Notify> = {
341 let key = verification_identity.clone();
346 if let Some((cached_result, cached_meta)) =
347 cache.get_with_companions(&group.credential, &group.detector_id, &group.companions)
348 {
349 return into_finding(group, cached_result, cached_meta);
350 }
351
352 match inflight.entry(key.clone()) {
353 dashmap::mapref::entry::Entry::Occupied(entry) => entry.get().clone(),
354 dashmap::mapref::entry::Entry::Vacant(entry) => {
355 if !try_reserve_inflight_slot(inflight_count, max_inflight_keys) {
356 note_inflight_cap_bypass(max_inflight_keys);
357 break None;
358 }
359 let notify = Arc::new(Notify::new());
360 entry.insert(notify.clone());
361 break Some(InflightGuard {
362 key,
363 inflight: inflight.clone(),
364 inflight_count: inflight_count.clone(),
365 notify,
366 });
367 }
368 }
369 };
370
371 notify_to_await.notified().await;
372 };
373
374 let (verification, metadata) = if let Some(verify_spec) = detector
375 .as_ref()
376 .and_then(|detector| detector.verify.as_ref())
377 {
378 verify_with_retry(
379 client,
380 verify_spec,
381 &group.credential,
382 &group.companions,
383 timeout,
384 shared.danger_allow_private_ips,
385 shared.danger_allow_http,
386 shared.proxy_in_use,
387 shared.insecure_tls,
388 shared.allow_script_verify,
389 shared.oob_session.as_ref(),
390 )
391 .await
392 } else {
393 (VerificationResult::Unverifiable, HashMap::new())
394 };
395
396 if verification_result_is_cacheable(&verification) {
401 cache.put_with_companions(
402 &group.credential,
403 &group.detector_id,
404 &group.companions,
405 verification.clone(),
406 metadata.clone(),
407 );
408 }
409
410 into_finding(group, verification, metadata)
411}
412
413pub(crate) fn verification_result_is_cacheable(result: &VerificationResult) -> bool {
427 matches!(
428 result,
429 VerificationResult::Live
430 | VerificationResult::Revoked
431 | VerificationResult::Dead
432 | VerificationResult::Unverifiable
433 | VerificationResult::Skipped
434 )
435}
436
437impl VerificationEngine {
438 pub fn new(
440 detectors: &[keyhog_core::DetectorSpec],
441 config: VerifyConfig,
442 ) -> Result<Self, VerifyError> {
443 for detector in detectors {
444 let errors = keyhog_core::json_selector::validate_detector_response_selectors(detector);
445 if !errors.is_empty() {
446 return Err(VerifyError::DetectorConfig(format!(
447 "detector {:?}: {}",
448 detector.id,
449 errors.join("; ")
450 )));
451 }
452 }
453 let mut builder = crate::harden_verifier_client_builder(
460 Client::builder()
461 .timeout(config.timeout)
462 .danger_accept_invalid_certs(config.insecure_tls),
463 );
464 builder = crate::apply_proxy_config(builder, config.proxy.as_deref())
465 .map_err(VerifyError::ProxyConfig)?;
466 let client = builder.build().map_err(VerifyError::ClientBuild)?;
467
468 let detector_map: HashMap<Arc<str>, keyhog_core::DetectorSpec> = detectors
469 .iter()
470 .cloned()
471 .map(|mut detector| {
472 if let Some(verify) = detector.verify.as_mut() {
473 if verify.service.trim().is_empty() {
474 verify.service.clone_from(&detector.service);
475 }
476 }
477 (detector.id.clone().into(), detector)
478 })
479 .collect();
480
481 let mut service_semaphores = HashMap::new();
482 for d in detectors {
483 service_semaphores
484 .entry(d.service.clone().into())
485 .or_insert_with(|| {
486 Arc::new(Semaphore::new(config.max_concurrent_per_service.max(1)))
487 });
488 }
489
490 Ok(Self {
491 client,
492 detectors: Arc::new(detector_map),
493 service_semaphores: Arc::new(service_semaphores),
494 max_concurrent_per_service: config.max_concurrent_per_service.max(1),
495 global_semaphore: Arc::new(Semaphore::new(config.max_concurrent_global.max(1))),
496 timeout: config.timeout,
497 cache: Arc::new(cache::VerificationCache::default_ttl()),
498 inflight: Arc::new(DashMap::new()),
499 inflight_count: Arc::new(AtomicUsize::new(0)),
500 max_inflight_keys: config.max_inflight_keys.max(1),
501 danger_allow_private_ips: config.danger_allow_private_ips,
502 danger_allow_http: config.danger_allow_http,
503 insecure_tls: config.insecure_tls,
504 allow_script_verify: config.allow_script_verify,
505 proxy_in_use: crate::proxy_is_active(config.proxy.as_deref()),
518 oob_session: None,
519 })
520 }
521
522 pub async fn verify_all(&self, groups: Vec<DedupedMatch>) -> Vec<VerifiedFinding> {
524 let max_active = self.global_semaphore.available_permits().max(1);
525 let total = groups.len();
526 let shared = Arc::new(VerifyTaskShared {
527 global_semaphore: self.global_semaphore.clone(),
528 service_semaphores: self.service_semaphores.clone(),
529 max_concurrent_per_service: self.max_concurrent_per_service,
530 client: self.client.clone(),
531 detectors: self.detectors.clone(),
532 timeout: self.timeout,
533 cache: self.cache.clone(),
534 inflight: self.inflight.clone(),
535 inflight_count: self.inflight_count.clone(),
536 max_inflight_keys: self.max_inflight_keys,
537 danger_allow_private_ips: self.danger_allow_private_ips,
538 danger_allow_http: self.danger_allow_http,
539 insecure_tls: self.insecure_tls,
540 allow_script_verify: self.allow_script_verify,
541 proxy_in_use: self.proxy_in_use,
542 oob_session: self.oob_session.clone(),
543 });
544 let mut join_set = JoinSet::new();
545 let mut task_groups = HashMap::new();
546 let mut pending = groups.into_iter();
547
548 while join_set.len() < max_active {
549 keyhog_profile::record_annotation(
551 keyhog_profile::AnnotationId::QueueDepth,
552 (join_set.len() + pending.len()) as u64,
553 );
554 let Some(group) = pending.next() else {
555 break;
556 };
557 spawn_tracked_verify_task(&mut join_set, &mut task_groups, shared.clone(), group);
558 }
559
560 let mut out = Vec::with_capacity(total);
561 while let Some(result) = join_set.join_next_with_id().await {
562 match result {
563 Ok((task_id, finding)) => {
564 task_groups.remove(&task_id);
565 out.push(finding);
566 }
567 Err(join_error) => {
568 if let Some(finding) = finding_for_join_error(join_error, &mut task_groups) {
569 out.push(finding);
570 }
571 }
572 }
573 if let Some(group) = pending.next() {
574 keyhog_profile::record_annotation(
576 keyhog_profile::AnnotationId::QueueDepth,
577 (join_set.len() + pending.len()) as u64,
578 );
579 spawn_tracked_verify_task(&mut join_set, &mut task_groups, shared.clone(), group);
580 }
581 }
582 out
583 }
584
585 pub async fn enable_oob(
595 &mut self,
596 config: crate::oob::OobConfig,
597 ) -> Result<(), crate::oob::InteractshError> {
598 if let Some(old) = self.oob_session.take() {
599 old.shutdown().await;
600 }
601 let session = crate::oob::OobSession::start_with_network_policy(
602 self.client.clone(),
603 config,
604 self.timeout,
605 self.proxy_in_use,
606 self.insecure_tls,
607 )
608 .await?;
609 self.oob_session = Some(session);
610 Ok(())
611 }
612
613 pub async fn shutdown_oob(&mut self) {
616 if let Some(session) = self.oob_session.take() {
617 session.shutdown().await;
618 }
619 }
620}
621
622impl Drop for VerificationEngine {
623 fn drop(&mut self) {
624 if let Some(session) = self.oob_session.take() {
635 session.abort_poller_for_drop();
636 }
637 }
638}