1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fs;
5use std::net::SocketAddr;
6use std::path::Path;
7
8use crate::crypto::ZeroizingKey;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11pub enum PeerStatus {
12 Connected,
13 HandshakeInProgress,
14 Standby,
15 Disconnected,
16}
17
18impl PeerStatus {
19 pub fn badge_text(&self) -> &'static str {
20 match self {
21 PeerStatus::Connected => "ONLINE",
22 PeerStatus::HandshakeInProgress => "HANDSHAKE",
23 PeerStatus::Standby => "STANDBY",
24 PeerStatus::Disconnected => "OFFLINE",
25 }
26 }
27}
28
29#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
30pub enum PathType {
31 DirectIPv4(SocketAddr),
32 DirectIPv6(SocketAddr),
33 Relay(String),
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct CandidatePath {
38 pub path_type: PathType,
39 pub rtt_ms: Option<f64>,
40 pub last_success: Option<DateTime<Utc>>,
41 pub consecutive_failures: u32,
42 pub is_active: bool,
43}
44
45impl CandidatePath {
46 pub fn new(path_type: PathType) -> Self {
47 Self {
48 path_type,
49 rtt_ms: None,
50 last_success: None,
51 consecutive_failures: 0,
52 is_active: false,
53 }
54 }
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct PeerConfig {
59 pub callsign: String,
60 pub node_id: String,
61 pub public_key_base64: String,
62 pub endpoint: Option<String>,
63 pub overlay_ip: Option<String>,
64 pub created_at: DateTime<Utc>,
65}
66
67#[derive(Debug, Clone)]
68pub struct PeerState {
69 pub config: PeerConfig,
70 pub parsed_endpoint: Option<SocketAddr>,
71 pub status: PeerStatus,
72 pub rtt_ms: Option<f64>,
73 pub bytes_sent: u64,
74 pub bytes_recv: u64,
75 pub last_handshake: Option<DateTime<Utc>>,
76 pub last_ping_sent: Option<DateTime<Utc>>,
77 pub sequence_counter: u64,
78 pub zk_attested: bool,
79 pub zk_verified_at: Option<DateTime<Utc>>,
80 pub roaming_events: u64,
81 pub session_key: Option<ZeroizingKey>,
83 pub session_epoch: u64,
84 pub last_rekey: Option<DateTime<Utc>>,
85 pub rekey_count: u64,
86 pub pending_ephemeral_secret: Option<[u8; 32]>,
87 pub candidate_paths: Vec<CandidatePath>,
89 pub heartbeat_mode: HeartbeatMode,
91 pub last_data_activity: Option<DateTime<Utc>>,
92}
93
94#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
96pub enum HeartbeatMode {
97 Active,
99 Idle,
101 PowerSave,
103}
104
105impl HeartbeatMode {
106 pub fn interval(&self) -> std::time::Duration {
107 match self {
108 HeartbeatMode::Active => std::time::Duration::from_secs(2),
109 HeartbeatMode::Idle => std::time::Duration::from_secs(25),
110 HeartbeatMode::PowerSave => std::time::Duration::from_secs(120),
111 }
112 }
113}
114
115impl PeerState {
116 pub fn new(config: PeerConfig) -> Self {
117 let parsed_endpoint: Option<SocketAddr> = config.endpoint.as_deref().and_then(|ep| ep.parse().ok());
118 let mut candidate_paths = Vec::new();
119 if let Some(ep) = parsed_endpoint {
120 let ptype = if ep.is_ipv6() {
121 PathType::DirectIPv6(ep)
122 } else {
123 PathType::DirectIPv4(ep)
124 };
125 let mut cp = CandidatePath::new(ptype);
126 cp.is_active = true;
127 candidate_paths.push(cp);
128 }
129
130 Self {
131 config,
132 parsed_endpoint,
133 status: PeerStatus::Disconnected,
134 rtt_ms: None,
135 bytes_sent: 0,
136 bytes_recv: 0,
137 last_handshake: None,
138 last_ping_sent: None,
139 sequence_counter: 0,
140 zk_attested: false,
141 zk_verified_at: None,
142 roaming_events: 0,
143 session_key: None,
144 session_epoch: 0,
145 last_rekey: None,
146 rekey_count: 0,
147 pending_ephemeral_secret: None,
148 candidate_paths,
149 heartbeat_mode: HeartbeatMode::Idle,
150 last_data_activity: None,
151 }
152 }
153
154 pub fn record_data_activity(&mut self) {
156 self.last_data_activity = Some(Utc::now());
157 self.heartbeat_mode = HeartbeatMode::Active;
158 }
159
160 pub fn update_heartbeat_mode(&mut self, now: DateTime<Utc>) -> HeartbeatMode {
162 if let Some(last_activity) = self.last_data_activity {
163 let elapsed_secs = (now - last_activity).num_seconds();
164 if elapsed_secs < 30 {
165 self.heartbeat_mode = HeartbeatMode::Active;
166 } else if elapsed_secs < 120 {
167 self.heartbeat_mode = HeartbeatMode::Idle;
168 } else {
169 self.heartbeat_mode = HeartbeatMode::PowerSave;
170 }
171 } else {
172 self.heartbeat_mode = HeartbeatMode::Idle;
173 }
174 self.heartbeat_mode
175 }
176
177 pub fn is_heartbeat_due(&mut self, now: DateTime<Utc>) -> bool {
179 self.update_heartbeat_mode(now);
180 let interval = chrono::Duration::from_std(self.heartbeat_mode.interval()).unwrap_or(chrono::Duration::seconds(25));
181 if let Some(last_ping) = self.last_ping_sent {
182 now - last_ping >= interval
183 } else {
184 true
185 }
186 }
187
188 pub fn add_or_update_path(&mut self, path_type: PathType) {
190 if !self.candidate_paths.iter().any(|p| p.path_type == path_type) {
191 let mut cp = CandidatePath::new(path_type);
192 if self.candidate_paths.is_empty() {
193 cp.is_active = true;
194 }
195 self.candidate_paths.push(cp);
196 }
197 }
198
199 pub fn record_path_success(&mut self, path_type: &PathType, rtt: f64) {
201 for path in &mut self.candidate_paths {
202 if &path.path_type == path_type {
203 path.rtt_ms = Some(rtt);
204 path.last_success = Some(Utc::now());
205 path.consecutive_failures = 0;
206 }
207 }
208 }
209
210 pub fn record_path_failure(&mut self, path_type: &PathType, failure_threshold: u32) -> bool {
212 let mut failed_active = false;
213 for path in &mut self.candidate_paths {
214 if &path.path_type == path_type {
215 path.consecutive_failures += 1;
216 if path.is_active && path.consecutive_failures >= failure_threshold {
217 path.is_active = false;
218 failed_active = true;
219 }
220 }
221 }
222 if failed_active {
223 self.failover_to_best_path();
224 true
225 } else {
226 false
227 }
228 }
229
230 pub fn failover_to_best_path(&mut self) -> Option<PathType> {
232 for p in &mut self.candidate_paths {
233 p.is_active = false;
234 }
235
236 if let Some(best) = self.candidate_paths.iter_mut().min_by(|a, b| {
237 let score_a = a.consecutive_failures as f64 * 1000.0 + a.rtt_ms.unwrap_or(500.0);
238 let score_b = b.consecutive_failures as f64 * 1000.0 + b.rtt_ms.unwrap_or(500.0);
239 score_a.partial_cmp(&score_b).unwrap_or(std::cmp::Ordering::Equal)
240 }) {
241 best.is_active = true;
242 match &best.path_type {
243 PathType::DirectIPv4(addr) | PathType::DirectIPv6(addr) => {
244 self.parsed_endpoint = Some(*addr);
245 }
246 PathType::Relay(_) => {}
247 }
248 Some(best.path_type.clone())
249 } else {
250 None
251 }
252 }
253
254 pub fn active_path(&self) -> Option<&CandidatePath> {
255 self.candidate_paths.iter().find(|p| p.is_active)
256 }
257}
258
259#[derive(Debug, Clone, Default, Serialize, Deserialize)]
260pub struct PeersFile {
261 #[serde(default)]
262 pub peers: Vec<PeerConfig>,
263}
264
265pub struct PeerTable {
266 pub peers: HashMap<String, PeerState>,
267}
268
269impl PeerTable {
270 pub fn new() -> Self {
271 Self {
272 peers: HashMap::new(),
273 }
274 }
275
276 pub fn load_from_file(path: &Path) -> Result<Self, String> {
277 let mut table = Self::new();
278 if path.exists() {
279 let content = fs::read_to_string(path)
280 .map_err(|e| format!("Failed to read peers file: {}", e))?;
281 let file: PeersFile = toml::from_str(&content)
282 .map_err(|e| format!("Failed to parse peers.toml: {}", e))?;
283
284 for cfg in file.peers {
285 let pubkey = cfg.public_key_base64.clone();
286 table.peers.insert(pubkey, PeerState::new(cfg));
287 }
288 }
289 Ok(table)
290 }
291
292 pub fn save_to_file(&self, path: &Path) -> Result<(), String> {
293 let file = PeersFile {
294 peers: self.peers.values().map(|p| p.config.clone()).collect(),
295 };
296 let content = toml::to_string_pretty(&file)
297 .map_err(|e| format!("Failed to serialize peers.toml: {}", e))?;
298 fs::write(path, content)
299 .map_err(|e| format!("Failed to write peers.toml: {}", e))?;
300 Ok(())
301 }
302
303 pub fn add_peer(&mut self, config: PeerConfig) {
304 let pubkey = config.public_key_base64.clone();
305 self.peers.insert(pubkey, PeerState::new(config));
306 }
307
308 pub fn remove_peer(&mut self, pubkey_base64: &str) -> Option<PeerState> {
309 self.peers.remove(pubkey_base64)
310 }
311
312 pub fn get_mut_by_pubkey(&mut self, pubkey_base64: &str) -> Option<&mut PeerState> {
313 self.peers.get_mut(pubkey_base64)
314 }
315
316 pub fn get_by_pubkey(&self, pubkey_base64: &str) -> Option<&PeerState> {
317 self.peers.get(pubkey_base64)
318 }
319
320 pub fn list(&self) -> Vec<&PeerState> {
321 let mut list: Vec<&PeerState> = self.peers.values().collect();
322 list.sort_by(|a, b| a.config.callsign.cmp(&b.config.callsign));
323 list
324 }
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330
331 #[test]
332 fn test_candidate_path_fast_failover() {
333 let cfg = PeerConfig {
334 callsign: "backup-node".to_string(),
335 node_id: "sbm-0xbackup".to_string(),
336 public_key_base64: "dGVzdF9rZXk=".to_string(),
337 endpoint: Some("198.51.100.1:58888".to_string()),
338 overlay_ip: None,
339 created_at: Utc::now(),
340 };
341
342 let mut peer = PeerState::new(cfg);
343 let ipv4_path = PathType::DirectIPv4("198.51.100.1:58888".parse().unwrap());
344 let ipv6_path = PathType::DirectIPv6("[2001:db8::1]:58888".parse().unwrap());
345 let relay_path = PathType::Relay("sbm-0xrelaynode".to_string());
346
347 peer.add_or_update_path(ipv6_path.clone());
348 peer.add_or_update_path(relay_path.clone());
349
350 assert_eq!(peer.candidate_paths.len(), 3);
351 assert_eq!(peer.active_path().unwrap().path_type, ipv4_path);
352
353 assert!(!peer.record_path_failure(&ipv4_path, 2));
355 assert!(peer.record_path_failure(&ipv4_path, 2)); let active = peer.active_path().expect("Should have active path");
359 assert_ne!(active.path_type, ipv4_path);
360 assert_eq!(active.path_type, ipv6_path);
361 assert_eq!(peer.parsed_endpoint, Some("[2001:db8::1]:58888".parse().unwrap()));
362
363 peer.record_path_success(&ipv6_path, 25.0);
365 peer.record_path_success(&relay_path, 10.0);
366
367 let best = peer.failover_to_best_path().unwrap();
369 assert_eq!(best, relay_path);
370 assert_eq!(peer.active_path().unwrap().path_type, relay_path);
371 }
372
373 #[test]
374 fn test_adaptive_heartbeat_backoff_and_wake() {
375 let cfg = PeerConfig {
376 callsign: "mobile-node".to_string(),
377 node_id: "sbm-0xmobile".to_string(),
378 public_key_base64: "bW9iaWxlX2tleQ==".to_string(),
379 endpoint: Some("192.168.1.100:58888".to_string()),
380 overlay_ip: None,
381 created_at: Utc::now(),
382 };
383
384 let mut peer = PeerState::new(cfg);
385 let t0 = Utc::now();
386
387 assert_eq!(peer.update_heartbeat_mode(t0), HeartbeatMode::Idle);
389
390 peer.record_data_activity();
392 assert_eq!(peer.heartbeat_mode, HeartbeatMode::Active);
393 assert_eq!(peer.update_heartbeat_mode(t0 + chrono::Duration::seconds(10)), HeartbeatMode::Active);
394 assert_eq!(peer.heartbeat_mode.interval(), std::time::Duration::from_secs(2));
395
396 assert_eq!(peer.update_heartbeat_mode(t0 + chrono::Duration::seconds(45)), HeartbeatMode::Idle);
398 assert_eq!(peer.heartbeat_mode.interval(), std::time::Duration::from_secs(25));
399
400 assert_eq!(peer.update_heartbeat_mode(t0 + chrono::Duration::seconds(150)), HeartbeatMode::PowerSave);
402 assert_eq!(peer.heartbeat_mode.interval(), std::time::Duration::from_secs(120));
403
404 peer.record_data_activity();
406 assert_eq!(peer.heartbeat_mode, HeartbeatMode::Active);
407 }
408}
409
410
411