use super::*;
use crate::dht::routing::NodeId;
use std::net::{IpAddr, Ipv4Addr};
fn metrics() -> Arc<Metrics> {
Arc::new(Metrics::new())
}
fn peer(byte: u8) -> NodeId {
NodeId::from_bytes([byte; 32])
}
fn ip(a: u8, b: u8, c: u8, d: u8) -> IpAddr {
IpAddr::V4(Ipv4Addr::new(a, b, c, d))
}
fn adr005_default_cfg() -> ProbeRateLimitConfig {
ProbeRateLimitConfig::from(&decdn_common::config::ResolvedProbe::default())
}
#[test]
fn per_peer_burst_of_five_admits_five_then_rejects() {
let lim = ProbeRateLimiter::new(&adr005_default_cfg(), metrics());
let p = peer(1);
let i = Some(ip(10, 0, 0, 1));
for n in 0..5 {
assert_eq!(lim.check(&p, i), Ok(()), "probe {n} within burst");
}
assert_eq!(lim.check(&p, i), Err(ProbeRejectLayer::PerPeer));
}
#[test]
fn rejection_increments_probe_scrape_counter_not_dht() {
let mut cfg = adr005_default_cfg();
cfg.per_peer_burst = 1;
cfg.per_peer_rate_per_sec = 1e-9;
let metrics = metrics();
let lim = ProbeRateLimiter::new(&cfg, Arc::clone(&metrics));
let p = peer(1);
let i = Some(ip(10, 0, 0, 1));
assert_eq!(lim.check(&p, i), Ok(()));
assert_eq!(lim.check(&p, i), Err(ProbeRejectLayer::PerPeer));
let text = metrics.encode().unwrap();
assert!(
text.contains("decdn_probe_rate_limit_rejected_per_peer_total 1"),
"probe per-peer rejection must appear in /metrics; got:\n{text}"
);
assert!(
text.contains("decdn_dht_rate_limit_rejected_per_peer_total 0"),
"probe rejection must NOT bump the DHT counter; got:\n{text}"
);
}
#[test]
fn per_ip_rejection_drives_probe_per_ip_counter() {
let mut cfg = adr005_default_cfg();
cfg.per_ip_burst = 1;
cfg.per_ip_rate_per_sec = 1e-9;
let metrics = metrics();
let lim = ProbeRateLimiter::new(&cfg, Arc::clone(&metrics));
let i = Some(ip(10, 0, 0, 2));
assert_eq!(lim.check(&peer(3), i), Ok(()));
assert_eq!(lim.check(&peer(4), i), Err(ProbeRejectLayer::PerIp));
let text = metrics.encode().unwrap();
assert!(
text.contains("decdn_probe_rate_limit_rejected_per_ip_total 1"),
"probe per-IP rejection must drive the per-IP counter; got:\n{text}"
);
}
#[test]
fn global_rejection_drives_probe_global_counter() {
let mut cfg = adr005_default_cfg();
cfg.global_burst = 1;
cfg.global_rate_per_sec = 1e-9;
let metrics = metrics();
let lim = ProbeRateLimiter::new(&cfg, Arc::clone(&metrics));
assert_eq!(lim.check(&peer(1), Some(ip(10, 0, 0, 1))), Ok(()));
assert_eq!(
lim.check(&peer(2), Some(ip(10, 0, 0, 2))),
Err(ProbeRejectLayer::Global)
);
let text = metrics.encode().unwrap();
assert!(
text.contains("decdn_probe_rate_limit_rejected_global_total 1"),
"probe global rejection must drive the global counter; got:\n{text}"
);
}