use std::net::IpAddr;
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 is_ipv6_literal(&self) -> bool {
matches!(self.host.parse::<IpAddr>(), Ok(IpAddr::V6(_)))
}
fn family_then_kind_rank(&self) -> (u8, u8) {
let family_rank = u8::from(!self.is_ipv6_literal());
(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);
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProviderRecord {
pub content_key: String,
pub provider_peer_id: String,
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())
}
}
#[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 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"
);
}
}