ant_quic/bootstrap_cache/
entry.rs1use crate::nat_traversal_api::PeerId;
4use serde::{Deserialize, Serialize};
5use std::collections::HashSet;
6use std::net::SocketAddr;
7use std::time::{Duration, SystemTime};
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct CachedPeer {
12 #[serde(with = "peer_id_serde")]
14 pub peer_id: PeerId,
15
16 pub addresses: Vec<SocketAddr>,
18
19 pub capabilities: PeerCapabilities,
21
22 pub first_seen: SystemTime,
24
25 pub last_seen: SystemTime,
27
28 pub last_attempt: Option<SystemTime>,
30
31 pub stats: ConnectionStats,
33
34 #[serde(default = "default_quality_score")]
36 pub quality_score: f64,
37
38 pub source: PeerSource,
40}
41
42fn default_quality_score() -> f64 {
43 0.5
44}
45
46#[derive(Debug, Clone, Default, Serialize, Deserialize)]
48pub struct PeerCapabilities {
49 pub supports_relay: bool,
51
52 pub supports_coordination: bool,
54
55 #[serde(default)]
57 pub protocols: HashSet<String>,
58
59 pub nat_type: Option<NatType>,
61
62 #[serde(default)]
64 pub external_addresses: Vec<SocketAddr>,
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
69pub enum NatType {
70 None,
72 FullCone,
74 AddressRestrictedCone,
76 PortRestrictedCone,
78 Symmetric,
80 Unknown,
82}
83
84#[derive(Debug, Clone, Default, Serialize, Deserialize)]
86pub struct ConnectionStats {
87 pub success_count: u32,
89
90 pub failure_count: u32,
92
93 pub avg_rtt_ms: u32,
95
96 pub min_rtt_ms: u32,
98
99 pub max_rtt_ms: u32,
101
102 pub bytes_relayed: u64,
104
105 pub coordinations_completed: u32,
107}
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
111pub enum PeerSource {
112 Seed,
114 Connection,
116 Relay,
118 Coordination,
120 Merge,
122 #[default]
124 Unknown,
125}
126
127#[derive(Debug, Clone)]
129pub struct ConnectionOutcome {
130 pub success: bool,
132 pub rtt_ms: Option<u32>,
134 pub capabilities_discovered: Option<PeerCapabilities>,
136}
137
138impl CachedPeer {
139 pub fn new(peer_id: PeerId, addresses: Vec<SocketAddr>, source: PeerSource) -> Self {
141 let now = SystemTime::now();
142 Self {
143 peer_id,
144 addresses,
145 capabilities: PeerCapabilities::default(),
146 first_seen: now,
147 last_seen: now,
148 last_attempt: None,
149 stats: ConnectionStats::default(),
150 quality_score: 0.5, source,
152 }
153 }
154
155 pub fn record_success(&mut self, rtt_ms: u32, caps: Option<PeerCapabilities>) {
157 self.last_seen = SystemTime::now();
158 self.last_attempt = Some(SystemTime::now());
159 self.stats.success_count = self.stats.success_count.saturating_add(1);
160
161 if self.stats.avg_rtt_ms == 0 {
163 self.stats.avg_rtt_ms = rtt_ms;
164 self.stats.min_rtt_ms = rtt_ms;
165 self.stats.max_rtt_ms = rtt_ms;
166 } else {
167 self.stats.avg_rtt_ms = (self.stats.avg_rtt_ms * 7 + rtt_ms) / 8;
168 self.stats.min_rtt_ms = self.stats.min_rtt_ms.min(rtt_ms);
169 self.stats.max_rtt_ms = self.stats.max_rtt_ms.max(rtt_ms);
170 }
171
172 if let Some(caps) = caps {
173 self.capabilities = caps;
174 }
175 }
176
177 pub fn record_failure(&mut self) {
179 self.last_attempt = Some(SystemTime::now());
180 self.stats.failure_count = self.stats.failure_count.saturating_add(1);
181 }
182
183 pub fn calculate_quality(&mut self, weights: &super::config::QualityWeights) {
185 let total_attempts = self.stats.success_count + self.stats.failure_count;
186
187 let success_rate = if total_attempts > 0 {
189 self.stats.success_count as f64 / total_attempts as f64
190 } else {
191 0.5 };
193
194 let rtt_score = if self.stats.avg_rtt_ms > 0 {
197 1.0 - (self.stats.avg_rtt_ms as f64 / 1000.0).min(1.0)
198 } else {
199 0.5 };
201
202 let age_secs = self
204 .last_seen
205 .duration_since(SystemTime::UNIX_EPOCH)
206 .ok()
207 .and_then(|last_seen_epoch| {
208 SystemTime::now()
209 .duration_since(SystemTime::UNIX_EPOCH)
210 .ok()
211 .map(|now_epoch| {
212 now_epoch
213 .as_secs()
214 .saturating_sub(last_seen_epoch.as_secs())
215 })
216 })
217 .unwrap_or(0) as f64;
218
219 let freshness = (-age_secs * 0.693 / 86400.0).exp();
221
222 let mut cap_bonus: f64 = 0.0;
224 if self.capabilities.supports_relay {
225 cap_bonus += 0.3;
226 }
227 if self.capabilities.supports_coordination {
228 cap_bonus += 0.3;
229 }
230 if matches!(
231 self.capabilities.nat_type,
232 Some(NatType::None) | Some(NatType::FullCone)
233 ) {
234 cap_bonus += 0.4; }
236 let cap_score = cap_bonus.min(1.0);
237
238 self.quality_score = (success_rate * weights.success_rate
240 + rtt_score * weights.rtt
241 + freshness * weights.freshness
242 + cap_score * weights.capabilities)
243 .clamp(0.0, 1.0);
244 }
245
246 pub fn is_stale(&self, threshold: Duration) -> bool {
248 self.last_seen
249 .elapsed()
250 .map(|age| age > threshold)
251 .unwrap_or(true)
252 }
253
254 pub fn success_rate(&self) -> f64 {
256 let total = self.stats.success_count + self.stats.failure_count;
257 if total == 0 {
258 0.5
259 } else {
260 self.stats.success_count as f64 / total as f64
261 }
262 }
263
264 pub fn merge_addresses(&mut self, other: &CachedPeer) {
266 for addr in &other.addresses {
267 if !self.addresses.contains(addr) {
268 self.addresses.push(*addr);
269 }
270 }
271 if self.addresses.len() > 10 {
273 self.addresses.truncate(10);
274 }
275 }
276}
277
278mod peer_id_serde {
280 use super::PeerId;
281 use serde::{Deserialize, Deserializer, Serialize, Serializer};
282
283 pub fn serialize<S>(peer_id: &PeerId, serializer: S) -> Result<S::Ok, S::Error>
284 where
285 S: Serializer,
286 {
287 hex::encode(peer_id.0).serialize(serializer)
288 }
289
290 pub fn deserialize<'de, D>(deserializer: D) -> Result<PeerId, D::Error>
291 where
292 D: Deserializer<'de>,
293 {
294 let s = String::deserialize(deserializer)?;
295 let bytes = hex::decode(&s).map_err(serde::de::Error::custom)?;
296 if bytes.len() != 32 {
297 return Err(serde::de::Error::custom("PeerId must be 32 bytes"));
298 }
299 let mut arr = [0u8; 32];
300 arr.copy_from_slice(&bytes);
301 Ok(PeerId(arr))
302 }
303}
304
305#[cfg(test)]
306mod tests {
307 use super::*;
308
309 #[test]
310 fn test_cached_peer_new() {
311 let peer_id = PeerId([1u8; 32]);
312 let peer = CachedPeer::new(
313 peer_id,
314 vec!["127.0.0.1:9000".parse().unwrap()],
315 PeerSource::Seed,
316 );
317
318 assert_eq!(peer.peer_id, peer_id);
319 assert_eq!(peer.addresses.len(), 1);
320 assert_eq!(peer.source, PeerSource::Seed);
321 assert!((peer.quality_score - 0.5).abs() < f64::EPSILON);
322 }
323
324 #[test]
325 fn test_record_success() {
326 let mut peer = CachedPeer::new(
327 PeerId([1u8; 32]),
328 vec!["127.0.0.1:9000".parse().unwrap()],
329 PeerSource::Seed,
330 );
331
332 peer.record_success(100, None);
333 assert_eq!(peer.stats.success_count, 1);
334 assert_eq!(peer.stats.avg_rtt_ms, 100);
335 assert_eq!(peer.stats.min_rtt_ms, 100);
336 assert_eq!(peer.stats.max_rtt_ms, 100);
337
338 peer.record_success(200, None);
339 assert_eq!(peer.stats.success_count, 2);
340 assert_eq!(peer.stats.avg_rtt_ms, 112);
342 assert_eq!(peer.stats.min_rtt_ms, 100);
343 assert_eq!(peer.stats.max_rtt_ms, 200);
344 }
345
346 #[test]
347 fn test_record_failure() {
348 let mut peer = CachedPeer::new(
349 PeerId([1u8; 32]),
350 vec!["127.0.0.1:9000".parse().unwrap()],
351 PeerSource::Seed,
352 );
353
354 peer.record_failure();
355 assert_eq!(peer.stats.failure_count, 1);
356 assert!(peer.last_attempt.is_some());
357 }
358
359 #[test]
360 fn test_success_rate() {
361 let mut peer = CachedPeer::new(
362 PeerId([1u8; 32]),
363 vec!["127.0.0.1:9000".parse().unwrap()],
364 PeerSource::Seed,
365 );
366
367 assert!((peer.success_rate() - 0.5).abs() < f64::EPSILON);
369
370 peer.record_success(100, None);
371 assert!((peer.success_rate() - 1.0).abs() < f64::EPSILON);
372
373 peer.record_failure();
374 assert!((peer.success_rate() - 0.5).abs() < f64::EPSILON);
375 }
376
377 #[test]
378 fn test_quality_calculation() {
379 let weights = super::super::config::QualityWeights::default();
380 let mut peer = CachedPeer::new(
381 PeerId([1u8; 32]),
382 vec!["127.0.0.1:9000".parse().unwrap()],
383 PeerSource::Seed,
384 );
385
386 peer.calculate_quality(&weights);
388 assert!(peer.quality_score > 0.3 && peer.quality_score < 0.7);
389
390 for _ in 0..5 {
392 peer.record_success(50, None); }
394 peer.calculate_quality(&weights);
395 assert!(peer.quality_score > 0.6);
396 }
397
398 #[test]
399 fn test_peer_serialization() {
400 let peer = CachedPeer::new(
401 PeerId([0xab; 32]),
402 vec!["127.0.0.1:9000".parse().unwrap()],
403 PeerSource::Seed,
404 );
405
406 let json = serde_json::to_string(&peer).unwrap();
407 let deserialized: CachedPeer = serde_json::from_str(&json).unwrap();
408
409 assert_eq!(deserialized.peer_id, peer.peer_id);
410 assert_eq!(deserialized.addresses, peer.addresses);
411 assert_eq!(deserialized.source, peer.source);
412 }
413}