dht/src/peer_lookup.rs

use crate::budget::RateBucket;
use crate::krpc::{encode_get_peers_query, for_each_response_node, for_each_response_peer};
use crate::node_id::TransactionId;
use crate::protocol::DhtResponse;
use crate::routing_snapshot::{RoutingSnapshot, xor_distance_cmp};
use crate::runtime_stats::DhtRuntimeStats;
use crate::server::HashDiscovered;
use crate::types::{NetMode, NodeTuple, PeerLookupOptions};
use ahash::{AHashMap, AHashSet};
use arc_swap::ArcSwap;
use bytes::BytesMut;
#[cfg(feature = "metrics")]
use metrics::counter;
use std::collections::VecDeque;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

const LOOKUP_TID_TAG: u8 = 0xa5;
const MAX_QUERIES_PER_LOOKUP: usize = 12;
const MAX_CONCURRENT_QUERIES_PER_LOOKUP: usize = 4;
const MAX_FRONTIER_NODES: usize = 64;
const MAX_PEERS_PER_LOOKUP: usize = 12;
const LOOKUP_TIMEOUT: Duration = Duration::from_secs(2);
const QUERY_TIMEOUT: Duration = Duration::from_millis(500);
const MAINTENANCE_INTERVAL: Duration = Duration::from_millis(25);
const REQUEST_CHANNEL_CAPACITY: usize = 1_024;
const RESPONSE_CHANNEL_CAPACITY: usize = 4_096;

#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
struct PendingKey {
    addr: SocketAddr,
    tid: TransactionId,
}

#[derive(Debug, Clone, Copy)]
struct PendingQuery {
    lookup_id: u64,
    deadline: Instant,
}

struct LookupState {
    info_hash: [u8; 20],
    info_hash_hex: String,
    frontier: Vec<NodeTuple>,
    seen_nodes: AHashSet<SocketAddr>,
    peers: AHashSet<SocketAddr>,
    queried: usize,
    outstanding: usize,
    deadline: Instant,
}

impl LookupState {
    fn pop_closest(&mut self) -> Option<NodeTuple> {
        let index = self
            .frontier
            .iter()
            .enumerate()
            .min_by(|(_, left), (_, right)| xor_distance_cmp(&left.id, &right.id, &self.info_hash))
            .map(|(index, _)| index)?;
        Some(self.frontier.swap_remove(index))
    }

    fn is_complete(&self, now: Instant) -> bool {
        self.deadline <= now
            || self.peers.len() >= MAX_PEERS_PER_LOOKUP
            || (self.outstanding == 0
                && (self.queried >= MAX_QUERIES_PER_LOOKUP || self.frontier.is_empty()))
    }
}

struct LookupResponse {
    remote_addr: SocketAddr,
    tid: TransactionId,
    response: DhtResponse,
}

#[derive(Clone)]
pub(crate) struct PeerLookupHandle {
    request_tx: mpsc::Sender<[u8; 20]>,
    response_tx: mpsc::Sender<LookupResponse>,
    runtime_stats: DhtRuntimeStats,
}

pub(crate) struct PeerLookupRuntime {
    pub(crate) options: PeerLookupOptions,
    pub(crate) stats: DhtRuntimeStats,
    pub(crate) shutdown: CancellationToken,
}

impl PeerLookupHandle {
    pub(crate) fn request_sender(&self) -> mpsc::Sender<[u8; 20]> {
        self.request_tx.clone()
    }

    pub(crate) fn route_response(
        &self,
        remote_addr: SocketAddr,
        tid: TransactionId,
        response: DhtResponse,
    ) {
        if self
            .response_tx
            .try_send(LookupResponse {
                remote_addr,
                tid,
                response,
            })
            .is_err()
        {
            self.runtime_stats.peer_lookup_response_dropped();
            #[cfg(feature = "metrics")]
            counter!("dht_peer_lookup_dropped_total", "reason" => "response_queue_full")
                .increment(1);
        }
    }
}

pub(crate) fn is_peer_lookup_tid(tid: &TransactionId) -> bool {
    tid[0] == LOOKUP_TID_TAG
}

