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