shadow_dht/

routing.rs

//! Kademlia routing table implementation

use shadow_core::{PeerId, PeerInfo};
use shadow_core::error::{Result, ShadowError};
use std::collections::VecDeque;
use std::time::{Instant, Duration};

/// K-bucket containing peers at a specific distance range
#[derive(Debug, Clone)]
pub struct KBucket {
    /// Maximum bucket size (k parameter)
    pub k: usize,
    /// Peers in this bucket (oldest first)
    pub peers: VecDeque<BucketEntry>,
    /// Last update time
    pub last_updated: Instant,
}

/// Entry in a K-bucket
#[derive(Debug, Clone)]
pub struct BucketEntry {
    pub peer: PeerInfo,
    pub last_seen: Instant,
    pub rtt: Option<Duration>,
    pub failed_queries: u32,
}

impl KBucket {
    /// Create new K-bucket
    pub fn new(k: usize) -> Self {
        Self {
            k,
            peers: VecDeque::new(),
            last_updated: Instant::now(),
        }
    }

    /// Try to add peer to bucket
    pub fn add_peer(&mut self, peer: PeerInfo) -> bool {
        // Check if peer already exists
        if let Some(pos) = self.peers.iter().position(|e| e.peer.id == peer.id) {
            // Move to end (most recently seen)
            let mut entry = self.peers.remove(pos).unwrap();
            entry.last_seen = Instant::now();
            self.peers.push_back(entry);
            self.last_updated = Instant::now();
            return true;
        }

        // If bucket not full, add peer
        if self.peers.len() < self.k {
            self.peers.push_back(BucketEntry {
                peer,
                last_seen: Instant::now(),
                rtt: None,
                failed_queries: 0,
            });
            self.last_updated = Instant::now();
            return true;
        }

        // Bucket full: try to replace stale peer
        // Check if oldest peer (front) is stale
        if let Some(oldest) = self.peers.front() {
            if oldest.last_seen.elapsed() > Duration::from_secs(900) || oldest.failed_queries > 3 {
                // Replace stale peer
                self.peers.pop_front();
                self.peers.push_back(BucketEntry {
                    peer,
                    last_seen: Instant::now(),
                    rtt: None,
                    failed_queries: 0,
                });
                self.last_updated = Instant::now();
                return true;
            }
        }

        false // Couldn't add peer
    }

    /// Remove peer from bucket
    pub fn remove_peer(&mut self, peer_id: &PeerId) -> bool {
        if let Some(pos) = self.peers.iter().position(|e| &e.peer.id == peer_id) {
            self.peers.remove(pos);
            self.last_updated = Instant::now();
            return true;
        }
        false
    }

    /// Get peer by ID
    pub fn get_peer(&self, peer_id: &PeerId) -> Option<&PeerInfo> {
        self.peers.iter()
            .find(|e| &e.peer.id == peer_id)
            .map(|e| &e.peer)
    }

    /// Mark peer as failed
    pub fn mark_failed(&mut self, peer_id: &PeerId) {
        if let Some(entry) = self.peers.iter_mut().find(|e| &e.peer.id == peer_id) {
            entry.failed_queries += 1;
        }
    }

    /// Get all peers in bucket
    pub fn all_peers(&self) -> Vec<PeerInfo> {
        self.peers.iter().map(|e| e.peer.clone()).collect()
    }

    /// Get number of peers
    pub fn len(&self) -> usize {
        self.peers.len()
    }

    /// Check if bucket is empty
    pub fn is_empty(&self) -> bool {
        self.peers.is_empty()
    }
}

/// Kademlia routing table
pub struct RoutingTable {
    /// Our peer ID
    local_id: PeerId,
    /// K-buckets (160 buckets for 160-bit IDs)
    buckets: Vec<KBucket>,
    /// K parameter (bucket size)
    k: usize,
}

impl RoutingTable {
    /// Create new routing table
    pub fn new(local_id: PeerId, k: usize) -> Self {
        let buckets = (0..160)
            .map(|_| KBucket::new(k))
            .collect();

        Self {
            local_id,
            buckets,
            k,
        }
    }

    /// Add peer to routing table
    pub fn add_peer(&mut self, peer: PeerInfo) -> Result<bool> {
        if peer.id == self.local_id {
            return Ok(false); // Don't add ourselves
        }

        let bucket_index = self.bucket_index(&peer.id)?;
        Ok(self.buckets[bucket_index].add_peer(peer))
    }

