typebit 0.1.8

//! Peer connection state and wire-level bookkeeping.
//!
//! A [`Peer`] owns everything about one TCP connection: phase, availability,
//! extension state, rates, reputation (anti-leech ledger), and the outgoing
//! byte buffer. The choke/unchoke **algorithm** itself lives in
//! [`crate::leech`] (reciprocity scoring, snubs, bans); this module only
//! supplies the per-connection state it operates on.

use crate::bitfield::Bitfield;
use crate::leech::PeerReputation;
use crate::monitoring::DiscoverySource;
use crate::platform::{ConnId, NetAddr};
use crate::wire::{ExtHandshake, Message, MessageStream};
use alloc::vec::Vec;

/// Lifecycle phase of a peer connection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PeerPhase {
    /// Connecting (non-blocking connect in progress).
    Connecting,
    /// Our handshake sent; waiting for theirs.
    Handshake,
    /// Handshake complete; exchanging bitfield/extension messages.
    Ready,
    /// Dead.
    Closed,
}

/// Per-connection peer state.
#[derive(Debug)]
pub struct Peer {
    /// Host connection handle.
    pub id: ConnId,
    /// Remote endpoint.
    pub addr: NetAddr,
    /// Lifecycle phase.
    pub phase: PeerPhase,
    /// Their peer id (after handshake).
    pub peer_id: Option<[u8; 20]>,
    /// Their reserved bits.
    pub reserved: [u8; 8],
    /// We are choking them.
    pub am_choking: bool,
    /// We are interested in them.
    pub am_interested: bool,
    /// They are choking us.
    pub peer_choking: bool,
    /// They are interested in us.
    pub peer_interested: bool,
    /// Their piece availability (empty if have_all/have_none).
    pub have: Bitfield,
    /// `have_all` received (fast extension).
    pub have_all: bool,
    /// `have_none` received.
    pub have_none: bool,
    /// Their extended-handshake data.
    pub ext: Option<ExtHandshake>,
    /// `ut_metadata` extended id (if any).
    pub ext_metadata: Option<u8>,
    /// `ut_pex` extended id (if any).
    pub ext_pex: Option<u8>,
    /// They are a seed (have everything).
    pub is_seed: bool,
    /// Supports fast extension.
    pub fast: bool,
    /// Supports DHT.
    pub supports_dht: bool,
    /// Supports v2 (BEP-52).
    pub supports_v2: bool,
    /// Bytes received total.
    pub down_total: u64,
    /// Bytes sent total.
    pub up_total: u64,
    /// Download rate (bytes/s, smoothed).
    pub down_rate: u32,
    /// Upload rate (bytes/s, smoothed).
    pub up_rate: u32,
    /// Timestamp of last inbound activity.
    pub last_active: u64,
    /// Timestamp of last data received (for snub detection).
    pub last_data_in: u64,
    /// Timestamp of the last *real* interaction (piece request/data) — keep-alives
    /// and metadata fetches don't count; drives the metadata-served idle disconnect.
    pub last_real_at: u64,
    /// Timestamp of connection start.
    pub connected_at: u64,
    /// ut_metadata pieces of OUR info dict we have served to this peer.
    pub metadata_served_pieces: u32,
    /// Total ut_metadata pieces of our info dict (0 = nothing to serve).
    pub metadata_total_pieces: u32,
    /// True once every metadata piece has been served (ms of completion).
    pub metadata_served: bool,
    /// When the last metadata piece was served (ms; 0 = never).
    pub metadata_served_at: u64,
    /// Snubbed (no data for a while despite requests).
    pub snubbed: bool,
    /// This peer holds an optimistic-unchoke slot.
    pub optimistic: bool,
    /// Outstanding request count on this connection.
    pub requests_in_flight: u32,
    /// Consecutive request timeouts on this connection (mechanism 2). Reset
    /// to zero as soon as the peer actually delivers a block; when it
    /// reaches the session's `max_request_timeouts` the peer is
    /// disconnected and blacklisted.
    pub requests_timed_out: u32,
    /// Max request pipeline depth for THIS connection (mechanism 1). The
    /// session copies its configured `request_pipeline` here at attach time
    /// so the configured value is actually enforced (the old code ignored
    /// the config and always used the crate constant).
    pub max_pipeline: u32,
    /// The piece this peer is committed to filling
    pub current_piece: Option<u32>,
    /// Last time they requested a block from us (ms; 0 = never). Drives
    /// the idle-slot (bandwidth squatting) detection.
    pub last_request_at: u64,
    /// Blocks we have served them on this connection.
    pub served_requests: u32,
    /// Cancels for blocks we never had outstanding to them (abuse signal).
    pub spurious_cancels: u32,
    /// Structurally invalid requests (abuse signal).
    pub invalid_requests: u32,
    /// Outgoing byte buffer (engine drains via host).
    pub out: Vec<u8>,
    /// Incoming message stream.
    pub msgs: MessageStream,
    /// Handshake prefix buffer.
    pub handshake_buf: Vec<u8>,
    /// How this peer was found.
    pub source: DiscoverySource,
    /// Anti-leech reputation ledger (client, corrupt blocks, violations).
    pub rep: PeerReputation,
    /// Last PEX message sent.
    pub pex_sent_at: u64,
    /// Rate-sampling window accumulators.
    pub window_down: u64,
    /// Bytes uploaded in the current window.
    pub window_up: u64,
    /// When the current rate window started.
    pub window_started: u64,
}