pub(crate) fn spawn_peer_lookup(
    netmode: NetMode,
    local_id: [u8; 20],
    sockets: &std::collections::HashMap<SocketAddr, Arc<UdpSocket>>,
    snapshot: Arc<ArcSwap<RoutingSnapshot>>,
    hash_tx: mpsc::Sender<HashDiscovered>,
    runtime: PeerLookupRuntime,
) -> PeerLookupHandle {
    let PeerLookupRuntime {
        options,
        stats,
        shutdown,
    } = runtime;
    let (request_tx, request_rx) = mpsc::channel(REQUEST_CHANNEL_CAPACITY);
    let (response_tx, response_rx) = mpsc::channel(RESPONSE_CHANNEL_CAPACITY);
    let socket_v4 = sockets
        .iter()
        .find_map(|(addr, socket)| addr.is_ipv4().then(|| socket.clone()));
    let socket_v6 = sockets
        .iter()
        .find_map(|(addr, socket)| addr.is_ipv6().then(|| socket.clone()));
    let actor = PeerLookupActor {
        netmode,
        local_id,
        socket_v4,
        socket_v6,
        snapshot,
        hash_tx,
        request_rx,
        response_rx,
        request_budget: RateBucket::per_second(
            options.max_lookups_per_second,
            options.burst,
            true,
            Instant::now(),
        ),
        max_active_lookups: options.max_active_lookups,
        active: AHashMap::new(),
        pending: AHashMap::new(),
        pending_expiry: VecDeque::new(),
        next_lookup_id: 1,
        next_tid: 1,
        runtime_stats: stats.clone(),
        shutdown,
    };
    tokio::spawn(actor.run());
    PeerLookupHandle {
        request_tx,
        response_tx,
        runtime_stats: stats,
    }
}

struct PeerLookupActor {
    netmode: NetMode,
    local_id: [u8; 20],
    socket_v4: Option<Arc<UdpSocket>>,
    socket_v6: Option<Arc<UdpSocket>>,
    snapshot: Arc<ArcSwap<RoutingSnapshot>>,
    hash_tx: mpsc::Sender<HashDiscovered>,
    request_rx: mpsc::Receiver<[u8; 20]>,
    response_rx: mpsc::Receiver<LookupResponse>,
    request_budget: RateBucket,
    max_active_lookups: usize,
    active: AHashMap<u64, LookupState>,
    pending: AHashMap<PendingKey, PendingQuery>,
    pending_expiry: VecDeque<(Instant, PendingKey)>,
    next_lookup_id: u64,
    next_tid: u64,
    runtime_stats: DhtRuntimeStats,
    shutdown: CancellationToken,
}

impl PeerLookupActor {
    async fn run(mut self) {
        let mut maintenance = tokio::time::interval(MAINTENANCE_INTERVAL);
        maintenance.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        loop {
            tokio::select! {
                biased;
                _ = self.shutdown.cancelled() => break,
                response = self.response_rx.recv() => {
                    let Some(response) = response else { break };
                    self.handle_response(response, Instant::now()).await;
                }
                _ = maintenance.tick() => self.expire(Instant::now()).await,
                request = self.request_rx.recv() => {
                    let Some(info_hash) = request else { break };
                    self.start_lookup(info_hash, Instant::now()).await;
                }
            }
        }
    }

    async fn start_lookup(&mut self, info_hash: [u8; 20], now: Instant) {
        self.runtime_stats.peer_lookup_requested();
        #[cfg(feature = "metrics")]
        counter!("dht_peer_lookup_requests_total").increment(1);
        if self.active.len() >= self.max_active_lookups || !self.request_budget.try_take_one(now) {
            self.runtime_stats.peer_lookup_rate_limited();
            #[cfg(feature = "metrics")]
            counter!("dht_peer_lookup_dropped_total", "reason" => "rate_limit").increment(1);
            return;
        }

        let filter_ipv6 = match (self.socket_v4.is_some(), self.socket_v6.is_some()) {
            (true, false) => Some(false),
            (false, true) => Some(true),
            _ => None,
        };
        let frontier =
            self.snapshot
                .load()
                .closest_nodes(&info_hash, MAX_QUERIES_PER_LOOKUP, filter_ipv6);
        if frontier.is_empty() {
            self.runtime_stats.peer_lookup_empty();
            return;
        }

        let lookup_id = self.next_lookup_id;
        self.next_lookup_id = self.next_lookup_id.wrapping_add(1).max(1);
        let seen_nodes = frontier.iter().map(|node| node.addr).collect();
        self.active.insert(
            lookup_id,
            LookupState {
                info_hash,
                info_hash_hex: hex::encode(info_hash),
                frontier,
                seen_nodes,
                peers: AHashSet::new(),
                queried: 0,
                outstanding: 0,
                deadline: now + LOOKUP_TIMEOUT,
            },
        );
        self.runtime_stats.peer_lookup_started();
        #[cfg(feature = "metrics")]
        counter!("dht_peer_lookup_started_total").increment(1);
        self.dispatch_more(lookup_id, now).await;
    }

