use super::*; pub(crate) fn auth_probe_record_failure_for_testing( shared: &ProxySharedState, peer_ip: IpAddr, now: Instant, ) { auth_probe_record_failure_in(shared, peer_ip, now); } pub(crate) fn auth_probe_fail_streak_for_testing_in_shared( shared: &ProxySharedState, peer_ip: IpAddr, ) -> Option { let peer_ip = normalize_auth_probe_ip(peer_ip); shared .handshake .auth_probe .get(&peer_ip) .map(|entry| entry.fail_streak) } pub(crate) fn clear_auth_probe_state_for_testing_in_shared(shared: &ProxySharedState) { shared.handshake.auth_probe.clear(); shared.handshake.auth_probe_slots.reset_for_testing(); match shared.handshake.auth_probe_saturation.lock() { Ok(mut saturation) => { *saturation = None; } Err(poisoned) => { let mut saturation = poisoned.into_inner(); *saturation = None; shared.handshake.auth_probe_saturation.clear_poison(); } } } pub(crate) fn auth_probe_state_for_testing_in_shared( shared: &ProxySharedState, ) -> &DashMap { &shared.handshake.auth_probe } pub(crate) fn auth_probe_slots_for_testing_in_shared(shared: &ProxySharedState) -> usize { shared.handshake.auth_probe_slots.used() } pub(crate) fn auth_probe_saturation_state_for_testing_in_shared( shared: &ProxySharedState, ) -> &Mutex> { &shared.handshake.auth_probe_saturation } pub(crate) fn auth_probe_saturation_state_lock_for_testing_in_shared( shared: &ProxySharedState, ) -> std::sync::MutexGuard<'_, Option> { shared .handshake .auth_probe_saturation .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } pub(crate) fn clear_unknown_sni_warn_state_for_testing_in_shared(shared: &ProxySharedState) { let mut guard = shared .handshake .unknown_sni_warn_next_allowed .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); *guard = None; } pub(crate) fn should_emit_unknown_sni_warn_for_testing_in_shared( shared: &ProxySharedState, now: Instant, ) -> bool { should_emit_unknown_sni_warn_in(shared, now) } pub(crate) fn clear_warned_secrets_for_testing_in_shared(shared: &ProxySharedState) { if let Ok(mut guard) = shared.handshake.invalid_secret_warned.lock() { guard.clear(); } } pub(crate) fn warned_secrets_for_testing_in_shared( shared: &ProxySharedState, ) -> &Mutex> { &shared.handshake.invalid_secret_warned } pub(crate) fn auth_probe_is_throttled_for_testing_in_shared( shared: &ProxySharedState, peer_ip: IpAddr, ) -> bool { auth_probe_is_throttled_in(shared, peer_ip, Instant::now()) } pub(crate) fn auth_probe_saturation_is_throttled_for_testing_in_shared( shared: &ProxySharedState, ) -> bool { auth_probe_saturation_is_throttled_in(shared, Instant::now()) } pub(crate) fn auth_probe_saturation_is_throttled_at_for_testing_in_shared( shared: &ProxySharedState, now: Instant, ) -> bool { auth_probe_saturation_is_throttled_in(shared, now) } #[test] fn parallel_distinct_failures_respect_exact_auth_probe_capacity() { const ATTEMPTS: usize = 10_000; let shared = ProxySharedState::new(); std::thread::scope(|scope| { for worker in 0..16 { let shared = Arc::clone(&shared); scope.spawn(move || { for index in (worker..ATTEMPTS).step_by(16) { let octets = (index as u32).to_be_bytes(); let peer_ip = IpAddr::V4(std::net::Ipv4Addr::new( octets[1], octets[2], octets[3], worker as u8, )); auth_probe_record_failure_in(shared.as_ref(), peer_ip, Instant::now()); } }); } }); assert_eq!(shared.handshake.auth_probe.len(), AUTH_PROBE_TRACK_MAX_ENTRIES); assert_eq!( auth_probe_slots_for_testing_in_shared(shared.as_ref()), AUTH_PROBE_TRACK_MAX_ENTRIES ); }