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>,
}

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

    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(),
        }
    }

    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),
            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) {
        self.our_state.chocked = choked;
    }

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

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

    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()
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
        assert_eq!(2 + 2, 4);
    }
}