    async fn dispatch_more(&mut self, lookup_id: u64, now: Instant) {
        loop {
            let Some((node, info_hash)) = self.active.get_mut(&lookup_id).and_then(|state| {
                if state.deadline <= now
                    || state.peers.len() >= MAX_PEERS_PER_LOOKUP
                    || state.queried >= MAX_QUERIES_PER_LOOKUP
                    || state.outstanding >= MAX_CONCURRENT_QUERIES_PER_LOOKUP
                {
                    return None;
                }
                let node = state.pop_closest()?;
                state.queried += 1;
                Some((node, state.info_hash))
            }) else {
                break;
            };

            let tid = self.next_transaction_id();
            let key = PendingKey {
                addr: node.addr,
                tid,
            };
            let mut buffer = BytesMut::with_capacity(128);
            encode_get_peers_query(&mut buffer, &tid, &info_hash, &self.local_id);
            let socket = if node.addr.is_ipv4() {
                self.socket_v4.clone()
            } else {
                self.socket_v6.clone()
            };
            let sent = match socket {
                Some(socket) => socket.send_to(&buffer, node.addr).await.is_ok(),
                None => false,
            };
            if !sent {
                self.runtime_stats.peer_lookup_send_failed();
                continue;
            }

            let deadline = now + QUERY_TIMEOUT;
            self.pending.insert(
                key,
                PendingQuery {
                    lookup_id,
                    deadline,
                },
            );
            self.pending_expiry.push_back((deadline, key));
            if let Some(state) = self.active.get_mut(&lookup_id) {
                state.outstanding += 1;
            }
            self.runtime_stats.udp_sent(buffer.len());
            self.runtime_stats.peer_lookup_query();
            #[cfg(feature = "metrics")]
            counter!("dht_peer_lookup_queries_total").increment(1);
        }
        self.finish_if_complete(lookup_id, now);
    }

    async fn handle_response(&mut self, event: LookupResponse, now: Instant) {
        let key = PendingKey {
            addr: event.remote_addr,
            tid: event.tid,
        };
        let Some(pending) = self.pending.remove(&key) else {
            return;
        };
        let Some(state) = self.active.get_mut(&pending.lookup_id) else {
            return;
        };
        state.outstanding = state.outstanding.saturating_sub(1);
        self.runtime_stats.peer_lookup_response();

        let mut discovered = Vec::new();
        for_each_response_peer(&event.response, self.netmode, |peer| {
            if state.peers.len() < MAX_PEERS_PER_LOOKUP && state.peers.insert(peer) {
                discovered.push(peer);
            }
        });
        let mut response_nodes = Vec::new();
        for_each_response_node(&event.response, self.netmode, |node| {
            response_nodes.push(node)
        });
        for node in response_nodes {
            if state.frontier.len() >= MAX_FRONTIER_NODES {
                break;
            }
            if state.seen_nodes.insert(node.addr) {
                state.frontier.push(node);
            }
        }
        let hash = state.info_hash_hex.clone();
        let lookup_id = pending.lookup_id;
        let _ = state;

        for peer in discovered {
            let event = HashDiscovered {
                info_hash: hash.clone(),
                peer_addr: peer,
                discovered_at: now,
            };
            if self.hash_tx.try_send(event).is_ok() {
                self.runtime_stats.peer_lookup_peer_found();
                #[cfg(feature = "metrics")]
                counter!("dht_peer_lookup_peers_found_total").increment(1);
            } else {
                self.runtime_stats.peer_lookup_output_dropped();
                #[cfg(feature = "metrics")]
                counter!("dht_peer_lookup_dropped_total", "reason" => "hash_queue_full")
                    .increment(1);
            }
        }
        self.dispatch_more(lookup_id, now).await;
    }

