1  //! Routing Table
  2  //!
  3  //! Greedy routing based on tree coordinates and DHT-style key distance.
  4  //!
  5  //! This is **Pure Logic** - no IO, no networking.
  6  //! Takes State → Returns Decisions.
  7  
  8  use std::collections::HashMap;
  9  
 10  use ed25519_dalek::VerifyingKey;
 11  use serde::{Deserialize, Serialize};
 12  use thiserror::Error;
 13  
 14  use crate::address::{address_for_key, Address};
 15  use crate::coords::{RouteDirection, TreeCoords};
 16  
 17  /// Errors from routing operations
 18  #[derive(Error, Debug, Clone)]
 19  pub enum RoutingError {
 20      #[error("No route to destination")]
 21      NoRoute,
 22      
 23      #[error("Peer not found: {0}")]
 24      PeerNotFound(String),
 25      
 26      #[error("Invalid address")]
 27      InvalidAddress,
 28  
 29      #[error("Routing table full (max {0} peers)")]
 30      TableFull(usize),
 31  
 32      #[error("Too many peers from same prefix (max {0})")]
 33      PrefixLimitReached(usize),
 34  }
 35  
 36  /// A peer's identity (32-byte Ed25519 public key)
 37  pub type PeerKey = [u8; 32];
 38  
 39  /// Information about a known peer
 40  #[derive(Debug, Clone, Serialize, Deserialize)]
 41  pub struct PeerInfo {
 42      /// The peer's public key
 43      pub key: PeerKey,
 44      /// The peer's tree coordinates  
 45      pub coords: TreeCoords,
 46      /// The peer's derived IPv6 address
 47      pub address: Address,
 48      /// Port number on our side (which child slot they occupy)
 49      pub port: u64,
 50      /// Is this our parent in the tree?
 51      pub is_parent: bool,
 52      /// Link quality metric (lower is better)
 53      pub cost: u64,
 54  }
 55  
 56  impl PeerInfo {
 57      /// Create a new peer info from a verifying key
 58      pub fn new(key: &VerifyingKey, coords: TreeCoords, port: u64, is_parent: bool, cost: u64) -> Self {
 59          let address = address_for_key(key);
 60          Self {
 61              key: *key.as_bytes(),
 62              coords,
 63              address,
 64              port,
 65              is_parent,
 66              cost,
 67          }
 68      }
 69  
 70      /// Create from raw key bytes (for testing)
 71      pub fn from_key_bytes(key: PeerKey, coords: TreeCoords, port: u64, is_parent: bool, cost: u64) -> Self {
 72          // Derive address from key bytes
 73          // Note: This creates a dummy key for address derivation
 74          let mut addr_bytes = [0u8; 16];
 75          addr_bytes[0] = 0x02;
 76          // Use first byte of key as ones count, rest as address bits
 77          addr_bytes[1] = key[0] & 0x7F; // max 127 ones
 78          addr_bytes[2..].copy_from_slice(&key[..14]);
 79          
 80          Self {
 81              key,
 82              coords,
 83              address: Address::from_bytes(addr_bytes),
 84              port,
 85              is_parent,
 86              cost,
 87          }
 88      }
 89  }
 90  
 91  /// Routing table configuration for Sybil resistance
 92  #[derive(Debug, Clone)]
 93  pub struct TableConfig {
 94      /// Maximum total peers (0 = unlimited)
 95      pub max_peers: usize,
 96      /// Maximum peers sharing the same first byte of pubkey (0 = unlimited)
 97      pub max_peers_per_prefix: usize,
 98  }
 99  