    /// Remove peer from routing table
    pub fn remove_peer(&mut self, peer_id: &PeerId) -> Result<bool> {
        let bucket_index = self.bucket_index(peer_id)?;
        Ok(self.buckets[bucket_index].remove_peer(peer_id))
    }

    /// Find closest K peers to target ID
    pub fn find_closest(&self, target: &PeerId, count: usize) -> Vec<PeerInfo> {
        let mut peers: Vec<(PeerInfo, [u8; 32])> = Vec::new();

        // Collect all peers with their XOR distances to target
        for bucket in &self.buckets {
            for peer in bucket.all_peers() {
                let distance = target.xor_distance(&peer.id);
                peers.push((peer, distance));
            }
        }

        // Sort by distance to target (compare byte arrays directly)
        peers.sort_by(|a, b| a.1.cmp(&b.1));

        // Return top K
        peers.into_iter()
            .take(count)
            .map(|(peer, _)| peer)
            .collect()
    }

    /// Get bucket index for peer ID (based on XOR distance)
    fn bucket_index(&self, peer_id: &PeerId) -> Result<usize> {
        let distance = self.local_id.xor_distance(peer_id);

        // Find first differing bit (bucket index)
        for (i, &byte) in distance.iter().enumerate() {
            if byte != 0 {
                // Find position of first 1 bit
                let bit_pos = 7 - byte.leading_zeros() as usize;
                let bucket_idx = i * 8 + bit_pos;
                return Ok(bucket_idx.min(159)); // Max bucket index
            }
        }

        // All bits same (shouldn't happen with different IDs)
        Err(ShadowError::Dht("Identical peer IDs".into()))
    }

    /// Get all known peers
    pub fn all_peers(&self) -> Vec<PeerInfo> {
        let mut all = Vec::new();
        for bucket in &self.buckets {
            all.extend(bucket.all_peers());
        }
        all
    }

    /// Get total number of peers
    pub fn peer_count(&self) -> usize {
        self.buckets.iter().map(|b| b.len()).sum()
    }

    /// Mark peer as failed
    pub fn mark_failed(&mut self, peer_id: &PeerId) -> Result<()> {
        let bucket_index = self.bucket_index(peer_id)?;
        self.buckets[bucket_index].mark_failed(peer_id);
        Ok(())
    }

    /// Get our local ID
    pub fn local_id(&self) -> &PeerId {
        &self.local_id
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_kbucket_add() {
        let mut bucket = KBucket::new(3);

        let peer1 = PeerInfo::new(
            PeerId::random(),
            vec!["127.0.0.1:9000".to_string()],
            [0u8; 32],
            [0u8; 32],
        );

        assert!(bucket.add_peer(peer1.clone()));
        assert_eq!(bucket.len(), 1);

        // Adding same peer should move to end
        assert!(bucket.add_peer(peer1.clone()));
        assert_eq!(bucket.len(), 1);
    }

    #[test]
    fn test_kbucket_full() {
        let mut bucket = KBucket::new(2);

        let peer1 = PeerInfo::new(PeerId::random(), vec![], [0u8; 32], [0u8; 32]);
        let peer2 = PeerInfo::new(PeerId::random(), vec![], [0u8; 32], [0u8; 32]);
        let peer3 = PeerInfo::new(PeerId::random(), vec![], [0u8; 32], [0u8; 32]);

        assert!(bucket.add_peer(peer1));
        assert!(bucket.add_peer(peer2));
        assert!(!bucket.add_peer(peer3)); // Bucket full, can't add

        assert_eq!(bucket.len(), 2);
    }

    #[test]
    fn test_routing_table() {
        let local_id = PeerId::random();
        let mut table = RoutingTable::new(local_id, 20);

        let peer1 = PeerInfo::new(
            PeerId::random(),
            vec!["127.0.0.1:9001".to_string()],
            [0u8; 32],
            [0u8; 32],
        );

        assert!(table.add_peer(peer1.clone()).unwrap());
        assert_eq!(table.peer_count(), 1);

        // Don't add ourselves
        let self_peer = PeerInfo::new(local_id, vec![], [0u8; 32], [0u8; 32]);
        assert!(!table.add_peer(self_peer).unwrap());
    }

    #[test]
    fn test_find_closest() {
        let local_id = PeerId::random();
        let mut table = RoutingTable::new(local_id, 20);

        // Add several peers
        for _ in 0..10 {
            let peer = PeerInfo::new(PeerId::random(), vec![], [0u8; 32], [0u8; 32]);
            table.add_peer(peer).unwrap();
        }

        let target = PeerId::random();
        let closest = table.find_closest(&target, 5);

        assert_eq!(closest.len(), 5);
    }
}