use std::net::{IpAddr, SocketAddr};
use dig_ip::Family;
use serde::{Deserialize, Serialize};
use dig_nat::PeerId;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AddressKind {
Direct,
Mapped,
Reflexive,
Relay,
}
impl AddressKind {
pub fn rank(self) -> u8 {
match self {
AddressKind::Direct => 0,
AddressKind::Mapped => 1,
AddressKind::Reflexive => 2,
AddressKind::Relay => 3,
}
}
pub fn is_dialable(self) -> bool {
!matches!(self, AddressKind::Relay)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CandidateAddr {
pub host: String,
pub port: u16,
pub kind: AddressKind,
}
impl CandidateAddr {
pub fn direct(host: impl Into<String>, port: u16) -> Self {
CandidateAddr {
host: host.into(),
port,
kind: AddressKind::Direct,
}
}
pub fn relay_marker() -> Self {
CandidateAddr {
host: String::new(),
port: 0,
kind: AddressKind::Relay,
}
}
fn family_rank(&self) -> u8 {
let family = self
.host
.parse::<IpAddr>()
.ok()
.map(|ip| Family::of(&SocketAddr::new(ip, self.port)));
match family {
Some(Family::V6) => 0,
Some(Family::V4) => 1,
None => 2,
}
}
fn dial_identity(&self) -> (String, u16) {
let host = match self.host.parse::<IpAddr>() {
Ok(IpAddr::V6(v6)) => v6
.to_ipv4_mapped()
.map(IpAddr::V4)
.unwrap_or(IpAddr::V6(v6))
.to_string(),
Ok(ip) => ip.to_string(),
Err(_) => self.host.clone(),
};
(host, self.port)
}
fn is_ipv6_literal(&self) -> bool {
self.family_rank() == 0
}
fn family_then_kind_rank(&self) -> (u8, u8) {
(self.family_rank(), self.kind.rank())
}
}
pub(crate) fn sort_addresses_ipv6_first(addresses: &mut [CandidateAddr]) {
addresses.sort_by_key(CandidateAddr::family_then_kind_rank);
}
pub const MAX_ADDRESSES_PER_RECORD: usize = 8;
pub(crate) fn sort_and_cap_addresses(addresses: &mut Vec<CandidateAddr>) {
sort_addresses_ipv6_first(addresses);
addresses.truncate(MAX_ADDRESSES_PER_RECORD);
}
pub const MAX_DIAL_CANDIDATES: usize = 4;
pub fn dial_candidates(addresses: &[CandidateAddr]) -> Vec<&CandidateAddr> {
let mut candidates: Vec<&CandidateAddr> =
addresses.iter().filter(|a| a.kind.is_dialable()).collect();
candidates.sort_by_key(|a| a.family_then_kind_rank());
let mut seen = std::collections::HashSet::new();
candidates.retain(|a| seen.insert(a.dial_identity()));
reserve_fallback_slot_and_cap(&mut candidates);
candidates
}
fn reserve_fallback_slot_and_cap(candidates: &mut Vec<&CandidateAddr>) {
if candidates.len() <= MAX_DIAL_CANDIDATES {
return;
}
let kept_excludes_every_fallback = candidates[..MAX_DIAL_CANDIDATES]
.iter()
.all(|a| a.is_ipv6_literal());
let fallback = kept_excludes_every_fallback
.then(|| candidates.iter().find(|a| !a.is_ipv6_literal()).copied())
.flatten();
match fallback {
Some(fallback) => {
candidates.truncate(MAX_DIAL_CANDIDATES - 1);
candidates.push(fallback);
}
None => candidates.truncate(MAX_DIAL_CANDIDATES),
}
}
pub(crate) fn deserialize_capped_addresses<'de, D>(
deserializer: D,
) -> Result<Vec<CandidateAddr>, D::Error>
where
D: serde::Deserializer<'de>,
{
let mut addresses = Vec::<CandidateAddr>::deserialize(deserializer)?;
sort_and_cap_addresses(&mut addresses);
Ok(addresses)
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProviderRecord {
pub content_key: String,
pub provider_peer_id: String,
#[serde(deserialize_with = "deserialize_capped_addresses")]
pub addresses: Vec<CandidateAddr>,
pub expires_at: u64,
}
impl ProviderRecord {
pub fn new(
content_key: &crate::key::Key,
provider: &PeerId,
mut addresses: Vec<CandidateAddr>,
expires_at: u64,
) -> Self {
sort_and_cap_addresses(&mut addresses);
ProviderRecord {
content_key: content_key.to_hex(),
provider_peer_id: provider.to_hex(),
addresses,
expires_at,
}
}
pub fn provider_peer_id(&self) -> Option<PeerId> {
PeerId::from_hex(&self.provider_peer_id)
}
pub fn is_expired(&self, now: u64) -> bool {
now >= self.expires_at
}
pub fn best_address(&self) -> Option<&CandidateAddr> {
self.addresses.iter().find(|a| a.kind.is_dialable())
}
pub fn dial_candidates(&self) -> Vec<&CandidateAddr> {
dial_candidates(&self.addresses)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::key::Key;
fn pid(b: u8) -> PeerId {
PeerId::from_bytes([b; 32])
}
#[test]
fn record_round_trips_through_json() {
let key = Key::from_bytes([0xAB; 32]);
let rec = ProviderRecord::new(
&key,
&pid(0x07),
vec![CandidateAddr::direct("203.0.113.7", 9444)],
1_000,
);
let json = serde_json::to_string(&rec).unwrap();
let back: ProviderRecord = serde_json::from_str(&json).unwrap();
assert_eq!(rec, back);
assert_eq!(back.provider_peer_id().unwrap(), pid(0x07));
assert_eq!(back.content_key, key.to_hex());
}
#[test]
fn ttl_expiry() {
let rec = ProviderRecord::new(&Key::from_bytes([0u8; 32]), &pid(1), vec![], 100);
assert!(!rec.is_expired(99));
assert!(rec.is_expired(100));
assert!(rec.is_expired(101));
}
#[test]
fn address_kind_wire_tokens_are_lowercase() {
assert_eq!(
serde_json::to_string(&AddressKind::Direct).unwrap(),
"\"direct\""
);
assert_eq!(
serde_json::to_string(&AddressKind::Reflexive).unwrap(),
"\"reflexive\""
);
assert_eq!(
serde_json::to_string(&AddressKind::Mapped).unwrap(),
"\"mapped\""
);
assert_eq!(
serde_json::to_string(&AddressKind::Relay).unwrap(),
"\"relay\""
);
}
#[test]
fn best_address_prefers_most_direct() {
let key = Key::from_bytes([0u8; 32]);
let rec = ProviderRecord::new(
&key,
&pid(1),
vec![
CandidateAddr {
host: "r".into(),
port: 1,
kind: AddressKind::Reflexive,
},
CandidateAddr::direct("d", 2),
CandidateAddr::relay_marker(),
],
10,
);
assert_eq!(rec.best_address().unwrap().kind, AddressKind::Direct);
}
#[test]
fn best_address_none_when_only_relay() {
let key = Key::from_bytes([0u8; 32]);
let rec = ProviderRecord::new(&key, &pid(1), vec![CandidateAddr::relay_marker()], 10);
assert!(rec.best_address().is_none());
}
#[test]
fn address_rank_ordering() {
assert!(AddressKind::Direct.rank() < AddressKind::Mapped.rank());
assert!(AddressKind::Mapped.rank() < AddressKind::Reflexive.rank());
assert!(AddressKind::Reflexive.rank() < AddressKind::Relay.rank());
assert!(!AddressKind::Relay.is_dialable());
assert!(AddressKind::Direct.is_dialable());
}
#[test]
fn provider_record_new_sorts_addresses_ipv6_first() {
let key = Key::from_bytes([0u8; 32]);
let rec = ProviderRecord::new(
&key,
&pid(1),
vec![
CandidateAddr::direct("203.0.113.7", 9444), CandidateAddr::direct("2001:db8::1", 9444), CandidateAddr {
host: "198.51.100.2".into(),
port: 1,
kind: AddressKind::Reflexive,
}, CandidateAddr {
host: "2001:db8::2".into(),
port: 1,
kind: AddressKind::Reflexive,
}, ],
10,
);
let hosts: Vec<&str> = rec.addresses.iter().map(|a| a.host.as_str()).collect();
assert_eq!(
hosts,
vec!["2001:db8::1", "2001:db8::2", "203.0.113.7", "198.51.100.2"],
"addresses must be IPv6-first, then ranked by AddressKind"
);
}
#[test]
fn family_key_derives_from_dig_ip_family() {
let key = Key::from_bytes([0u8; 32]);
let rec = ProviderRecord::new(
&key,
&pid(1),
vec![
CandidateAddr::direct("::ffff:203.0.113.9", 9444), CandidateAddr::direct("2001:db8::1", 9444), ],
10,
);
let hosts: Vec<&str> = rec.addresses.iter().map(|a| a.host.as_str()).collect();
assert_eq!(
hosts,
vec!["2001:db8::1", "::ffff:203.0.113.9"],
"an IPv4-mapped IPv6 address must sort as V4 (dig_ip::Family), after a genuine IPv6"
);
}
#[test]
fn directness_kind_rank_preserved_as_tiebreak_within_a_family() {
let key = Key::from_bytes([0u8; 32]);
let rec = ProviderRecord::new(
&key,
&pid(1),
vec![
CandidateAddr {
host: "2001:db8::3".into(),
port: 1,
kind: AddressKind::Reflexive,
},
CandidateAddr {
host: "2001:db8::2".into(),
port: 1,
kind: AddressKind::Mapped,
},
CandidateAddr::direct("2001:db8::1", 9444),
],
10,
);
let hosts: Vec<&str> = rec.addresses.iter().map(|a| a.host.as_str()).collect();
assert_eq!(
hosts,
vec!["2001:db8::1", "2001:db8::2", "2001:db8::3"],
"within one family, addresses must stay ordered by AddressKind::rank (most-direct first)"
);
}
#[test]
fn best_address_prefers_ipv6_over_ipv4_at_same_rank() {
let key = Key::from_bytes([0u8; 32]);
let rec = ProviderRecord::new(
&key,
&pid(1),
vec![
CandidateAddr::direct("203.0.113.7", 9444), CandidateAddr::direct("2001:db8::1", 9444), ],
10,
);
assert_eq!(rec.best_address().unwrap().host, "2001:db8::1");
}
#[test]
fn provider_record_new_caps_addresses_at_the_constant() {
let key = Key::from_bytes([0u8; 32]);
let many: Vec<CandidateAddr> = (0..1000)
.map(|i| CandidateAddr::direct(format!("203.0.113.{}", i % 255), 9444))
.collect();
let rec = ProviderRecord::new(&key, &pid(1), many, 10);
assert_eq!(rec.addresses.len(), MAX_ADDRESSES_PER_RECORD);
}
#[test]
fn provider_record_new_cap_keeps_most_preferred_after_sort() {
let key = Key::from_bytes([0u8; 32]);
let mut addrs: Vec<CandidateAddr> = Vec::new();
addrs.push(CandidateAddr::direct("2001:db8::1", 9444));
for i in 0..1000u32 {
addrs.push(CandidateAddr {
host: format!("198.51.100.{}", i % 255),
port: 1,
kind: AddressKind::Relay,
});
}
let rec = ProviderRecord::new(&key, &pid(1), addrs, 10);
assert_eq!(rec.addresses.len(), MAX_ADDRESSES_PER_RECORD);
assert_eq!(
rec.addresses[0].host, "2001:db8::1",
"the single most-preferred (IPv6 direct) candidate must survive truncation"
);
}
fn record_json_with_addresses(n: usize) -> String {
let addrs: Vec<String> = (0..n)
.map(|i| {
format!(
r#"{{"host":"198.51.100.{}","port":1,"kind":"relay"}}"#,
i % 255
)
})
.collect();
format!(
r#"{{"content_key":"{}","provider_peer_id":"{}","addresses":[{}],"expires_at":1}}"#,
"aa".repeat(32),
"bb".repeat(32),
addrs.join(",")
)
}
#[test]
fn deserialization_bounds_the_address_count() {
let rec: ProviderRecord = serde_json::from_str(&record_json_with_addresses(1000)).unwrap();
assert_eq!(rec.addresses.len(), MAX_ADDRESSES_PER_RECORD);
}
#[test]
fn deserialization_bound_is_one_off_exact() {
let at_cap: ProviderRecord =
serde_json::from_str(&record_json_with_addresses(MAX_ADDRESSES_PER_RECORD)).unwrap();
assert_eq!(at_cap.addresses.len(), MAX_ADDRESSES_PER_RECORD);
let over_by_one: ProviderRecord =
serde_json::from_str(&record_json_with_addresses(MAX_ADDRESSES_PER_RECORD + 1))
.unwrap();
assert_eq!(over_by_one.addresses.len(), MAX_ADDRESSES_PER_RECORD);
}
#[test]
fn deserialization_keeps_the_most_preferred_addresses() {
let mut addrs: Vec<String> =
vec![r#"{"host":"2001:db8::1","port":9444,"kind":"direct"}"#.to_string()];
for i in 0..1000 {
addrs.push(format!(
r#"{{"host":"198.51.100.{}","port":1,"kind":"relay"}}"#,
i % 255
));
}
addrs.rotate_left(1);
let json = format!(
r#"{{"content_key":"{}","provider_peer_id":"{}","addresses":[{}],"expires_at":1}}"#,
"aa".repeat(32),
"bb".repeat(32),
addrs.join(",")
);
let rec: ProviderRecord = serde_json::from_str(&json).unwrap();
assert_eq!(rec.addresses.len(), MAX_ADDRESSES_PER_RECORD);
assert_eq!(
rec.addresses[0].host, "2001:db8::1",
"the most-preferred candidate must survive the bound regardless of wire position"
);
}
fn record_with(addresses: Vec<CandidateAddr>) -> ProviderRecord {
ProviderRecord::new(&Key::from_bytes([0u8; 32]), &pid(1), addresses, 10)
}
#[test]
fn dial_candidates_order_v6_then_v4_then_unresolvable() {
let rec = record_with(vec![
CandidateAddr::direct("not-a-literal", 9444),
CandidateAddr::direct("203.0.113.7", 9444),
CandidateAddr::direct("2001:db8::1", 9444),
]);
let hosts: Vec<&str> = rec
.dial_candidates()
.iter()
.map(|a| a.host.as_str())
.collect();
assert_eq!(
hosts,
vec!["2001:db8::1", "203.0.113.7", "not-a-literal"],
"dial order is IPv6, then IPv4, then anything unresolvable (§5.2)"
);
}
#[test]
fn dial_candidates_keep_the_ipv4_fallback_behind_an_ipv6_candidate() {
let rec = record_with(vec![
CandidateAddr::direct("2001:db8::1", 9444),
CandidateAddr::direct("172.31.79.22", 9444),
]);
let candidates = rec.dial_candidates();
assert_eq!(candidates.len(), 2, "the fallback must not be dropped");
assert_eq!(candidates[0].host, "2001:db8::1");
assert_eq!(candidates[1].host, "172.31.79.22");
}
#[test]
fn dial_candidates_treat_v4_mapped_v6_as_ipv4() {
let rec = record_with(vec![
CandidateAddr::direct("::ffff:203.0.113.9", 9444),
CandidateAddr::direct("2001:db8::1", 9444),
]);
let hosts: Vec<&str> = rec
.dial_candidates()
.iter()
.map(|a| a.host.as_str())
.collect();
assert_eq!(hosts, vec!["2001:db8::1", "::ffff:203.0.113.9"]);
}
#[test]
fn dial_candidates_exclude_relay_markers() {
let rec = record_with(vec![
CandidateAddr::relay_marker(),
CandidateAddr::direct("2001:db8::1", 9444),
]);
let candidates = rec.dial_candidates();
assert_eq!(
candidates.len(),
1,
"a relay marker is not directly dialable"
);
assert_eq!(candidates[0].host, "2001:db8::1");
}
#[test]
fn dial_candidates_are_bounded_and_deduped() {
let mut addresses = vec![CandidateAddr::direct("2001:db8::1", 9444); 3];
addresses.extend((0..5).map(|i| CandidateAddr::direct(format!("10.0.0.{i}"), 9444)));
let rec = record_with(addresses);
let candidates = rec.dial_candidates();
assert_eq!(candidates.len(), MAX_DIAL_CANDIDATES);
assert_eq!(
candidates
.iter()
.filter(|a| a.host == "2001:db8::1")
.count(),
1,
"a repeated host:port contributes exactly one dial attempt"
);
}
#[test]
fn dial_candidates_of_a_relay_only_record_are_empty() {
let rec = record_with(vec![CandidateAddr::relay_marker()]);
assert!(rec.dial_candidates().is_empty());
}
#[test]
fn unresolvable_host_sorts_after_an_ipv4_literal_in_the_stored_order() {
let rec = record_with(vec![
CandidateAddr::direct("not-a-literal", 1),
CandidateAddr::direct("203.0.113.7", 1),
]);
let hosts: Vec<&str> = rec.addresses.iter().map(|a| a.host.as_str()).collect();
assert_eq!(hosts, vec!["203.0.113.7", "not-a-literal"]);
}
#[test]
fn dial_candidates_reserve_a_slot_for_the_ipv4_fallback() {
let rec = record_with(vec![
CandidateAddr::direct("2001:db8::1", 9444),
CandidateAddr::direct("2001:db8::2", 9444),
CandidateAddr::direct("2001:db8::3", 9444),
CandidateAddr::direct("2001:db8::4", 9444),
CandidateAddr::direct("203.0.113.7", 9444),
]);
let candidates = rec.dial_candidates();
assert_eq!(candidates.len(), MAX_DIAL_CANDIDATES);
assert!(
candidates.iter().any(|a| a.host == "203.0.113.7"),
"the IPv4 fallback tier must keep a slot inside the cap, got {:?}",
candidates.iter().map(|a| &a.host).collect::<Vec<_>>()
);
assert_eq!(
candidates[0].host, "2001:db8::1",
"IPv6 still leads — the reservation costs the LEAST preferred v6 slot, not the order"
);
}
#[test]
fn dial_candidates_reserve_the_fallback_only_when_it_would_be_lost() {
let rec = record_with(vec![
CandidateAddr::direct("2001:db8::1", 9444),
CandidateAddr::direct("2001:db8::2", 9444),
CandidateAddr::direct("2001:db8::3", 9444),
CandidateAddr::direct("203.0.113.7", 9444),
]);
let hosts: Vec<&str> = rec
.dial_candidates()
.iter()
.map(|a| a.host.as_str())
.collect();
assert_eq!(
hosts,
vec!["2001:db8::1", "2001:db8::2", "2001:db8::3", "203.0.113.7"]
);
}
#[test]
fn dial_candidates_dedupe_equivalent_spellings_of_one_address() {
let rec = record_with(vec![
CandidateAddr::direct("2001:db8::1", 9444),
CandidateAddr::direct("2001:0db8::1", 9444),
CandidateAddr::direct("2001:db8:0:0:0:0:0:1", 9444),
CandidateAddr::direct("2001:DB8::1", 9444),
CandidateAddr::direct("203.0.113.7", 9444),
]);
let candidates = rec.dial_candidates();
assert_eq!(
candidates.len(),
2,
"four spellings of one IPv6 address are ONE dial attempt, got {:?}",
candidates.iter().map(|a| &a.host).collect::<Vec<_>>()
);
assert!(candidates.iter().any(|a| a.host == "203.0.113.7"));
}
#[test]
fn dial_candidates_treat_a_v4_mapped_spelling_as_the_same_address_as_its_ipv4() {
let rec = record_with(vec![
CandidateAddr::direct("::ffff:203.0.113.7", 9444),
CandidateAddr::direct("203.0.113.7", 9444),
]);
assert_eq!(rec.dial_candidates().len(), 1);
}
#[test]
fn dial_candidates_keep_distinct_ports_of_one_host_apart() {
let rec = record_with(vec![
CandidateAddr::direct("2001:db8::1", 9444),
CandidateAddr::direct("2001:db8::1", 9445),
]);
assert_eq!(rec.dial_candidates().len(), 2);
}
}