impl Peer {
    /// Create a new peer (outbound or inbound). `max_pipeline` is the
    /// request pipeline depth this connection is allowed (mechanism 1);
    /// pass the session's configured `request_pipeline`.
    pub fn new(
        id: ConnId,
        addr: NetAddr,
        piece_count: u32,
        source: DiscoverySource,
        max_pipeline: u32,
    ) -> Self {
        Peer {
            id,
            addr,
            phase: PeerPhase::Connecting,
            peer_id: None,
            reserved: [0u8; 8],
            am_choking: true,
            am_interested: false,
            peer_choking: true,
            peer_interested: false,
            have: Bitfield::new(piece_count),
            have_all: false,
            have_none: false,
            ext: None,
            ext_metadata: None,
            ext_pex: None,
            is_seed: false,
            fast: false,
            supports_dht: false,
            supports_v2: false,
            down_total: 0,
            up_total: 0,
            down_rate: 0,
            up_rate: 0,
            last_active: 0,
            last_data_in: 0,
            last_real_at: 0,
            connected_at: 0,
            metadata_served_pieces: 0,
            metadata_total_pieces: 0,
            metadata_served: false,
            metadata_served_at: 0,
            snubbed: false,
            optimistic: false,
            requests_in_flight: 0,
            requests_timed_out: 0,
            max_pipeline: max_pipeline.max(1),
            current_piece: None,
            last_request_at: 0,
            served_requests: 0,
            spurious_cancels: 0,
            invalid_requests: 0,
            out: Vec::with_capacity(16 * 1024),
            msgs: MessageStream::new(),
            handshake_buf: Vec::new(),
            source,
            rep: PeerReputation::default(),
            pex_sent_at: 0,
            window_down: 0,
            window_up: 0,
            window_started: 0,
        }
    }

    /// Whether we know the peer has piece `p`.
    pub fn has_piece(&self, p: u32) -> bool {
        self.have_all || self.have.get(p)
    }

    /// Whether we should be interested (they have a piece we WANT).
    ///
    /// "Want" is **selective-download aware**: a piece we lack is only
    /// wanted when its per-piece priority is non-zero (it belongs to a
    /// selected file). A peer that only carries skipped-file pieces must
    /// not trigger `Interested` — otherwise we'd hold unchoke slots and
    /// waste handshakes on peers we can never download from.
    ///
    /// Deliberately **independent of choke state**: gating interest on
    /// `peer_choking` deadlocks fresh downloads (a seed unchokes only
    /// interested peers, but we'd never declare interest while choked).
    /// Send `Interested` as soon as their availability shows a wanted piece.
    pub fn should_be_interested(&self, our_have: &Bitfield, priorities: &[i64]) -> bool {
        if self.have_none {
            return false;
        }
        if self.have_all {
            // Interested iff some SELECTED piece is still missing.
            let n = our_have.len().min(priorities.len() as u32);
            for p in 0..n {
                if priorities[p as usize] > 0 && !our_have.get(p) {
                    return true;
                }
            }
            return false;
        }
        // any piece they have that we lack AND want
        let n = self.have.len().min(our_have.len());
        let mut i = self.have.first_set();
        while let Some(p) = i {
            if p >= n {
                break;
            }
            if priorities.get(p as usize).copied().unwrap_or(1) > 0 && !our_have.get(p) {
                return true;
            }
            i = self.have.next_set_from(p + 1);
        }
        false
    }

    /// Max request pipeline depth for this peer (mechanism 1: the session's
    /// configured per-peer pipeline, copied at attach time).
    pub fn max_pipeline(&self) -> u32 {
        self.max_pipeline
    }

