torment/torment-peer/src/lib.rs

use crate::message::{Handshake, Message, PieceMessage, SelectionMessage};
use std::collections::{HashSet, VecDeque};
use std::net::SocketAddr;
use torment_core::infohash::v1::U160;
use torment_core::metainfo::MetaInfo;
use torment_core::Bitfield;

pub mod message;

#[derive(Copy, Clone, Debug)]
pub struct PeerState {
    chocked: bool,
    interested: bool,
}

impl PeerState {
    const fn new() -> PeerState {
        PeerState {
            chocked: true,
            interested: false,
        }
    }
}

impl Default for PeerState {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Copy, Clone)]
pub enum PeerProtocol {
    TCP,
    MuTP,
}

#[derive(Debug, Clone)]
pub struct Peer {
    id: U160,
    info_hash: U160,
    address: SocketAddr,
    protocol: PeerProtocol,
    our_state: PeerState,
    their_state: PeerState,
    has: Bitfield,
    request_state: HashSet<SelectionMessage>,
    request_queue: VecDeque<SelectionMessage>,
    requested_queue: HashSet<SelectionMessage>,
    received_pieces: VecDeque<PieceMessage>,
    queue: VecDeque<Message>,
    have_queue: VecDeque<u32>,
}

impl Peer {
    pub fn id(&self) -> U160 {
        self.id
    }

    pub fn dump_queue(&mut self) -> VecDeque<Message> {
        std::mem::replace(&mut self.queue, VecDeque::new())
    }

    pub fn new(
        address: SocketAddr,
        protocol: PeerProtocol,
        header: Handshake,
        meta_info: &MetaInfo,
    ) -> Self {
        Peer {
            id: header.peer_id(),
            info_hash: header.info_hash(),
            address,
            protocol,
            our_state: PeerState::new(),
            their_state: PeerState::new(),
            has: Bitfield::with_size(meta_info.pieces()),
            request_state: Default::default(),
            request_queue: Default::default(),
            requested_queue: Default::default(),
            received_pieces: Default::default(),
            queue: Default::default(),
            have_queue: Default::default(),
        }
    }

    pub fn has(&mut self) -> Bitfield {
        self.has.clone()
    }

    pub fn process(&mut self, message: Message) -> bool {
        match message {
            Message::Choke => self.their_state.chocked = true,
            Message::Unchoke => self.their_state.chocked = false,
            Message::Interested => self.their_state.interested = true,
            Message::NotInterested => self.their_state.interested = false,
            Message::Have(nr) => {
                self.has.set(nr);
                self.have_queue.push_back(nr);
            }
            Message::Bitfield(bitfield) => self.has.add(bitfield),
            Message::Request(selection) => {
                self.request_queue.push_back(selection);
                self.request_state.insert(selection);
            }
            Message::Cancel(selection) => {
                self.request_state.remove(&selection);
            }
            Message::Piece(piece) => {
                self.requested_queue.remove(&SelectionMessage::new(
                    piece.index(),
                    piece.offset(),
                    piece.length(),
                ));
                self.received_pieces.push_back(piece);
            }
        }

        true
    }

    pub fn count_open_requests(&self) -> usize {
        self.requested_queue.len()
    }

    pub fn is_choked(&self) -> bool {
        self.their_state.chocked
    }

    pub fn is_interested(&self) -> bool {
        self.their_state.interested
    }

    pub fn we_interested(&self) -> bool {
        self.our_state.interested
    }

    pub fn we_choked(&self) -> bool {
        self.our_state.chocked
    }

    pub fn set_we_choked(&mut self, choked: bool) {
        if self.our_state.chocked != choked {
            self.queue.push_back(if choked {
                Message::Choke
            } else {
                Message::Unchoke
            })
        }
        self.our_state.chocked = choked;
    }

    pub fn send_bitfield(&mut self, bitfield: &Bitfield) {
        self.queue.push_back(Message::Bitfield(bitfield.clone()));
    }

    pub fn set_we_interested(&mut self, interested: bool) {
        self.our_state.interested = interested;
    }

    pub fn request(&mut self, index: u32, offset: u32, length: u32) {
        let msg = SelectionMessage::new(index, offset, length);
        if self.requested(msg) {
            self.queue.push_back(Message::Request(msg));
        }
    }

    pub fn requested(&mut self, selection: SelectionMessage) -> bool {
        self.requested_queue.insert(selection)
    }

    pub fn unchoke(&mut self) {
        if !self.our_state.chocked {
            return;
        }

        self.queue.push_back(Message::Unchoke);
    }

    pub fn choke(&mut self) {
        if self.our_state.chocked {
            return;
        }

        self.queue.push_back(Message::Choke);
    }

    pub fn interested(&mut self) {
        if self.our_state.interested {
            return;
        }

        self.our_state.interested = true;
        self.queue.push_back(Message::Interested);
    }

    pub fn lost_interest(&mut self) {
        if !self.our_state.interested {
            return;
        }

        self.queue.push_back(Message::NotInterested);
    }

    pub fn next_piece(&mut self) -> Option<PieceMessage> {
        self.received_pieces.pop_front()
    }

    pub fn next_request(&mut self) -> Option<SelectionMessage> {
        while let Some(item) = self.request_queue.pop_front() {
            if self.request_state.remove(&item) {
                return Some(item);
            }
        }
        None
    }

    pub fn addr(&self) -> SocketAddr {
        self.address
    }

    pub fn has_piece(&self, index: u32) -> bool {
        self.has.get(index as u32)
    }

    pub fn next(&mut self) -> Option<Message> {
        self.queue.pop_front()
    }

    pub fn next_have(&mut self) -> Option<u32> {
        self.have_queue.pop_front()
    }
}