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 evidence: keyhog_core::EvidenceVerdict::review_unattributed(),
280 };
281 let abort_handle = join_set.spawn(async { std::future::pending::<VerifiedFinding>().await });
282 task_groups.insert(abort_handle.id(), group);
283 abort_handle.abort();
284 match join_set.join_next_with_id().await {
285 Some(Err(join_error)) => finding_for_join_error(join_error, &mut task_groups),
286 _ => None,
287 }
288}
289
290async fn verify_group_task(shared: Arc<VerifyTaskShared>, group: DedupedMatch) -> VerifiedFinding {
291 let global = &shared.global_semaphore;
292 let service_sem = shared
293 .service_semaphores
294 .get(&*group.service)
295 .cloned()
296 .unwrap_or_else(|| Arc::new(Semaphore::new(shared.max_concurrent_per_service))); let client = &shared.client;
298 let detector = shared.detectors.get(&*group.detector_id).cloned();
299 let timeout = shared.timeout;
300
301 let cache = &shared.cache;
302 let inflight = &shared.inflight;
303 let inflight_count = &shared.inflight_count;
304 let max_inflight_keys = shared.max_inflight_keys;
305 let Ok(_global_permit) = keyhog_profile::instrument_future(
306 keyhog_profile::Stage::LiveVerification,
307 global.acquire(),
308 )
309 .await
310 else {
311 return into_finding(
312 group,
313 VerificationResult::Error("semaphore closed".into()),
314 HashMap::new(),
315 );
316 };
317 let Ok(_service_permit) = keyhog_profile::instrument_future(
318 keyhog_profile::Stage::LiveVerification,
319 service_sem.acquire(),
320 )
321 .await
322 else {
323 return into_finding(
324 group,
325 VerificationResult::Error("service semaphore closed".into()),
326 HashMap::new(),
327 );
328 };
329
330 let verification_identity =
331 cache::verification_identity(&group.credential, &group.detector_id, &group.companions);
332 if let Some((cached_result, cached_meta)) =
333 cache.get_with_companions(&group.credential, &group.detector_id, &group.companions)
334 {
335 return into_finding(group, cached_result, cached_meta);
336 }
337
338 let _inflight_guard = loop {
339 let notify_to_await: Arc<Notify> = {
342 let key = verification_identity.clone();
347 if let Some((cached_result, cached_meta)) =
348 cache.get_with_companions(&group.credential, &group.detector_id, &group.companions)
349 {
350 return into_finding(group, cached_result, cached_meta);
351 }
352
353 match inflight.entry(key.clone()) {
354 dashmap::mapref::entry::Entry::Occupied(entry) => entry.get().clone(),
355 dashmap::mapref::entry::Entry::Vacant(entry) => {
356 if !try_reserve_inflight_slot(inflight_count, max_inflight_keys) {
357 note_inflight_cap_bypass(max_inflight_keys);
358 break None;
359 }
360 let notify = Arc::new(Notify::new());
361 entry.insert(notify.clone());
362 break Some(InflightGuard {
363 key,
364 inflight: inflight.clone(),
365 inflight_count: inflight_count.clone(),
366 notify,
367 });
368 }
369 }
370 };
371
372 notify_to_await.notified().await;
373 };
374
375 let (verification, metadata) = if let Some(verify_spec) = detector
376 .as_ref()
377 .and_then(|detector| detector.verify.as_ref())
378 {
379 verify_with_retry(
380 client,
381 verify_spec,
382 &group.credential,
383 &group.companions,
384 timeout,
385 shared.danger_allow_private_ips,
386 shared.danger_allow_http,
387 shared.proxy_in_use,
388 shared.insecure_tls,
389 shared.allow_script_verify,
390 shared.oob_session.as_ref(),
391 )
392 .await
393 } else {
394 (VerificationResult::Unverifiable, HashMap::new())
395 };
396
397 if verification_result_is_cacheable(&verification) {
402 cache.put_with_companions(
403 &group.credential,
404 &group.detector_id,
405 &group.companions,
406 verification.clone(),
407 metadata.clone(),
408 );
409 }
410
411 into_finding(group, verification, metadata)
412}
413
414pub(crate) fn verification_result_is_cacheable(result: &VerificationResult) -> bool {
428 matches!(
429 result,
430 VerificationResult::Live
431 | VerificationResult::Revoked
432 | VerificationResult::Dead
433 | VerificationResult::Unverifiable
434 | VerificationResult::Skipped
435 )
436}
437
438impl VerificationEngine {
439 pub fn new(
441 detectors: &[keyhog_core::DetectorSpec],
442 config: VerifyConfig,
443 ) -> Result<Self, VerifyError> {
444 for detector in detectors {
445 let errors = keyhog_core::json_selector::validate_detector_response_selectors(detector);
446 if !errors.is_empty() {
447 return Err(VerifyError::DetectorConfig(format!(
448 "detector {:?}: {}",
449 detector.id,
450 errors.join("; ")
451 )));
452 }
453 }
454 let mut builder = crate::harden_verifier_client_builder(
461 Client::builder()
462 .timeout(config.timeout)
463 .danger_accept_invalid_certs(config.insecure_tls),
464 );
465 builder = crate::apply_proxy_config(builder, config.proxy.as_deref())
466 .map_err(VerifyError::ProxyConfig)?;
467 let client = builder.build().map_err(VerifyError::ClientBuild)?;
468
469 let detector_map: HashMap<Arc<str>, keyhog_core::DetectorSpec> = detectors
470 .iter()
471 .cloned()
472 .map(|mut detector| {
473 if let Some(verify) = detector.verify.as_mut() {
474 if verify.service.trim().is_empty() {
475 verify.service.clone_from(&detector.service);
476 }
477 }
478 (detector.id.clone().into(), detector)
479 })
480 .collect();
481
482 let mut service_semaphores = HashMap::new();
483 for d in detectors {
484 service_semaphores
485 .entry(d.service.clone().into())
486 .or_insert_with(|| {
487 Arc::new(Semaphore::new(config.max_concurrent_per_service.max(1)))
488 });
489 }
490
491 Ok(Self {
492 client,
493 detectors: Arc::new(detector_map),
494 service_semaphores: Arc::new(service_semaphores),
495 max_concurrent_per_service: config.max_concurrent_per_service.max(1),
496 global_semaphore: Arc::new(Semaphore::new(config.max_concurrent_global.max(1))),
497 timeout: config.timeout,
498 cache: Arc::new(cache::VerificationCache::default_ttl()),
499 inflight: Arc::new(DashMap::new()),
500 inflight_count: Arc::new(AtomicUsize::new(0)),
501 max_inflight_keys: config.max_inflight_keys.max(1),
502 danger_allow_private_ips: config.danger_allow_private_ips,
503 danger_allow_http: config.danger_allow_http,
504 insecure_tls: config.insecure_tls,
505 allow_script_verify: config.allow_script_verify,
506 proxy_in_use: crate::proxy_is_active(config.proxy.as_deref()),
519 oob_session: None,
520 })
521 }
522
523 pub async fn verify_all(&self, groups: Vec<DedupedMatch>) -> Vec<VerifiedFinding> {
525 let max_active = self.global_semaphore.available_permits().max(1);
526 let total = groups.len();
527 let shared = Arc::new(VerifyTaskShared {
528 global_semaphore: self.global_semaphore.clone(),
529 service_semaphores: self.service_semaphores.clone(),
530 max_concurrent_per_service: self.max_concurrent_per_service,
531 client: self.client.clone(),
532 detectors: self.detectors.clone(),
533 timeout: self.timeout,
534 cache: self.cache.clone(),
535 inflight: self.inflight.clone(),
536 inflight_count: self.inflight_count.clone(),
537 max_inflight_keys: self.max_inflight_keys,
538 danger_allow_private_ips: self.danger_allow_private_ips,
539 danger_allow_http: self.danger_allow_http,
540 insecure_tls: self.insecure_tls,
541 allow_script_verify: self.allow_script_verify,
542 proxy_in_use: self.proxy_in_use,
543 oob_session: self.oob_session.clone(),
544 });
545 let mut join_set = JoinSet::new();
546 let mut task_groups = HashMap::new();
547 let mut pending = groups.into_iter();
548
549 while join_set.len() < max_active {
550 keyhog_profile::record_annotation(
552 keyhog_profile::AnnotationId::QueueDepth,
553 (join_set.len() + pending.len()) as u64,
554 );
555 let Some(group) = pending.next() else {
556 break;
557 };
558 spawn_tracked_verify_task(&mut join_set, &mut task_groups, shared.clone(), group);
559 }
560
561 let mut out = Vec::with_capacity(total);
562 while let Some(result) = join_set.join_next_with_id().await {
563 match result {
564 Ok((task_id, finding)) => {
565 task_groups.remove(&task_id);
566 out.push(finding);
567 }
568 Err(join_error) => {
569 if let Some(finding) = finding_for_join_error(join_error, &mut task_groups) {
570 out.push(finding);
571 }
572 }
573 }
574 if let Some(group) = pending.next() {
575 keyhog_profile::record_annotation(
577 keyhog_profile::AnnotationId::QueueDepth,
578 (join_set.len() + pending.len()) as u64,
579 );
580 spawn_tracked_verify_task(&mut join_set, &mut task_groups, shared.clone(), group);
581 }
582 }
583 out
584 }
585
586 pub async fn enable_oob(
596 &mut self,
597 config: crate::oob::OobConfig,
598 ) -> Result<(), crate::oob::InteractshError> {
599 if let Some(old) = self.oob_session.take() {
600 old.shutdown().await;
601 }
602 let session = crate::oob::OobSession::start_with_network_policy(
603 self.client.clone(),
604 config,
605 self.timeout,
606 self.proxy_in_use,
607 self.insecure_tls,
608 )
609 .await?;
610 self.oob_session = Some(session);
611 Ok(())
612 }
613
614 pub async fn shutdown_oob(&mut self) {
617 if let Some(session) = self.oob_session.take() {
618 session.shutdown().await;
619 }
620 }
621}
622
623impl Drop for VerificationEngine {
624 fn drop(&mut self) {
625 if let Some(session) = self.oob_session.take() {
636 session.abort_poller_for_drop();
637 }
638 }
639}