    /// Queue a message for sending.
    pub fn send(&mut self, m: &Message) {
        self.out.extend_from_slice(&m.encode());
    }

    /// Whether the out buffer is empty.
    pub fn flushed(&self) -> bool {
        self.out.is_empty()
    }

    /// Reset the rate window (called on a fixed cadence).
    pub fn roll_window(&mut self, now: u64) {
        let span_ms = now.saturating_sub(self.window_started).max(1);
        self.down_rate = ((self.window_down * 1000) / span_ms) as u32;
        self.up_rate = ((self.window_up * 1000) / span_ms) as u32;
        self.window_down = 0;
        self.window_up = 0;
        self.window_started = now;
    }

    /// Record downloaded bytes.
    pub fn on_data_in(&mut self, n: usize, now: u64) {
        self.down_total += n as u64;
        self.window_down += n as u64;
        self.last_active = now;
        self.last_data_in = now;
        self.snubbed = false;
    }

    /// Record uploaded bytes.
    pub fn on_data_out(&mut self, n: usize, now: u64) {
        self.up_total += n as u64;
        self.window_up += n as u64;
        self.last_active = now;
    }

    /// Refresh the snub flag: we are snubbed when we expect data from them
    /// (they unchoked us and we have outstanding requests) but they stay
    /// quiet for `timeout_ms`. The flag clears as soon as the situation
    /// improves (data arrives or expectations end).
    pub fn refresh_snub(&mut self, now: u64, timeout_ms: u64) {
        self.snubbed = self.phase == PeerPhase::Ready
            && !self.peer_choking
            && self.requests_in_flight > 0
            && now.saturating_sub(self.last_data_in) > timeout_ms;
    }
}

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

    fn peer(id: ConnId) -> Peer {
        let mut p = Peer::new(
            id,
            NetAddr::V4([127, 0, 0, 1], 6881),
            64,
            DiscoverySource::Tracker,
            32,
        );
        p.phase = PeerPhase::Ready;
        p
    }

    #[test]
    fn snub_detection_and_clear() {
        let mut p = peer(1);
        p.peer_choking = false;
        p.requests_in_flight = 2;
        p.last_data_in = 0;
        p.refresh_snub(1000, 500);
        assert!(p.snubbed);
        // data arrives → cleared
        p.on_data_in(16, 1001);
        assert!(!p.snubbed);
        // no outstanding requests → never snubbed
        p.requests_in_flight = 0;
        p.last_data_in = 0;
        p.refresh_snub(2000, 500);
        assert!(!p.snubbed);
    }

    #[test]
    fn interested_detection() {
        let mut p = peer(1);
        p.peer_choking = false;
        p.have.set(5);
        let mut ours = Bitfield::new(64);
        ours.set(5);
        let all = [1i64; 64];
        assert!(!p.should_be_interested(&ours, &all));
        ours.clear(5);
        assert!(p.should_be_interested(&ours, &all));
        // choking must NOT suppress interest (interest ≠ unchoke) — a seed
        // that unchokes only interested peers would otherwise deadlock us.
        p.peer_choking = true;
        assert!(p.should_be_interested(&ours, &all));
        // have_none (they have nothing) still suppresses interest
        p.have_none = true;
        assert!(!p.should_be_interested(&ours, &all));
    }

    #[test]
    fn interest_is_selective_download_aware() {
        // The peer owns pieces 5 and 6. We lack both, but piece 6 belongs
        // to a SKIPPED file (priority 0) — we must NOT declare interest.
        let mut p = peer(1);
        p.have.set(5);
        p.have.set(6);
        let ours = Bitfield::new(64);
        let mut prio = [1i64; 64];
        prio[6] = 0; // skipped
        assert!(p.should_be_interested(&ours, &prio), "wanted piece 5");
        prio[5] = 0; // now every piece the peer has is skipped
        assert!(
            !p.should_be_interested(&ours, &prio),
            "a peer with only skipped-file pieces must not trigger interest"
        );
        // have_all seed with every selected piece already verified → no
        // interest either (avoids useless unchoke slots).
        let mut p2 = peer(2);
        p2.have_all = true;
        let mut have = Bitfield::new(64);
        for (i, &pr) in prio.iter().enumerate() {
            if pr > 0 {
                have.set(i as u32);
            }
        }
        assert!(!p2.should_be_interested(&have, &prio));
        // Missing a SELECTED piece (7) → interested again.
        have.clear(7);
        assert!(p2.should_be_interested(&have, &prio));
    }
}