use crate::torrent_manager::TorrentManager; use crate::tracker_manager::TrackerManager; use std::collections::{HashMap, HashSet, VecDeque}; use std::net::{IpAddr, SocketAddr}; use std::option::Option::Some; use torment_core::infohash::v1::U160; use torment_core::peer_storage::{PeerStorage, PeerStorageHolder}; use torment_peer::message::{Handshake, Message}; use torment_peer::PeerProtocol; pub struct QueuedMessage { pub message: Message, pub info_hash: U160, pub peer_id: U160, pub addr: SocketAddr, } pub struct SessionManager { id: U160, torrents: HashMap, torrents_ids: HashSet, peer_socket: HashMap<(U160, U160), SocketAddr>, message_queue: VecDeque, tracker_manager: TrackerManager, pub peer_storage: PeerStorageHolder, } impl SessionManager { pub fn peer_count(&self, info_hash: U160) -> usize { self.torrents[&info_hash].peer_count() } pub fn torrents(&self) -> Vec { self.torrents_ids.iter().copied().collect() } pub fn torrent_manager(&self, info_hash: U160) -> &TorrentManager { &self.torrents[&info_hash] } pub fn id(&self) -> U160 { self.id } pub fn tracker_manager(&self) -> &TrackerManager { &self.tracker_manager } pub fn tracker_manager_mut(&mut self) -> &mut TrackerManager { &mut self.tracker_manager } pub fn is_done(&self, info_hash: U160) -> bool { if let Some(torrent) = self.torrents.get(&info_hash) { torrent.is_done() } else { false } } pub fn dump_queue(&mut self) -> VecDeque { std::mem::replace(&mut self.message_queue, VecDeque::new()) } pub fn add_torrent_manager(&mut self, torrent: TorrentManager) { self.torrents_ids.insert(torrent.info_hash()); self.torrents.insert(torrent.info_hash(), torrent); } #[allow(dead_code)] fn remove_torrent_manager(&mut self, info_hash: U160) -> Option { self.torrents_ids.remove(&info_hash); self.torrents.remove(&info_hash) } pub fn handshake( &mut self, handshake: Handshake, addr: SocketAddr, protocol: PeerProtocol, ) -> bool { if let Some(torrent_manager) = self.torrents.get_mut(&handshake.info_hash()) { self.peer_socket .insert((handshake.peer_id(), handshake.info_hash()), addr); torrent_manager.handshake(handshake, addr, protocol) } else { false } } pub fn process(&mut self, torrent: U160, peer_id: U160, message: Message) -> bool { if let Some(torrent_manager) = self.torrents.get_mut(&torrent) { if torrent_manager.process(peer_id, message) { let queue = torrent_manager.dump_queue(); for (peer_id, message) in queue { if let Some(addr) = self.peer_socket.get(&(peer_id, torrent)).copied() { self.message_queue.push_back(QueuedMessage { message, addr, peer_id, info_hash: torrent, }); } } true } else { false } } else { false } } pub fn new(ip: IpAddr, port: u16, peer_id: U160) -> SessionManager { let peer_storage = PeerStorageHolder::new(PeerStorage::new()); SessionManager { id: peer_id, torrents: Default::default(), tracker_manager: TrackerManager::new(ip, port, peer_id, peer_storage.clone()), peer_storage, torrents_ids: Default::default(), peer_socket: Default::default(), message_queue: Default::default(), } } pub fn next(&mut self) -> Option { self.message_queue.pop_front() } pub fn house_keeping(&mut self) { for (info_hash, torrent) in &mut self.torrents { torrent.house_keeping(); let queue = torrent.dump_queue(); for (peer_id, message) in queue { if let Some(addr) = self.peer_socket.get(&(peer_id, *info_hash)).copied() { self.message_queue.push_back(QueuedMessage { message, addr, peer_id, info_hash: *info_hash, }); } } } } pub fn announce(&mut self) { println!("Announcing torrents..."); for info_hash in self.torrents_ids.clone() { let trackers = self.torrents[&info_hash].trackers().clone(); for tracker_tier in trackers { for tracker in tracker_tier { self.tracker_manager .announce(tracker, &self.torrents[&info_hash]); } } } } }