    async fn expire(&mut self, now: Instant) {
        let mut affected = AHashSet::new();
        while let Some((deadline, key)) = self.pending_expiry.front().copied() {
            if deadline > now {
                break;
            }
            self.pending_expiry.pop_front();
            let should_remove = self
                .pending
                .get(&key)
                .is_some_and(|pending| pending.deadline == deadline);
            if !should_remove {
                continue;
            }
            let pending = self.pending.remove(&key).expect("pending lookup exists");
            if let Some(state) = self.active.get_mut(&pending.lookup_id) {
                state.outstanding = state.outstanding.saturating_sub(1);
                affected.insert(pending.lookup_id);
            }
            self.runtime_stats.peer_lookup_timeout();
        }
        for lookup_id in affected {
            self.dispatch_more(lookup_id, now).await;
        }
        let expired: Vec<_> = self
            .active
            .iter()
            .filter_map(|(lookup_id, state)| (state.deadline <= now).then_some(*lookup_id))
            .collect();
        for lookup_id in expired {
            self.finish_lookup(lookup_id);
        }
    }

    fn finish_if_complete(&mut self, lookup_id: u64, now: Instant) {
        if self
            .active
            .get(&lookup_id)
            .is_some_and(|state| state.is_complete(now))
        {
            self.finish_lookup(lookup_id);
        }
    }

    fn finish_lookup(&mut self, lookup_id: u64) {
        self.active.remove(&lookup_id);
        self.pending
            .retain(|_, pending| pending.lookup_id != lookup_id);
    }

    fn next_transaction_id(&mut self) -> TransactionId {
        let mut tid = self.next_tid.to_be_bytes();
        tid[0] = LOOKUP_TID_TAG;
        self.next_tid = self.next_tid.wrapping_add(1).max(1);
        tid
    }
}

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

    #[test]
    fn lookup_transaction_ids_have_a_reserved_tag() {
        let tid = [LOOKUP_TID_TAG, 1, 2, 3, 4, 5, 6, 7];
        assert!(is_peer_lookup_tid(&tid));
        assert!(!is_peer_lookup_tid(&[0; 8]));
    }

    #[tokio::test]
    async fn lookup_response_feeds_discovered_peer_back_to_metadata_scheduler() {
        let local_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let remote_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let remote_addr = remote_socket.local_addr().unwrap();
        let local_addr = local_socket.local_addr().unwrap();
        let sockets = std::collections::HashMap::from([(local_addr, local_socket)]);
        let snapshot = Arc::new(ArcSwap::from_pointee(RoutingSnapshot::from_nodes(
            vec![NodeTuple {
                id: [9; 20],
                addr: remote_addr,
            }],
            1,
        )));
        let (hash_tx, mut hash_rx) = mpsc::channel(4);
        let stats = DhtRuntimeStats::default();
        let shutdown = CancellationToken::new();
        let handle = spawn_peer_lookup(
            NetMode::Ipv4Only,
            [7; 20],
            &sockets,
            snapshot,
            hash_tx,
            PeerLookupRuntime {
                options: PeerLookupOptions::default(),
                stats: stats.clone(),
                shutdown: shutdown.clone(),
            },
        );
        handle.request_tx.send([3; 20]).await.unwrap();

        let mut buffer = [0u8; 512];
        let (len, source) =
            tokio::time::timeout(Duration::from_secs(1), remote_socket.recv_from(&mut buffer))
                .await
                .unwrap()
                .unwrap();
        assert_eq!(source, local_addr);
        let query: DhtMessage = serde_bencode::from_bytes(&buffer[..len]).unwrap();
        assert_eq!(query.q.as_deref(), Some("get_peers"));
        let tid: TransactionId = query.t.as_ref().try_into().unwrap();
        handle.route_response(
            remote_addr,
            tid,
            DhtResponse {
                id: Some(serde_bytes::ByteBuf::from(vec![9; 20])),
                nodes: None,
                nodes6: None,
                values: Some(vec![serde_bytes::ByteBuf::from(vec![
                    8, 8, 4, 4, 0x1a, 0xe1,
                ])]),
            },
        );

        let event = tokio::time::timeout(Duration::from_secs(1), hash_rx.recv())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(event.info_hash, hex::encode([3; 20]));
        assert_eq!(event.peer_addr, "8.8.4.4:6881".parse().unwrap());
        let snapshot = stats.snapshot();
        assert_eq!(snapshot.peer_lookup_started, 1);
        assert_eq!(snapshot.peer_lookup_queries, 1);
        assert_eq!(snapshot.peer_lookup_responses, 1);
        assert_eq!(snapshot.peer_lookup_peers_found, 1);
        shutdown.cancel();
    }
}