Skip to main content

dig_ip/
candidate.rs

1//! Family-tagged peer candidate addresses aggregated from every discovery source.
2//!
3//! CLAUDE.md ยง5.2 requires that peer addresses be "stored/returned with the family recorded so the
4//! preference is explicit". A DIG node learns a peer's candidate addresses from many places โ€” relay
5//! introduction, PEX, the DHT, DNS (AAAA/A), STUN reflexive discovery, advertised listen addresses,
6//! prior successful dials โ€” and must aggregate ALL of them ("use as many methods as available"),
7//! tag each with its family, and dedup. [`PeerCandidates`] is that aggregate.
8
9use std::collections::{BTreeSet, HashSet};
10use std::net::SocketAddr;
11
12use crate::family::Family;
13
14/// Where a candidate address was learned โ€” provenance, kept for observability and so a future policy
15/// could weight sources. It does not affect the family-intersection rule (that is family-only).
16#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
17pub enum CandidateSource {
18    /// Introduced by the relay during rendezvous.
19    RelayIntroduction,
20    /// Learned via peer exchange (PEX).
21    Pex,
22    /// Learned from the distributed hash table.
23    Dht,
24    /// A DNS AAAA (IPv6) record.
25    DnsAAAA,
26    /// A DNS A (IPv4) record.
27    DnsA,
28    /// A server-reflexive address discovered via STUN.
29    StunReflexive,
30    /// An address the peer advertised as a listen address.
31    ListenAddr,
32    /// An address that succeeded on a prior dial.
33    PriorDial,
34}
35
36/// A single family-tagged candidate address for a peer.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct Candidate {
39    /// The address to dial.
40    pub addr: SocketAddr,
41    /// The address family (derived from `addr`; stored so preference is explicit end-to-end).
42    pub family: Family,
43    /// Where this address was learned.
44    pub source: CandidateSource,
45}
46
47impl Candidate {
48    /// Build a candidate, deriving its family from the address.
49    pub fn new(addr: SocketAddr, source: CandidateSource) -> Candidate {
50        Candidate {
51            addr,
52            family: Family::of(&addr),
53            source,
54        }
55    }
56}
57
58/// A peer's aggregated, family-tagged, de-duplicated candidate addresses.
59///
60/// Addresses are kept in insertion (discovery) order within each family so the caller's
61/// source ordering is preserved; a duplicate address keeps the FIRST source it was seen from.
62#[derive(Debug, Clone, Default, PartialEq, Eq)]
63pub struct PeerCandidates {
64    candidates: Vec<Candidate>,
65    seen: HashSet<SocketAddr>,
66}
67
68impl PeerCandidates {
69    /// An empty candidate set.
70    pub fn new() -> PeerCandidates {
71        PeerCandidates::default()
72    }
73
74    /// Add one discovered address (family derived, deduplicated). Returns `true` if it was new.
75    pub fn add(&mut self, addr: SocketAddr, source: CandidateSource) -> bool {
76        if !self.seen.insert(addr) {
77            return false;
78        }
79        self.candidates.push(Candidate::new(addr, source));
80        true
81    }
82
83    /// Add many addresses from one discovery source (dedup applies).
84    pub fn extend<I>(&mut self, addrs: I, source: CandidateSource)
85    where
86        I: IntoIterator<Item = SocketAddr>,
87    {
88        for addr in addrs {
89            self.add(addr, source);
90        }
91    }
92
93    /// The families this peer offers at least one candidate on.
94    pub fn families(&self) -> BTreeSet<Family> {
95        self.candidates.iter().map(|c| c.family).collect()
96    }
97
98    /// The candidate addresses of `family`, in discovery order.
99    pub fn of_family(&self, family: Family) -> impl Iterator<Item = SocketAddr> + '_ {
100        self.candidates
101            .iter()
102            .filter(move |c| c.family == family)
103            .map(|c| c.addr)
104    }
105
106    /// Every candidate, in discovery order (family-tagged).
107    pub fn all(&self) -> &[Candidate] {
108        &self.candidates
109    }
110
111    /// Whether this peer has no candidates at all.
112    pub fn is_empty(&self) -> bool {
113        self.candidates.is_empty()
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    fn sa(s: &str) -> SocketAddr {
122        s.parse().unwrap()
123    }
124
125    #[test]
126    fn derives_family_on_add() {
127        let mut p = PeerCandidates::new();
128        p.add(sa("[2001:db8::1]:443"), CandidateSource::Dht);
129        p.add(sa("203.0.113.1:443"), CandidateSource::DnsA);
130        assert_eq!(p.families(), BTreeSet::from([Family::V6, Family::V4]));
131    }
132
133    #[test]
134    fn dedups_and_keeps_first_source() {
135        let mut p = PeerCandidates::new();
136        assert!(p.add(sa("203.0.113.1:443"), CandidateSource::Pex));
137        assert!(!p.add(sa("203.0.113.1:443"), CandidateSource::Dht));
138        assert_eq!(p.all().len(), 1);
139        assert_eq!(p.all()[0].source, CandidateSource::Pex);
140    }
141
142    #[test]
143    fn of_family_preserves_discovery_order() {
144        let mut p = PeerCandidates::new();
145        p.extend(
146            [sa("[2001:db8::2]:443"), sa("[2001:db8::1]:443")],
147            CandidateSource::ListenAddr,
148        );
149        let v6: Vec<_> = p.of_family(Family::V6).collect();
150        assert_eq!(v6, vec![sa("[2001:db8::2]:443"), sa("[2001:db8::1]:443")]);
151        assert!(p.of_family(Family::V4).next().is_none());
152    }
153
154    #[test]
155    fn v4_mapped_v6_aggregates_as_v4() {
156        let mut p = PeerCandidates::new();
157        p.add(
158            sa("[::ffff:203.0.113.9]:443"),
159            CandidateSource::StunReflexive,
160        );
161        assert_eq!(p.families(), BTreeSet::from([Family::V4]));
162    }
163}