Skip to main content

sb_mesh/
peer.rs

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    // Phase 14: In-Band Rekeying & PFS Ratchet
82    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    // Phase 15: Multi-Path Fast Failover
88    pub candidate_paths: Vec<CandidatePath>,
89    // Phase 20: Adaptive Power-Saving Keepalive
90    pub heartbeat_mode: HeartbeatMode,
91    pub last_data_activity: Option<DateTime<Utc>>,
92}
93
94/// Sleep-Aware Adaptive Heartbeat Mode for Battery & Bandwidth Conservation
95#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
96pub enum HeartbeatMode {
97    /// High-activity mode (<30s since last traffic): fast 2s keepalive
98    Active,
99    /// Moderate-idle mode (30s - 120s since last traffic): standard 25s keepalive
100    Idle,
101    /// Sleep-aware battery-saver mode (>120s since last traffic): 120s backoff keepalive
102    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    /// Record packet/data activity, waking the peer and resetting heartbeat mode to Active
155    pub fn record_data_activity(&mut self) {
156        self.last_data_activity = Some(Utc::now());
157        self.heartbeat_mode = HeartbeatMode::Active;
158    }
159
160    /// Evaluates current idle duration and updates heartbeat mode accordingly
161    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    /// Checks if a heartbeat ping is due based on adaptive interval
178    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    /// Add or update candidate network path
189    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    /// Record success on a candidate path, updating RTT and resetting failures
200    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    /// Record failure on a candidate path; triggers fast failover if threshold reached
211    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    /// Selects the best candidate path (lowest latency, fewest failures) as active
231    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        // IPv4 fails twice (threshold = 2)
354        assert!(!peer.record_path_failure(&ipv4_path, 2));
355        assert!(peer.record_path_failure(&ipv4_path, 2)); // threshold met -> failover!
356
357        // Active path should now failover to IPv6
358        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        // Now record IPv6 latency 25ms, Relay latency 10ms
364        peer.record_path_success(&ipv6_path, 25.0);
365        peer.record_path_success(&relay_path, 10.0);
366
367        // Fast failover selects best path by latency
368        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        // 1. Initially Idle (no activity recorded yet)
388        assert_eq!(peer.update_heartbeat_mode(t0), HeartbeatMode::Idle);
389
390        // 2. Data activity occurs -> immediately Active (2s interval)
391        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        // 3. After 45s idle -> drops to Idle (25s interval)
397        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        // 4. After 150s idle -> drops to PowerSave (120s backoff interval)
401        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        // 5. Wake on packet / traffic -> immediately returns to Active
405        peer.record_data_activity();
406        assert_eq!(peer.heartbeat_mode, HeartbeatMode::Active);
407    }
408}
409
410
411