1  //! Multi-Root Spanning Tree Support
  2  //!
  3  //! Addresses the single-root vulnerability by allowing multiple "landmark" roots.
  4  //! Each root forms an independent spanning tree. Nodes maintain coordinates 
  5  //! relative to multiple roots, enabling routing to continue if any root fails.
  6  //!
  7  //! **Design**:
  8  //! - Roots are high-uptime, well-connected nodes (self-elected or community-designated)
  9  //! - Each root has a unique `RootId` (hash of their public key)
 10  //! - Nodes can maintain coordinates in up to 3 trees simultaneously
 11  //! - Routing prefers the tree with lowest hop count, falls back to DHT
 12  
 13  use serde::{Deserialize, Serialize};
 14  use std::collections::HashMap;
 15  
 16  use crate::coords::TreeCoords;
 17  
 18  /// Maximum number of roots a node tracks simultaneously
 19  pub const MAX_TRACKED_ROOTS: usize = 3;
 20  
 21  /// Root identifier (hash of root's public key)
 22  pub type RootId = [u8; 32];
 23  
 24  /// Information about a spanning tree root
 25  #[derive(Debug, Clone, Serialize, Deserialize)]
 26  pub struct RootInfo {
 27      /// Root's public key hash
 28      pub id: RootId,
 29      /// Root's public key (full)
 30      pub public_key: [u8; 32],
 31      /// Our coordinates relative to this root
 32      pub our_coords: TreeCoords,
 33      /// Root's announced sequence number (prevents replay)
 34      pub sequence: u64,
 35      /// Last time we heard from this root (unix millis)
 36      pub last_seen: u64,
 37      /// Link quality to this root (lower is better)
 38      pub cost: u64,
 39  }
 40  
 41  /// Multi-root coordinator for a node
 42  #[derive(Debug, Clone, Default)]
 43  pub struct MultiRootCoordinator {
 44      /// Known roots, keyed by RootId
 45      roots: HashMap<RootId, RootInfo>,
 46      /// Our preferred root (lowest cost)
 47      preferred_root: Option<RootId>,
 48  }
 49  
 50  impl MultiRootCoordinator {
 51      /// Create a new coordinator
 52      pub fn new() -> Self {
 53          Self::default()
 54      }
 55  
 56      /// Add or update a root
 57      ///
 58      /// Returns true if this is a new root or an update to existing.
 59      pub fn update_root(&mut self, info: RootInfo) -> bool {
 60          let id = info.id;
 61          
 62          // Check if this is a stale announcement
 63          if let Some(existing) = self.roots.get(&id)
 64              && info.sequence <= existing.sequence {
 65                  return false; // Stale or replay
 66              }
 67  
 68          // Enforce max roots - evict worst if at capacity
 69          if self.roots.len() >= MAX_TRACKED_ROOTS && !self.roots.contains_key(&id) {
 70              self.evict_worst_root();
 71          }
 72  
 73          self.roots.insert(id, info);
 74          self.update_preferred_root();
 75          true
 76      }
 77  
 78      /// Remove a root (e.g., if it's been offline too long)
 79      pub fn remove_root(&mut self, id: &RootId) -> Option<RootInfo> {
 80          let removed = self.roots.remove(id);
 81          if removed.is_some() {
 82              self.update_preferred_root();
 83          }
 84          removed
 85      }
 86  
 87      /// Get our coordinates for the preferred root
 88      pub fn preferred_coords(&self) -> Option<&TreeCoords> {
 89          self.preferred_root
 90              .as_ref()
 91              .and_then(|id| self.roots.get(id))
 92              .map(|info| &info.our_coords)
 93      }
 94  
 95      /// Get our coordinates for a specific root
 96      pub fn coords_for_root(&self, root_id: &RootId) -> Option<&TreeCoords> {
 97          self.roots.get(root_id).map(|info| &info.our_coords)
 98      }
 99  
100      /// Get all known roots
101      pub fn roots(&self) -> impl Iterator<Item = &RootInfo> {
102          self.roots.values()
103      }
104  
105      /// Check if we're tracking any roots
106      pub fn has_roots(&self) -> bool {
107          !self.roots.is_empty()
108      }
109  
110      /// Get the preferred root info
111      pub fn preferred_root(&self) -> Option<&RootInfo> {
112          self.preferred_root
113              .as_ref()
114              .and_then(|id| self.roots.get(id))
115      }
116  
117      /// Evict roots that haven't been seen within the timeout
118      ///
119      /// Returns number of roots evicted.
120      pub fn prune_stale_roots(&mut self, timeout_ms: u64, now_ms: u64) -> usize {
121          let before = self.roots.len();
122          self.roots.retain(|_, info| {
123              now_ms.saturating_sub(info.last_seen) < timeout_ms
124          });
125          let evicted = before.saturating_sub(self.roots.len());
126          if evicted > 0 {
127              self.update_preferred_root();
128          }
129          evicted
130      }
131  
132      /// Evict the worst (highest cost) root
133      fn evict_worst_root(&mut self) {
134          if let Some(worst_id) = self.roots
135              .iter()
136              .max_by_key(|(_, info)| info.cost)
137              .map(|(id, _)| *id)
138          {
139              self.roots.remove(&worst_id);
140          }
141      }
142  
143      /// Update which root is preferred (lowest cost)
144      fn update_preferred_root(&mut self) {
145          self.preferred_root = self.roots
146              .iter()
147              .min_by_key(|(_, info)| info.cost)
148              .map(|(id, _)| *id);
149      }
150  
151      /// Get routing statistics
152      pub fn stats(&self) -> MultiRootStats {
153          MultiRootStats {
154              root_count: self.roots.len(),
155              has_preferred: self.preferred_root.is_some(),
156              avg_cost: if self.roots.is_empty() {
157                  0
158              } else {
159                  self.roots.values().map(|r| r.cost).sum::<u64>() / self.roots.len() as u64
160              },
161          }
162      }
163  }
164  
165  /// Statistics about multi-root state
166  #[derive(Debug, Clone)]
167  pub struct MultiRootStats {
168      pub root_count: usize,
169      pub has_preferred: bool,
170      pub avg_cost: u64,
171  }
172  
173  #[cfg(test)]
174  mod tests {
175      use super::*;
176  
177      fn make_root(seed: u8, cost: u64, sequence: u64) -> RootInfo {
178          let mut id = [0u8; 32];
179          let mut public_key = [0u8; 32];
180          id[0] = seed;
181          public_key[0] = seed;
182          
183          RootInfo {
184              id,
185              public_key,
186              our_coords: TreeCoords::from_path(vec![seed as u64]),
187              sequence,
188              last_seen: 1000,
189              cost,
190          }
191      }
192  
193      #[test]
194      fn test_add_roots() {
195          let mut coord = MultiRootCoordinator::new();
196          
197          assert!(coord.update_root(make_root(1, 100, 1)));
198          assert!(coord.update_root(make_root(2, 50, 1)));
199          
200          assert_eq!(coord.roots.len(), 2);
201          
202          // Preferred should be lowest cost
203          let preferred = coord.preferred_root().unwrap();
204          assert_eq!(preferred.id[0], 2);
205      }
206  
207      #[test]
208      fn test_reject_stale_sequence() {
209          let mut coord = MultiRootCoordinator::new();
210          
211          assert!(coord.update_root(make_root(1, 100, 5)));
212          assert!(!coord.update_root(make_root(1, 100, 3))); // Stale
213          assert!(!coord.update_root(make_root(1, 100, 5))); // Same seq
214          assert!(coord.update_root(make_root(1, 100, 6)));  // Newer
215      }
216  
217      #[test]
218      fn test_max_roots_eviction() {
219          let mut coord = MultiRootCoordinator::new();
220          
221          // Add MAX_TRACKED_ROOTS
222          for i in 0..MAX_TRACKED_ROOTS as u8 {
223              coord.update_root(make_root(i, 100 - i as u64, 1));
224          }
225          
226          assert_eq!(coord.roots.len(), MAX_TRACKED_ROOTS);
227          
228          // Add one more - should evict highest cost
229          coord.update_root(make_root(99, 10, 1));
230          
231          assert_eq!(coord.roots.len(), MAX_TRACKED_ROOTS);
232          // Root with cost 100 (seed 0) should be evicted
233          assert!(!coord.roots.contains_key(&[0u8; 32].map(|_| 0)));
234      }
235  
236      #[test]
237      fn test_prune_stale() {
238          let mut coord = MultiRootCoordinator::new();
239          
240          coord.update_root(RootInfo {
241              last_seen: 1000,
242              ..make_root(1, 100, 1)
243          });
244          coord.update_root(RootInfo {
245              last_seen: 5000,
246              ..make_root(2, 50, 1)
247          });
248          
249          // Prune with 3000ms timeout at time 6000
250          let evicted = coord.prune_stale_roots(3000, 6000);
251          
252          assert_eq!(evicted, 1);
253          assert_eq!(coord.roots.len(), 1);
254      }
255  }