use crate::torrent_manager::TorrentManager; use crossbeam_channel::{unbounded, Receiver, Sender}; use reqwest::blocking::Client; use reqwest::header::HeaderMap; use std::collections::HashMap; use std::net::IpAddr; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::RwLock; use std::thread::JoinHandle; use torment_core::infohash::v1::U160; use torment_core::peer_storage::PeerStorageHolder; use torment_tracker::{ block_do_announce, Tracker, TrackerAnnounceRequest, TrackerAnnounceResponse, }; use url::Url; pub struct TrackerManager { id: U160, ip: IpAddr, port: u16, counter: AtomicUsize, trackers: HashMap>, url_index: HashMap, peer_storage: PeerStorageHolder, tracker_offshore_state: TrackerOffshoreState, } struct TrackerOffshoreState { threads: Vec>, tracker_request_sender: Sender<(usize, TrackerAnnounceRequest)>, tracker_request_receiver: Receiver<(usize, TrackerAnnounceRequest)>, tracker_response_sender: Sender<(usize, TrackerAnnounceResponse)>, tracker_response_receiver: Receiver<(usize, TrackerAnnounceResponse)>, } impl TrackerOffshoreState { fn new() -> TrackerOffshoreState { let (req_send, req_recv) = unbounded(); let (resp_send, resp_recv) = unbounded(); TrackerOffshoreState { threads: vec![], tracker_request_sender: req_send, tracker_request_receiver: req_recv, tracker_response_sender: resp_send, tracker_response_receiver: resp_recv, } } } pub type TrackerId = usize; fn tracker_thread( input: Receiver<(usize, TrackerAnnounceRequest)>, output: Sender<(usize, TrackerAnnounceResponse)>, ) { let mut map = HeaderMap::new(); map.insert( reqwest::header::USER_AGENT, format!("eater's Torment/{}", env!("CARGO_PKG_VERSION")) .parse() .unwrap(), ); let client = Client::builder().default_headers(map).build().unwrap(); for (tracker_id, req) in input { output .send((tracker_id, block_do_announce(req, &client))) .unwrap(); } } impl TrackerManager { pub fn new( ip: IpAddr, port: u16, peer_id: U160, peer_storage: PeerStorageHolder, ) -> TrackerManager { TrackerManager { id: peer_id, ip, port, counter: Default::default(), trackers: Default::default(), url_index: Default::default(), peer_storage, tracker_offshore_state: TrackerOffshoreState::new(), } } pub fn get_tracker_id(&mut self, url: &Url) -> Option { // dht is not a tracker bye if url.scheme() == "dht" { return None; } let url_str = url.as_str(); if let Some(id) = self.url_index.get(url_str) { return Some(*id); } let new_id = self.counter.fetch_add(1, Ordering::AcqRel); self.url_index.insert(url_str.to_string(), new_id); self.trackers.insert( new_id, RwLock::new(Tracker::new(url.clone(), self.peer_storage.clone())), ); Some(new_id) } pub fn announce(&mut self, tracker_id: usize, torrent: &TorrentManager) { let id = torrent.info_hash(); { let tracker = &self.trackers[&tracker_id].read().unwrap(); if !tracker.needs_update(id) { return; } } let tracker = &mut self.trackers[&tracker_id].write().unwrap(); println!( "Queueing announce {} on {}", torrent.info_hash(), tracker.url() ); let req = tracker.create_update_request( id, self.id, self.ip, self.port, torrent.uploaded(), torrent.downloaded(), torrent.bytes_left(), ); self.tracker_offshore_state .tracker_request_sender .send((tracker_id, req)) .unwrap() } pub fn house_keeping(&mut self) { while self.tracker_offshore_state.threads.len() < 5 { let sender = self.tracker_offshore_state.tracker_response_sender.clone(); let receiver = self.tracker_offshore_state.tracker_request_receiver.clone(); self.tracker_offshore_state.threads.push( std::thread::Builder::new() .name(format!( "tracker/{}", self.tracker_offshore_state.threads.len() )) .spawn(|| tracker_thread(receiver, sender)) .unwrap(), ); } while let Ok((tracker_id, resp)) = self .tracker_offshore_state .tracker_response_receiver .try_recv() { self.trackers[&tracker_id].write().unwrap().process(resp); } } }