100  impl Default for TableConfig {
101      fn default() -> Self {
102          Self {
103              max_peers: 256,           // Reasonable limit for mesh connectivity
104              max_peers_per_prefix: 16, // Limits clustering in key space
105          }
106      }
107  }
108  
109  /// The routing table for a node
110  #[derive(Debug, Clone)]
111  pub struct RoutingTable {
112      /// Our own key (reserved for future direct routing)
113      #[allow(dead_code)]
114      self_key: PeerKey,
115      /// Our own coordinates
116      self_coords: TreeCoords,
117      /// Our own address
118      self_address: Address,
119      /// Known peers indexed by their key
120      peers: HashMap<PeerKey, PeerInfo>,
121      /// Parent peer key (if we have one)
122      parent: Option<PeerKey>,
123      /// Children indexed by port number
124      children: HashMap<u64, PeerKey>,
125      /// Configuration for Sybil resistance
126      config: TableConfig,
127  }
128  
129  impl RoutingTable {
130      /// Create a new routing table for a node
131      pub fn new(self_key: &VerifyingKey, self_coords: TreeCoords) -> Self {
132          Self::with_config(self_key, self_coords, TableConfig::default())
133      }
134  
135      /// Create a new routing table with custom configuration
136      pub fn with_config(self_key: &VerifyingKey, self_coords: TreeCoords, config: TableConfig) -> Self {
137          let self_address = address_for_key(self_key);
138          Self {
139              self_key: *self_key.as_bytes(),
140              self_coords,
141              self_address,
142              peers: HashMap::new(),
143              parent: None,
144              children: HashMap::new(),
145              config,
146          }
147      }
148  
149      /// Create from raw key bytes (for testing)
150      pub fn from_key_bytes(self_key: PeerKey, self_coords: TreeCoords) -> Self {
151          Self::from_key_bytes_with_config(self_key, self_coords, TableConfig::default())
152      }
153  
154      /// Create from raw key bytes with custom config (for testing)
155      pub fn from_key_bytes_with_config(self_key: PeerKey, self_coords: TreeCoords, config: TableConfig) -> Self {
156          let mut addr_bytes = [0u8; 16];
157          addr_bytes[0] = 0x02;
158          addr_bytes[1] = self_key[0] & 0x7F;
159          addr_bytes[2..].copy_from_slice(&self_key[..14]);
160          
161          Self {
162              self_key,
163              self_coords,
164              self_address: Address::from_bytes(addr_bytes),
165              peers: HashMap::new(),
166              parent: None,
167              children: HashMap::new(),
168              config,
169          }
170      }
171  
172      /// Get our own address
173      pub fn self_address(&self) -> &Address {
174          &self.self_address
175      }
176  
177      /// Get our own coordinates
178      pub fn self_coords(&self) -> &TreeCoords {
179          &self.self_coords
180      }
181  
182      /// Add or update a peer. Returns error if limits would be exceeded.
183      pub fn add_peer(&mut self, peer: PeerInfo) -> Result<(), RoutingError> {
184          let key = peer.key;
185          
186          // If updating existing peer, skip limit checks
187          if !self.peers.contains_key(&key) {
188              // Check total peer limit
189              if self.config.max_peers > 0 && self.peers.len() >= self.config.max_peers {
190                  return Err(RoutingError::TableFull(self.config.max_peers));
191              }
192              
193              // Check per-prefix limit (first byte of pubkey)
194              if self.config.max_peers_per_prefix > 0 {
195                  let prefix = key[0];
196                  let same_prefix = self.peers.keys().filter(|k| k[0] == prefix).count();
197                  if same_prefix >= self.config.max_peers_per_prefix {
198                      return Err(RoutingError::PrefixLimitReached(self.config.max_peers_per_prefix));
199                  }
200              }
201          }
202          
203          if peer.is_parent {
204              self.parent = Some(key);
205          } else {
206              self.children.insert(peer.port, key);
207          }
208          
209          self.peers.insert(key, peer);
210          Ok(())
211      }
212  
213      /// Remove a peer
214      pub fn remove_peer(&mut self, key: &PeerKey) -> Option<PeerInfo> {
215          if let Some(peer) = self.peers.remove(key) {
216              if peer.is_parent {
217                  self.parent = None;
218              } else {
219                  self.children.remove(&peer.port);
220              }
221              Some(peer)
222          } else {
223              None
224          }
225      }
226  
227      /// Get a peer by key
228      pub fn get_peer(&self, key: &PeerKey) -> Option<&PeerInfo> {
229          self.peers.get(key)
230      }
231  
232      /// Get all peers
233      pub fn peers(&self) -> impl Iterator<Item = &PeerInfo> {
234          self.peers.values()
235      }
236  
237      /// Get parent peer
238      pub fn parent(&self) -> Option<&PeerInfo> {
239          self.parent.as_ref().and_then(|k| self.peers.get(k))
240      }
241  
242      /// Get children
243      pub fn children(&self) -> impl Iterator<Item = &PeerInfo> {
244          self.children.values().filter_map(|k| self.peers.get(k))
245      }
246  
247      /// Find the next hop for a destination address.
248      ///
249      /// Implements greedy routing:
250      /// 1. If destination is us, return None (we're the destination)
251      /// 2. If destination is in our tree (descendant), route to child
252      /// 3. Otherwise, route toward parent OR use DHT-style closest peer
253      pub fn find_next_hop(&self, destination: &Address) -> Result<Option<PeerKey>, RoutingError> {
254          // Check if destination is us
255          if destination == &self.self_address {
256              return Ok(None);
257          }
258  
259          // First, try tree-based routing if we know the destination's coordinates
260          if let Some(dest_coords) = self.find_coords_for_address(destination) {
261              return self.route_by_coords(&dest_coords);
262          }
263  
264          // Fall back to DHT-style routing (closest key to destination)
265          self.route_by_key_distance(destination)
266      }
267  
268      /// Route using tree coordinates (greedy tree routing)
269      fn route_by_coords(&self, dest_coords: &TreeCoords) -> Result<Option<PeerKey>, RoutingError> {
270          match self.self_coords.route_direction(dest_coords) {
271              RouteDirection::Self_ => Ok(None),
272              
273              RouteDirection::Up => {
274                  // Send to parent
275                  self.parent.ok_or(RoutingError::NoRoute).map(Some)
276              }
277              
278              RouteDirection::Down(port) => {
279                  // Send to child on that port
280                  self.children
281                      .get(&port)
282                      .copied()
283                      .ok_or(RoutingError::NoRoute)
284                      .map(Some)
285              }
286          }
287      }
288  
289      /// Route using DHT-style key distance (XOR metric)
290      fn route_by_key_distance(&self, destination: &Address) -> Result<Option<PeerKey>, RoutingError> {
291          let dest_bytes = destination.as_bytes();
292          
293          // Find peer with smallest XOR distance to destination
294          let mut best_peer: Option<&PeerInfo> = None;
295          let mut best_distance = u128::MAX;
296          
297          for peer in self.peers.values() {
298              let peer_addr = peer.address.as_bytes();
299              let distance = xor_distance(peer_addr, dest_bytes);
300              
301              if distance < best_distance {
302                  best_distance = distance;
303                  best_peer = Some(peer);
304              }
305          }
306  
307          // Check if any peer is closer than we are
308          let self_distance = xor_distance(self.self_address.as_bytes(), dest_bytes);
309          
310          if let Some(peer) = best_peer
311              && best_distance < self_distance
312          {
313              return Ok(Some(peer.key));
314          }
315  
316          // No peer is closer - we might be the closest, but we're not the destination
317          // This could mean we need to wait for more routing info
318          Err(RoutingError::NoRoute)
319      }
320  
321      /// Find coordinates for an address (if we know them)
322      fn find_coords_for_address(&self, address: &Address) -> Option<TreeCoords> {
323          // Check if it's a known peer
324          for peer in self.peers.values() {
325              if &peer.address == address {
326                  return Some(peer.coords.clone());
327              }
328          }
329          None
330      }
331  
332      /// Get routing table statistics
333      pub fn stats(&self) -> RoutingStats {
334          RoutingStats {
335              peer_count: self.peers.len(),
336              has_parent: self.parent.is_some(),
337              child_count: self.children.len(),
338              tree_depth: self.self_coords.depth(),
339          }
340      }
341  }
342  
343  /// Routing table statistics
344  #[derive(Debug, Clone)]
345  pub struct RoutingStats {
346      pub peer_count: usize,
347      pub has_parent: bool,
348      pub child_count: usize,
349      pub tree_depth: usize,
350  }
351  
352  /// Calculate XOR distance between two addresses (for DHT routing)
353  fn xor_distance(a: &[u8; 16], b: &[u8; 16]) -> u128 {
354      let a_val = u128::from_be_bytes(*a);
355      let b_val = u128::from_be_bytes(*b);
356      a_val ^ b_val
357  }
358  
359  #[cfg(test)]
360  mod tests {
361      use super::*;
362  
363      fn make_key(seed: u8) -> PeerKey {
364          let mut key = [0u8; 32];
365          key[0] = seed;
366          key[1] = seed.wrapping_add(1);
367          for i in 2..32 {
368              key[i] = seed.wrapping_add(i as u8);
369          }
370          key
371      }
372  
373      #[test]
374      fn test_add_remove_peer() {
375          let self_key = make_key(0);
376          let mut table = RoutingTable::from_key_bytes(self_key, TreeCoords::root());
377          
378          let peer_key = make_key(1);
379          let peer = PeerInfo::from_key_bytes(
380              peer_key,
381              TreeCoords::from_path(vec![0]),
382              0,
383              false,
384              100,
385          );
386          
387          let _ = table.add_peer(peer);
388          assert!(table.get_peer(&peer_key).is_some());
389          
390          table.remove_peer(&peer_key);
391          assert!(table.get_peer(&peer_key).is_none());
392      }
393  
394      #[test]
395      fn test_tree_routing() {
396          let self_key = make_key(0);
397          let mut table = RoutingTable::from_key_bytes(
398              self_key,
399              TreeCoords::from_path(vec![1, 2]),
400          );
401          
402          // Add parent
403          let parent_key = make_key(1);
404          let parent = PeerInfo::from_key_bytes(
405              parent_key,
406              TreeCoords::from_path(vec![1]),
407              0,
408              true,
409              50,
410          );
411          let _ = table.add_peer(parent);
412          
413          // Add child
414          let child_key = make_key(2);
415          let child = PeerInfo::from_key_bytes(
416              child_key,
417              TreeCoords::from_path(vec![1, 2, 5]),
418              5,
419              false,
420              50,
421          );
422          let _ = table.add_peer(child);
423          
424          // Route to child's subtree should go to child
425          let child_dest = table.get_peer(&child_key).unwrap().address;
426          let next_hop = table.find_next_hop(&child_dest).unwrap();
427          assert_eq!(next_hop, Some(child_key));
428          
429          // Route to parent should go to parent  
430          let parent_dest = table.get_peer(&parent_key).unwrap().address;
431          let next_hop = table.find_next_hop(&parent_dest).unwrap();
432          assert_eq!(next_hop, Some(parent_key));
433      }
434  
435      #[test]
436      fn test_self_destination() {
437          let self_key = make_key(0);
438          let table = RoutingTable::from_key_bytes(self_key, TreeCoords::root());
439          
440          let result = table.find_next_hop(table.self_address());
441          assert_eq!(result.unwrap(), None);
442      }
443  
444      #[test]
445      fn test_stats() {
446          let self_key = make_key(0);
447          let mut table = RoutingTable::from_key_bytes(
448              self_key,
449              TreeCoords::from_path(vec![1, 2, 3]),
450          );
451          
452          let stats = table.stats();
453          assert_eq!(stats.peer_count, 0);
454          assert!(!stats.has_parent);
455          assert_eq!(stats.tree_depth, 3);
456          
457          // Add a parent
458          let parent = PeerInfo::from_key_bytes(
459              make_key(1),
460              TreeCoords::from_path(vec![1, 2]),
461              0,
462              true,
463              50,
464          );
465          let _ = table.add_peer(parent);
466          
467          let stats = table.stats();
468          assert_eq!(stats.peer_count, 1);
469          assert!(stats.has_parent);
470      }
471  
472      #[test]
473      fn test_sybil_max_peers() {
474          let config = TableConfig {
475              max_peers: 3,
476              max_peers_per_prefix: 0, // Disabled
477          };
478          let self_key = make_key(0);
479          let mut table = RoutingTable::from_key_bytes_with_config(
480              self_key,
481              TreeCoords::root(),
482              config,
483          );
484          
485          // Add 3 peers - should succeed
486          for i in 1..=3 {
487              let peer = PeerInfo::from_key_bytes(make_key(i), TreeCoords::root(), i as u64, false, 50);
488              assert!(table.add_peer(peer).is_ok());
489          }
490          
491          // 4th peer should fail
492          let peer = PeerInfo::from_key_bytes(make_key(4), TreeCoords::root(), 4, false, 50);
493          let err = table.add_peer(peer).unwrap_err();
494          assert!(matches!(err, RoutingError::TableFull(3)));
495          
496          // Updating existing peer should still work
497          let peer = PeerInfo::from_key_bytes(make_key(1), TreeCoords::root(), 1, false, 100);
498          assert!(table.add_peer(peer).is_ok());
499      }
500  
501      #[test]
502      fn test_sybil_prefix_limit() {
503          let config = TableConfig {
504              max_peers: 256,
505              max_peers_per_prefix: 2, // Only 2 peers per prefix
506          };
507          let self_key = make_key(0);
508          let mut table = RoutingTable::from_key_bytes_with_config(
509              self_key,
510              TreeCoords::root(),
511              config,
512          );
513          
514          // make_key puts the index in the first byte, so keys 1,2 share prefix 0x01, 0x02
515          // Create peers with same prefix (we'll force same first byte)
516          let mut key1 = [0u8; 32];
517          key1[0] = 0xAA;
518          key1[1] = 0x01;
519          let mut key2 = [0u8; 32];
520          key2[0] = 0xAA;
521          key2[1] = 0x02;
522          let mut key3 = [0u8; 32];
523          key3[0] = 0xAA;
524          key3[1] = 0x03;
525          let mut key4 = [0u8; 32];
526          key4[0] = 0xBB; // Different prefix
527          
528          // Add 2 peers with prefix 0xAA - should succeed
529          let peer1 = PeerInfo::from_key_bytes(key1, TreeCoords::root(), 1, false, 50);
530          let peer2 = PeerInfo::from_key_bytes(key2, TreeCoords::root(), 2, false, 50);
531          assert!(table.add_peer(peer1).is_ok());
532          assert!(table.add_peer(peer2).is_ok());
533          
534          // 3rd peer with same prefix should fail
535          let peer3 = PeerInfo::from_key_bytes(key3, TreeCoords::root(), 3, false, 50);
536          let err = table.add_peer(peer3).unwrap_err();
537          assert!(matches!(err, RoutingError::PrefixLimitReached(2)));
538          
539          // Peer with different prefix should succeed
540          let peer4 = PeerInfo::from_key_bytes(key4, TreeCoords::root(), 4, false, 50);
541          assert!(table.add_peer(peer4).is_ok());
542      }
543  }