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, request_queue: VecDeque, requested_queue: HashSet, received_pieces: VecDeque, queue: VecDeque, have_queue: VecDeque, } impl Peer { pub fn id(&self) -> U160 { self.id } pub fn dump_queue(&mut self) -> VecDeque { 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 { self.received_pieces.pop_front() } pub fn next_request(&mut self) -> Option { 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 { self.queue.pop_front() } pub fn next_have(&mut self) -> Option { self.have_queue.pop_front() } }