1use std::collections::{BTreeSet, HashSet};
10use std::net::SocketAddr;
11
12use crate::family::Family;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
17pub enum CandidateSource {
18 RelayIntroduction,
20 Pex,
22 Dht,
24 DnsAAAA,
26 DnsA,
28 StunReflexive,
30 ListenAddr,
32 PriorDial,
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct Candidate {
39 pub addr: SocketAddr,
41 pub family: Family,
43 pub source: CandidateSource,
45}
46
47impl Candidate {
48 pub fn new(addr: SocketAddr, source: CandidateSource) -> Candidate {
50 Candidate {
51 addr,
52 family: Family::of(&addr),
53 source,
54 }
55 }
56}
57
58#[derive(Debug, Clone, Default, PartialEq, Eq)]
63pub struct PeerCandidates {
64 candidates: Vec<Candidate>,
65 seen: HashSet<SocketAddr>,
66}
67
68impl PeerCandidates {
69 pub fn new() -> PeerCandidates {
71 PeerCandidates::default()
72 }
73
74 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 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 pub fn families(&self) -> BTreeSet<Family> {
95 self.candidates.iter().map(|c| c.family).collect()
96 }
97
98 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 pub fn all(&self) -> &[Candidate] {
108 &self.candidates
109 }
110
111 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}