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 spec.body_contains.is_some() || spec.body_not_contains.is_some() || spec.json_path.is_some()
71}
72
73pub(crate) fn resolve_live_verdict(is_live: bool, success_is_explicit: bool, body: &str) -> bool {
78 is_live && (success_is_explicit || !body_indicates_error(body))
79}
80
81static INFLIGHT_CAP_BYPASSES: AtomicUsize = AtomicUsize::new(0);
87static INFLIGHT_CAP_WARNED: std::sync::OnceLock<()> = std::sync::OnceLock::new();
88
89pub(crate) fn note_inflight_cap_bypass(max_inflight_keys: usize) -> usize {
90 let count = INFLIGHT_CAP_BYPASSES.fetch_add(1, Ordering::Relaxed) + 1;
91 if INFLIGHT_CAP_WARNED.set(()).is_ok() {
92 tracing::warn!(
93 max_inflight_keys,
94 "verifier inflight-dedup cap reached: verifying complete request identities \
95 WITHOUT the single-in-flight guard, so concurrent duplicate probes can hit the \
96 live API (rate-limit bans). Raise max_inflight_keys to restore dedup."
97 );
98 }
99 count
100}
101
102#[derive(Clone)]
103struct VerifyTaskShared {
104 global_semaphore: Arc<Semaphore>,
105 service_semaphores: Arc<HashMap<Arc<str>, Arc<Semaphore>>>,
106 max_concurrent_per_service: usize,
111 client: Client,
112 detectors: Arc<HashMap<Arc<str>, keyhog_core::DetectorSpec>>,
113 timeout: Duration,
114 cache: Arc<cache::VerificationCache>,
115 inflight: Arc<DashMap<cache::VerificationIdentity, Arc<Notify>>>,
116 inflight_count: Arc<AtomicUsize>,
117 max_inflight_keys: usize,
118 danger_allow_private_ips: bool,
119 danger_allow_http: bool,
120 insecure_tls: bool,
127 allow_script_verify: bool,
128 proxy_in_use: bool,
135 oob_session: Option<Arc<crate::oob::OobSession>>,
136}
137
138struct InflightGuard {
139 key: cache::VerificationIdentity,
140 inflight: Arc<DashMap<cache::VerificationIdentity, Arc<Notify>>>,
141 inflight_count: Arc<AtomicUsize>,
142 notify: Arc<Notify>,
143}
144
145impl Drop for InflightGuard {
146 fn drop(&mut self) {
147 self.inflight.remove(&self.key);
152 self.inflight_count.fetch_sub(1, Ordering::Release);
153 self.notify.notify_waiters();
154 }
155}
156
157fn try_reserve_inflight_slot(inflight_count: &AtomicUsize, max_inflight_keys: usize) -> bool {
158 let mut current = inflight_count.load(Ordering::Acquire);
159 loop {
160 if current >= max_inflight_keys {
161 return false;
162 }
163 match inflight_count.compare_exchange_weak(
164 current,
165 current + 1,
166 Ordering::AcqRel,
167 Ordering::Acquire,
168 ) {
169 Ok(_) => return true,
170 Err(observed) => current = observed,
171 }
172 }
173}
174
175async fn verify_group_task_safe(shared: VerifyTaskShared, group: DedupedMatch) -> VerifiedFinding {
176 let group_for_error = group.clone();
177 match std::panic::AssertUnwindSafe(verify_group_task(shared, group))
178 .catch_unwind()
179 .await
180 {
181 Ok(finding) => finding,
182 Err(e) => {
183 let reason = if let Some(s) = e.downcast_ref::<&str>() {
186 format!("verification task panicked: {s}")
187 } else if let Some(s) = e.downcast_ref::<String>() {
188 format!("verification task panicked: {s}")
189 } else {
190 "verification task panicked".to_string() };
192 tracing::error!(reason);
193 into_finding(
194 group_for_error,
195 VerificationResult::Error(reason),
196 HashMap::new(),
197 )
198 }
199 }
200}
201
202fn spawn_tracked_verify_task(
203 join_set: &mut JoinSet<VerifiedFinding>,
204 task_groups: &mut HashMap<TaskId, DedupedMatch>,
205 shared: VerifyTaskShared,
206 group: DedupedMatch,
207) {
208 let group_for_error = group.clone();
209 let abort_handle = join_set.spawn(verify_group_task_safe(shared, group));
210 task_groups.insert(abort_handle.id(), group_for_error);
211}
212
213fn finding_for_join_error(
214 join_error: JoinError,
215 task_groups: &mut HashMap<TaskId, DedupedMatch>,
216) -> Option<VerifiedFinding> {
217 let task_id = join_error.id();
218 tracing::error!(
219 %join_error,
220 %task_id,
221 "a verification task failed to join; preserving the credential group as a verification error"
222 );
223 match task_groups.remove(&task_id) {
224 Some(group) => Some(into_finding(
225 group,
226 VerificationResult::Error(format!("verification task failed to join: {join_error}")),
227 HashMap::new(),
228 )),
229 None => {
230 tracing::error!(
231 %task_id,
232 "a verification task failed to join but had no tracked credential group"
233 );
234 None
235 }
236 }
237}
238
239#[doc(hidden)]
240pub async fn tracked_join_error_preservation_for_test() -> Option<VerifiedFinding> {
241 let mut join_set = JoinSet::new();
242 let mut task_groups = HashMap::new();
243 let group = DedupedMatch {
244 detector_id: Arc::from("test-detector"),
245 detector_name: Arc::from("Test Detector"),
246 service: Arc::from("test-service"),
247 severity: Severity::High,
248 credential: SensitiveString::from("test-secret-for-join-error"),
249 credential_hash: CredentialHash::ZERO,
250 companions: HashMap::new(),
251 primary_location: MatchLocation {
252 source: Arc::from("test"),
253 file_path: Some(Arc::from("fixture.txt")),
254 line: Some(1),
255 offset: 0,
256 commit: None,
257 author: None,
258 date: None,
259 },
260 additional_locations: Vec::new(),
261 entropy: None,
262 confidence: Some(0.9),
263 };
264 let abort_handle = join_set.spawn(async { std::future::pending::<VerifiedFinding>().await });
265 task_groups.insert(abort_handle.id(), group);
266 abort_handle.abort();
267 match join_set.join_next_with_id().await {
268 Some(Err(join_error)) => finding_for_join_error(join_error, &mut task_groups),
269 _ => None,
270 }
271}
272
273async fn verify_group_task(shared: VerifyTaskShared, group: DedupedMatch) -> VerifiedFinding {
274 let global = shared.global_semaphore;
275 let service_sem = shared
276 .service_semaphores
277 .get(&*group.service)
278 .cloned()
279 .unwrap_or_else(|| Arc::new(Semaphore::new(shared.max_concurrent_per_service))); let client = shared.client;
281 let detector = shared.detectors.get(&*group.detector_id).cloned();
282 let timeout = shared.timeout;
283
284 let cache = shared.cache;
285 let inflight = shared.inflight;
286 let inflight_count = shared.inflight_count;
287 let max_inflight_keys = shared.max_inflight_keys;
288
289 let Ok(_global_permit) = global.acquire().await else {
290 return into_finding(
291 group,
292 VerificationResult::Error("semaphore closed".into()),
293 HashMap::new(),
294 );
295 };
296 let Ok(_service_permit) = service_sem.acquire().await else {
297 return into_finding(
298 group,
299 VerificationResult::Error("service semaphore closed".into()),
300 HashMap::new(),
301 );
302 };
303
304 let verification_identity =
305 cache::verification_identity(&group.credential, &group.detector_id, &group.companions);
306 if let Some((cached_result, cached_meta)) =
307 cache.get_with_companions(&group.credential, &group.detector_id, &group.companions)
308 {
309 return into_finding(group, cached_result, cached_meta);
310 }
311
312 let _inflight_guard = loop {
313 let notify_to_await: Arc<Notify> = {
316 let key = verification_identity.clone();
321 if let Some((cached_result, cached_meta)) =
322 cache.get_with_companions(&group.credential, &group.detector_id, &group.companions)
323 {
324 return into_finding(group, cached_result, cached_meta);
325 }
326
327 match inflight.entry(key.clone()) {
328 dashmap::mapref::entry::Entry::Occupied(entry) => entry.get().clone(),
329 dashmap::mapref::entry::Entry::Vacant(entry) => {
330 if !try_reserve_inflight_slot(&inflight_count, max_inflight_keys) {
331 note_inflight_cap_bypass(max_inflight_keys);
332 break None;
333 }
334 let notify = Arc::new(Notify::new());
335 entry.insert(notify.clone());
336 break Some(InflightGuard {
337 key,
338 inflight: inflight.clone(),
339 inflight_count: inflight_count.clone(),
340 notify,
341 });
342 }
343 }
344 };
345
346 notify_to_await.notified().await;
347 };
348
349 let (verification, metadata) = if let Some(verify_spec) = detector
350 .as_ref()
351 .and_then(|detector| detector.verify.as_ref())
352 {
353 verify_with_retry(
354 &client,
355 verify_spec,
356 &group.credential,
357 &group.companions,
358 timeout,
359 shared.danger_allow_private_ips,
360 shared.danger_allow_http,
361 shared.proxy_in_use,
362 shared.insecure_tls,
363 shared.allow_script_verify,
364 shared.oob_session.as_ref(),
365 )
366 .await
367 } else {
368 (VerificationResult::Unverifiable, HashMap::new())
369 };
370
371 if verification_result_is_cacheable(&verification) {
376 cache.put_with_companions(
377 &group.credential,
378 &group.detector_id,
379 &group.companions,
380 verification.clone(),
381 metadata.clone(),
382 );
383 }
384
385 into_finding(group, verification, metadata)
386}
387
388pub(crate) fn verification_result_is_cacheable(result: &VerificationResult) -> bool {
402 matches!(
403 result,
404 VerificationResult::Live
405 | VerificationResult::Revoked
406 | VerificationResult::Dead
407 | VerificationResult::Unverifiable
408 | VerificationResult::Skipped
409 )
410}
411
412impl VerificationEngine {
413 pub fn new(
415 detectors: &[keyhog_core::DetectorSpec],
416 config: VerifyConfig,
417 ) -> Result<Self, VerifyError> {
418 for detector in detectors {
419 let errors = keyhog_core::json_selector::validate_detector_response_selectors(detector);
420 if !errors.is_empty() {
421 return Err(VerifyError::DetectorConfig(format!(
422 "detector {:?}: {}",
423 detector.id,
424 errors.join("; ")
425 )));
426 }
427 }
428 let mut builder = crate::harden_verifier_client_builder(
435 Client::builder()
436 .timeout(config.timeout)
437 .danger_accept_invalid_certs(config.insecure_tls),
438 );
439 builder = crate::apply_proxy_config(builder, config.proxy.as_deref())
440 .map_err(VerifyError::ProxyConfig)?;
441 let client = builder.build().map_err(VerifyError::ClientBuild)?;
442
443 let detector_map: HashMap<Arc<str>, keyhog_core::DetectorSpec> = detectors
444 .iter()
445 .cloned()
446 .map(|mut detector| {
447 if let Some(verify) = detector.verify.as_mut() {
448 if verify.service.trim().is_empty() {
449 verify.service.clone_from(&detector.service);
450 }
451 }
452 (detector.id.clone().into(), detector)
453 })
454 .collect();
455
456 let mut service_semaphores = HashMap::new();
457 for d in detectors {
458 service_semaphores
459 .entry(d.service.clone().into())
460 .or_insert_with(|| {
461 Arc::new(Semaphore::new(config.max_concurrent_per_service.max(1)))
462 });
463 }
464
465 Ok(Self {
466 client,
467 detectors: Arc::new(detector_map),
468 service_semaphores: Arc::new(service_semaphores),
469 max_concurrent_per_service: config.max_concurrent_per_service.max(1),
470 global_semaphore: Arc::new(Semaphore::new(config.max_concurrent_global.max(1))),
471 timeout: config.timeout,
472 cache: Arc::new(cache::VerificationCache::default_ttl()),
473 inflight: Arc::new(DashMap::new()),
474 inflight_count: Arc::new(AtomicUsize::new(0)),
475 max_inflight_keys: config.max_inflight_keys.max(1),
476 danger_allow_private_ips: config.danger_allow_private_ips,
477 danger_allow_http: config.danger_allow_http,
478 insecure_tls: config.insecure_tls,
479 allow_script_verify: config.allow_script_verify,
480 proxy_in_use: crate::proxy_is_active(config.proxy.as_deref()),
493 oob_session: None,
494 })
495 }
496
497 pub async fn verify_all(&self, groups: Vec<DedupedMatch>) -> Vec<VerifiedFinding> {
499 let max_active = self.global_semaphore.available_permits().max(1);
500 let total = groups.len();
501 let shared = VerifyTaskShared {
502 global_semaphore: self.global_semaphore.clone(),
503 service_semaphores: self.service_semaphores.clone(),
504 max_concurrent_per_service: self.max_concurrent_per_service,
505 client: self.client.clone(),
506 detectors: self.detectors.clone(),
507 timeout: self.timeout,
508 cache: self.cache.clone(),
509 inflight: self.inflight.clone(),
510 inflight_count: self.inflight_count.clone(),
511 max_inflight_keys: self.max_inflight_keys,
512 danger_allow_private_ips: self.danger_allow_private_ips,
513 danger_allow_http: self.danger_allow_http,
514 insecure_tls: self.insecure_tls,
515 allow_script_verify: self.allow_script_verify,
516 proxy_in_use: self.proxy_in_use,
517 oob_session: self.oob_session.clone(),
518 };
519 let mut pending = groups.into_iter();
520 let mut join_set = JoinSet::new();
521 let mut task_groups = HashMap::new();
522
523 while join_set.len() < max_active {
524 let Some(group) = pending.next() else {
525 break;
526 };
527 spawn_tracked_verify_task(&mut join_set, &mut task_groups, shared.clone(), group);
528 }
529
530 let mut out = Vec::with_capacity(total);
531 while let Some(result) = join_set.join_next_with_id().await {
532 match result {
533 Ok((task_id, finding)) => {
534 task_groups.remove(&task_id);
535 out.push(finding);
536 }
537 Err(join_error) => {
538 if let Some(finding) = finding_for_join_error(join_error, &mut task_groups) {
539 out.push(finding);
540 }
541 }
542 }
543 if let Some(group) = pending.next() {
544 spawn_tracked_verify_task(&mut join_set, &mut task_groups, shared.clone(), group);
545 }
546 }
547 out
548 }
549
550 pub async fn enable_oob(
560 &mut self,
561 config: crate::oob::OobConfig,
562 ) -> Result<(), crate::oob::InteractshError> {
563 if let Some(old) = self.oob_session.take() {
564 old.shutdown().await;
565 }
566 let session = crate::oob::OobSession::start_with_network_policy(
567 self.client.clone(),
568 config,
569 self.timeout,
570 self.proxy_in_use,
571 self.insecure_tls,
572 )
573 .await?;
574 self.oob_session = Some(session);
575 Ok(())
576 }
577
578 pub async fn shutdown_oob(&mut self) {
581 if let Some(session) = self.oob_session.take() {
582 session.shutdown().await;
583 }
584 }
585}
586
587impl Drop for VerificationEngine {
588 fn drop(&mut self) {
589 if let Some(session) = self.oob_session.take() {
600 session.abort_poller_for_drop();
601 }
602 }
603}