quicproquo/crates/quicprochat-p2p/src/routing_table.rs

//! Distributed routing table built from mesh announcements.
//!
//! The [`RoutingTable`] stores [`RoutingEntry`] records keyed by 16-byte
//! truncated mesh addresses, enabling multi-hop packet forwarding through
//! the mesh network.

use std::collections::HashMap;
use std::time::{Duration, Instant};

use crate::announce::MeshAnnounce;
use crate::transport::TransportAddr;

/// A routing entry for a known mesh destination.
#[derive(Clone, Debug)]
pub struct RoutingEntry {
    /// Full 32-byte Ed25519 public key of the destination.
    pub identity_key: [u8; 32],
    /// 16-byte truncated mesh address.
    pub address: [u8; 16],
    /// Next-hop transport name (e.g. "tcp", "iroh-quic", "lora").
    pub next_hop_transport: String,
    /// Next-hop address to send through.
    pub next_hop_addr: TransportAddr,
    /// Number of hops to this destination.
    pub hops: u8,
    /// Estimated cost (lower is better). Currently computed as hops as f64.
    pub cost: f64,
    /// Capabilities of the destination node.
    pub capabilities: u16,
    /// Last announce sequence number seen from this node.
    pub last_sequence: u64,
    /// When this entry was last updated.
    pub last_seen: Instant,
    /// When this entry expires (based on announce TTL).
    pub expires_at: Instant,
}

/// Distributed routing table built from received mesh announcements.
pub struct RoutingTable {
    /// Entries keyed by 16-byte truncated address.
    entries: HashMap<[u8; 16], RoutingEntry>,
    /// Default entry TTL.
    default_ttl: Duration,
}

impl RoutingTable {
    /// Create a new empty routing table with the given default TTL for entries.
    pub fn new(default_ttl: Duration) -> Self {
        Self {
            entries: HashMap::new(),
            default_ttl,
        }
    }

    /// Update the routing table from a received mesh announcement.
    ///
    /// Returns `true` if this was a new or improved route.
    ///
    /// Logic:
    /// - If `sequence <= last_sequence` for this address, the announce is stale — ignored.
    /// - If the entry is new or has lower cost, it replaces the existing entry.
    pub fn update(
        &mut self,
        announce: &MeshAnnounce,
        received_via_transport: &str,
        received_from: TransportAddr,
    ) -> bool {
        let address = announce.address;
        let new_cost = announce.hop_count as f64;
        let now = Instant::now();

        let identity_key: [u8; 32] = match announce.identity_key.as_slice().try_into() {
            Ok(k) => k,
            Err(_) => return false,
        };

        if let Some(existing) = self.entries.get(&address) {
            // Stale announce — older or same sequence number.
            if announce.sequence <= existing.last_sequence {
                return false;
            }
            // Only replace if the new route is better or equal (newer sequence wins on tie).
            if new_cost > existing.cost && announce.sequence == existing.last_sequence + 1 {
                // Higher cost with only incremental sequence — still update since it's fresher.
            }
        }

        let entry = RoutingEntry {
            identity_key,
            address,
            next_hop_transport: received_via_transport.to_string(),
            next_hop_addr: received_from,
            hops: announce.hop_count,
            cost: new_cost,
            capabilities: announce.capabilities,
            last_sequence: announce.sequence,
            last_seen: now,
            expires_at: now + self.default_ttl,
        };

        self.entries.insert(address, entry);
        true
    }

    /// Look up a routing entry by 16-byte truncated mesh address.
    pub fn lookup(&self, address: &[u8; 16]) -> Option<&RoutingEntry> {
        self.entries.get(address)
    }

    /// Look up a routing entry by the full 32-byte Ed25519 public key.
    pub fn lookup_by_key(&self, identity_key: &[u8; 32]) -> Option<&RoutingEntry> {
        self.entries.values().find(|e| &e.identity_key == identity_key)
    }

    /// Remove all expired entries. Returns the number of entries removed.
    pub fn remove_expired(&mut self) -> usize {
        let now = Instant::now();
        let before = self.entries.len();
        self.entries.retain(|_, entry| entry.expires_at > now);
        before - self.entries.len()
    }

