use crate::PeerId;
use crate::dht::core_engine::{DhtCoreEngine, NodeInfo};
use crate::security::IPDiversityConfig;
use std::time::Instant;
fn make_node_with_id(id_bytes: [u8; 32], addr: &str) -> NodeInfo {
NodeInfo {
id: PeerId::from_bytes(id_bytes),
addresses: vec![addr.parse().unwrap()],
last_seen: Instant::now(),
}
}
fn bucket0_id(seq: u8) -> [u8; 32] {
let mut id = [0u8; 32];
id[0] = 0x80;
id[31] = seq;
id
}
#[tokio::test]
async fn test_ip_diversity_enforcement_ipv6() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
for i in 1..=5u8 {
let node = make_node_with_id(bucket0_id(i), &format!("/ip6/2001:db8::{i}/udp/9000/quic"));
engine.add_node_no_trust(node).await?;
}
let node6 = make_node_with_id(bucket0_id(6), "/ip6/2001:db8::6/udp/9000/quic");
let result = engine.add_node_no_trust(node6).await;
assert!(result.is_err());
assert!(
result.unwrap_err().to_string().contains("IP diversity:"),
"Error should indicate IP diversity limits"
);
Ok(())
}
#[tokio::test]
async fn test_ip_diversity_enforcement_ipv4() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
let node1 = make_node_with_id(bucket0_id(1), "/ip4/192.168.1.1/udp/9000/quic");
engine.add_node_no_trust(node1).await?;
let node2 = make_node_with_id(bucket0_id(2), "/ip4/192.168.1.1/udp/9001/quic");
engine.add_node_no_trust(node2).await?;
let node3 = make_node_with_id(bucket0_id(3), "/ip4/192.168.1.1/udp/9002/quic");
let result = engine.add_node_no_trust(node3).await;
assert!(result.is_err());
assert!(
result.unwrap_err().to_string().contains("IP diversity:"),
"Error should indicate IP diversity limits"
);
Ok(())
}
#[tokio::test]
async fn test_ipv4_subnet_24_limit() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
for i in 1..=5u8 {
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(i),
&format!("/ip4/192.168.1.{i}/udp/9000/quic"),
))
.await?;
}
let node6 = make_node_with_id(bucket0_id(6), "/ip4/192.168.1.6/udp/9000/quic");
let result = engine.add_node_no_trust(node6).await;
assert!(result.is_err());
assert!(
result.unwrap_err().to_string().contains("IP diversity:"),
"Error should indicate IP diversity limits"
);
Ok(())
}
#[tokio::test]
async fn test_mixed_ipv4_ipv6_enforcement() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(1),
"/ip4/192.168.1.1/udp/9000/quic",
))
.await?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(2),
"/ip6/2001:db8::1/udp/9000/quic",
))
.await?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(3),
"/ip4/192.168.1.1/udp/9001/quic",
))
.await?;
let result_v4 = engine
.add_node_no_trust(make_node_with_id(
bucket0_id(4),
"/ip4/192.168.1.1/udp/9002/quic",
))
.await;
assert!(result_v4.is_err());
for i in 5..=8u8 {
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(i),
&format!("/ip6/2001:db8::{i}/udp/9000/quic"),
))
.await?;
}
let result_v6 = engine
.add_node_no_trust(make_node_with_id(
bucket0_id(9),
"/ip6/2001:db8::9/udp/9000/quic",
))
.await;
assert!(result_v6.is_err());
Ok(())
}
#[tokio::test]
async fn test_ipv4_ip_override_raises_limit() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine.set_ip_diversity_config(IPDiversityConfig {
max_per_ip: Some(3),
max_per_subnet: Some(usize::MAX),
..IPDiversityConfig::default()
});
for i in 1..=3u8 {
let node = make_node_with_id(
bucket0_id(i),
&format!("/ip4/192.168.1.1/udp/{}/quic", 9000 + u16::from(i)),
);
engine.add_node_no_trust(node).await?;
}
let node4 = make_node_with_id(bucket0_id(4), "/ip4/192.168.1.1/udp/9003/quic");
let result = engine.add_node_no_trust(node4).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("IP diversity:"));
Ok(())
}
#[tokio::test]
async fn test_ipv4_subnet_override_lowers_limit() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine.set_ip_diversity_config(IPDiversityConfig {
max_per_subnet: Some(1),
..IPDiversityConfig::default()
});
let node1 = make_node_with_id(bucket0_id(1), "/ip4/10.0.1.1/udp/9000/quic");
engine.add_node_no_trust(node1).await?;
let node2 = make_node_with_id(bucket0_id(2), "/ip4/10.0.1.2/udp/9000/quic");
let result = engine.add_node_no_trust(node2).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("IP diversity:"));
Ok(())
}
#[tokio::test]
async fn test_ipv6_subnet_override_raises_limit() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine.set_ip_diversity_config(IPDiversityConfig {
max_per_subnet: Some(8),
..IPDiversityConfig::default()
});
for i in 1..=8u8 {
let node = make_node_with_id(bucket0_id(i), &format!("/ip6/2001:db8::{i}/udp/9000/quic"));
engine.add_node_no_trust(node).await?;
}
let node9 = make_node_with_id(bucket0_id(9), "/ip6/2001:db8::9/udp/9000/quic");
let result = engine.add_node_no_trust(node9).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("IP diversity:"));
Ok(())
}
#[tokio::test]
async fn test_ipv6_subnet_override_lowers_limit() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine.set_ip_diversity_config(IPDiversityConfig {
max_per_subnet: Some(1),
..IPDiversityConfig::default()
});
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(1),
"/ip6/2001:db8::1/udp/9000/quic",
))
.await?;
let node2 = make_node_with_id(bucket0_id(2), "/ip6/2001:db8::2/udp/9000/quic");
let result = engine.add_node_no_trust(node2).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("IP diversity:"));
Ok(())
}
#[tokio::test]
async fn test_no_override_uses_defaults() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(1),
"/ip4/192.168.1.1/udp/9000/quic",
))
.await?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(2),
"/ip4/192.168.1.1/udp/9001/quic",
))
.await?;
let node3 = make_node_with_id(bucket0_id(3), "/ip4/192.168.1.1/udp/9002/quic");
let result = engine.add_node_no_trust(node3).await;
assert!(result.is_err());
Ok(())
}
#[tokio::test]
async fn test_trust_protects_peer_from_swap() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine.set_ip_diversity_config(IPDiversityConfig::default());
let mut id_far = [0u8; 32];
id_far[0] = 0xFF; engine
.add_node_no_trust(make_node_with_id(id_far, "/ip4/10.0.1.1/udp/9000/quic"))
.await?;
let mut id_mid = [0u8; 32];
id_mid[0] = 0xFE;
engine
.add_node_no_trust(make_node_with_id(id_mid, "/ip4/10.0.1.1/udp/9001/quic"))
.await?;
let mut id_close = [0u8; 32];
id_close[0] = 0x80; let far_peer = PeerId::from_bytes(id_far);
let trust_fn = |peer_id: &PeerId| -> f64 {
if *peer_id == far_peer {
0.8 } else {
0.5 }
};
let result = engine
.add_node(
make_node_with_id(id_close, "/ip4/10.0.1.1/udp/9002/quic"),
&trust_fn,
)
.await;
assert!(result.is_err());
assert!(engine.has_node(&far_peer).await);
assert!(engine.has_node(&PeerId::from_bytes(id_mid)).await);
Ok(())
}
#[tokio::test]
async fn test_untrusted_peer_can_be_swapped() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine.set_ip_diversity_config(IPDiversityConfig::default());
let mut id_far = [0u8; 32];
id_far[0] = 0xFF;
engine
.add_node_no_trust(make_node_with_id(id_far, "/ip4/10.0.1.1/udp/9000/quic"))
.await?;
let mut id_mid = [0u8; 32];
id_mid[0] = 0xFE;
engine
.add_node_no_trust(make_node_with_id(id_mid, "/ip4/10.0.1.1/udp/9001/quic"))
.await?;
let mut id_close = [0u8; 32];
id_close[0] = 0x80;
let far_peer = PeerId::from_bytes(id_far);
let trust_fn = |peer_id: &PeerId| -> f64 {
if *peer_id == far_peer {
0.3 } else {
0.5
}
};
let result = engine
.add_node(
make_node_with_id(id_close, "/ip4/10.0.1.1/udp/9002/quic"),
&trust_fn,
)
.await;
assert!(result.is_ok());
assert!(engine.has_node(&PeerId::from_bytes(id_close)).await);
assert!(!engine.has_node(&far_peer).await);
assert!(engine.has_node(&PeerId::from_bytes(id_mid)).await);
Ok(())
}
#[tokio::test]
async fn test_self_insertion_rejected() -> anyhow::Result<()> {
let self_id = PeerId::from_bytes([0u8; 32]);
let mut engine = DhtCoreEngine::new_for_tests(self_id)?;
let self_node = NodeInfo {
id: self_id,
addresses: vec!["/ip4/10.0.0.1/udp/9000/quic".parse().unwrap()],
last_seen: Instant::now(),
};
let result = engine.add_node_no_trust(self_node).await;
assert!(result.is_err());
assert!(
result.unwrap_err().to_string().contains("cannot add self"),
"expected self-insertion rejection"
);
Ok(())
}
#[tokio::test]
async fn test_ipv4_mapped_ipv6_counts_as_ipv4() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(1),
"/ip4/192.168.1.1/udp/9000/quic",
))
.await?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(2),
"/ip4/192.168.1.1/udp/9001/quic",
))
.await?;
let node3 = make_node_with_id(bucket0_id(3), "/ip6/::ffff:192.168.1.1/udp/9002/quic");
let result = engine.add_node_no_trust(node3).await;
assert!(
result.is_err(),
"IPv4-mapped IPv6 should be treated as IPv4 and hit the exact-IP limit"
);
Ok(())
}
#[tokio::test]
async fn test_ipv6_exact_ip_limit() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(1),
"/ip6/2001:db8::1/udp/9000/quic",
))
.await?;
engine
.add_node_no_trust(make_node_with_id(
bucket0_id(2),
"/ip6/2001:db8::1/udp/9001/quic",
))
.await?;
let node3 = make_node_with_id(bucket0_id(3), "/ip6/2001:db8::1/udp/9002/quic");
let result = engine.add_node_no_trust(node3).await;
assert!(result.is_err());
assert!(
result.unwrap_err().to_string().contains("exact-IP"),
"expected exact-IP rejection for IPv6"
);
Ok(())
}
#[tokio::test]
async fn test_farther_candidate_cannot_swap() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
let mut id1 = [0u8; 32];
id1[0] = 0x80;
id1[31] = 0x01; let mut id2 = [0u8; 32];
id2[0] = 0x80;
id2[31] = 0x02;
engine
.add_node_no_trust(make_node_with_id(id1, "/ip4/10.0.1.1/udp/9000/quic"))
.await?;
engine
.add_node_no_trust(make_node_with_id(id2, "/ip4/10.0.1.1/udp/9001/quic"))
.await?;
let mut id_far = [0u8; 32];
id_far[0] = 0xFF;
let result = engine
.add_node_no_trust(make_node_with_id(id_far, "/ip4/10.0.1.1/udp/9002/quic"))
.await;
assert!(
result.is_err(),
"farther candidate should not be able to swap in"
);
assert!(
result.unwrap_err().to_string().contains("IP diversity:"),
"expected IP diversity rejection"
);
assert!(engine.has_node(&PeerId::from_bytes(id1)).await);
assert!(engine.has_node(&PeerId::from_bytes(id2)).await);
Ok(())
}
fn id_in_bucket(bucket: usize, seq: u8) -> [u8; 32] {
let mut id = [0u8; 32];
let byte_idx = bucket / 8;
let bit_idx = 7 - (bucket % 8);
id[byte_idx] = 1 << bit_idx;
id[31] |= seq; id
}
#[tokio::test]
async fn test_close_group_ip_diversity_rejects_excess() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
for i in 0..5u8 {
let bucket = 255 - (i as usize); let id = id_in_bucket(bucket, 0);
engine
.add_node_no_trust(make_node_with_id(
id,
&format!("/ip4/10.0.1.{}/udp/9000/quic", i + 1),
))
.await?;
}
let id6 = id_in_bucket(250, 0);
let result = engine
.add_node_no_trust(make_node_with_id(id6, "/ip4/10.0.1.6/udp/9000/quic"))
.await;
assert!(
result.is_err(),
"close-group /24 limit should reject 6th same-subnet peer"
);
assert!(
result.unwrap_err().to_string().contains("close-group"),
"expected close-group rejection"
);
Ok(())
}
#[tokio::test]
async fn test_close_group_swap_closer_evicts_farthest() -> anyhow::Result<()> {
let mut engine = DhtCoreEngine::new_for_tests(PeerId::from_bytes([0u8; 32]))?;
let mut peer_ids = Vec::new();
for i in 0..5u8 {
let bucket = 255 - (i as usize);
let id = id_in_bucket(bucket, 0);
peer_ids.push(id);
engine
.add_node_no_trust(make_node_with_id(
id,
&format!("/ip4/10.0.1.{}/udp/9000/quic", i + 1),
))
.await?;
}
let id_closer = id_in_bucket(253, 1); let farthest_id = PeerId::from_bytes(peer_ids[4]);
let result = engine
.add_node_no_trust(make_node_with_id(id_closer, "/ip4/10.0.1.7/udp/9000/quic"))
.await;
assert!(
result.is_ok(),
"closer same-subnet peer should swap in: {:?}",
result
);
assert!(engine.has_node(&PeerId::from_bytes(id_closer)).await);
assert!(
!engine.has_node(&farthest_id).await,
"farthest same-subnet peer should have been evicted from close group"
);
Ok(())
}