    /// Iterate over all routing entries.
    pub fn entries(&self) -> impl Iterator<Item = &RoutingEntry> {
        self.entries.values()
    }

    /// Number of entries in the routing table.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Whether the routing table is empty.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::announce::{compute_address, CAP_RELAY};
    use crate::identity::MeshIdentity;

    fn make_announce(identity: &MeshIdentity, sequence: u64, hop_count: u8) -> MeshAnnounce {
        let mut announce =
            MeshAnnounce::with_sequence(identity, CAP_RELAY, vec![], 8, sequence);
        announce.hop_count = hop_count;
        announce
    }

    #[test]
    fn insert_and_lookup() {
        let mut table = RoutingTable::new(Duration::from_secs(300));
        let id = MeshIdentity::generate();
        let announce = make_announce(&id, 1, 1);
        let addr = TransportAddr::Socket("127.0.0.1:9000".parse().unwrap());

        assert!(table.update(&announce, "tcp", addr.clone()));
        assert_eq!(table.len(), 1);

        let mesh_addr = compute_address(&id.public_key());
        let entry = table.lookup(&mesh_addr).expect("entry should exist");
        assert_eq!(entry.hops, 1);
        assert_eq!(entry.last_sequence, 1);
        assert_eq!(entry.next_hop_transport, "tcp");
        assert_eq!(entry.next_hop_addr, addr);
    }

    #[test]
    fn update_with_better_route() {
        let mut table = RoutingTable::new(Duration::from_secs(300));
        let id = MeshIdentity::generate();
        let addr = TransportAddr::Socket("127.0.0.1:9000".parse().unwrap());

        // First announce: 3 hops, sequence 1.
        let announce1 = make_announce(&id, 1, 3);
        assert!(table.update(&announce1, "tcp", addr.clone()));

        let mesh_addr = compute_address(&id.public_key());
        assert_eq!(table.lookup(&mesh_addr).unwrap().hops, 3);

        // Second announce: 1 hop, sequence 2 — should replace.
        let announce2 = make_announce(&id, 2, 1);
        assert!(table.update(&announce2, "tcp", addr));

        let entry = table.lookup(&mesh_addr).unwrap();
        assert_eq!(entry.hops, 1);
        assert_eq!(entry.last_sequence, 2);
    }

    #[test]
    fn reject_stale_sequence() {
        let mut table = RoutingTable::new(Duration::from_secs(300));
        let id = MeshIdentity::generate();
        let addr = TransportAddr::Socket("127.0.0.1:9000".parse().unwrap());

        // Insert with sequence 5.
        let announce1 = make_announce(&id, 5, 1);
        assert!(table.update(&announce1, "tcp", addr.clone()));

        // Try to update with sequence 3 — should be rejected.
        let announce2 = make_announce(&id, 3, 1);
        assert!(
            !table.update(&announce2, "tcp", addr),
            "stale sequence must be rejected"
        );

        let mesh_addr = compute_address(&id.public_key());
        assert_eq!(table.lookup(&mesh_addr).unwrap().last_sequence, 5);
    }

    #[test]
    fn expire_old_entries() {
        let mut table = RoutingTable::new(Duration::from_millis(1));
        let id = MeshIdentity::generate();
        let addr = TransportAddr::Socket("127.0.0.1:9000".parse().unwrap());

        let announce = make_announce(&id, 1, 1);
        table.update(&announce, "tcp", addr);
        assert_eq!(table.len(), 1);

        // Wait for TTL to expire.
        std::thread::sleep(Duration::from_millis(10));

        let removed = table.remove_expired();
        assert_eq!(removed, 1);
        assert!(table.is_empty());
    }

    #[test]
    fn lookup_by_key_works() {
        let mut table = RoutingTable::new(Duration::from_secs(300));
        let id = MeshIdentity::generate();
        let addr = TransportAddr::Socket("127.0.0.1:9000".parse().unwrap());

        let announce = make_announce(&id, 1, 2);
        table.update(&announce, "tcp", addr);

        let pk = id.public_key();
        let entry = table.lookup_by_key(&pk).expect("should find by key");
        assert_eq!(entry.identity_key, pk);
        assert_eq!(entry.hops, 2);
    }
}