use crate::infohash::v1::U160; use crate::LookupFilter; use std::collections::{HashMap, HashSet}; use std::net::SocketAddr; use std::ops::Add; use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant}; use url::Url; #[derive(Debug, Default)] pub struct PeerStorage { torrents: HashMap, } #[derive(Debug, Default)] pub struct PeerCollection { sets: Vec, peers: HashSet, peers6: HashSet, } #[derive(Debug)] pub struct PeerSet { source: PeerStorageSource, expires: Instant, peers: Vec, } impl PeerCollection { fn add_peers(&mut self, peers: Vec, source: PeerStorageSource) { let set = PeerSet { source, peers, expires: Instant::now().add(Duration::from_secs(60 * 30)), }; for peer in &set.peers { if peer.is_ipv4() { self.peers.insert(*peer); } else { self.peers6.insert(*peer); } } self.sets.push(set); } } #[derive(Debug, Clone)] pub enum PeerStorageSource { DHT, Tracker(Url), } impl PeerStorage { pub fn new() -> PeerStorage { PeerStorage { torrents: Default::default(), } } pub fn add_peers( &mut self, info_hash: U160, peers: Vec, source: PeerStorageSource, ) { self.torrents .entry(info_hash) .or_default() .add_peers(peers, source) } pub fn get_peers(&self, info_hash: U160, filter: LookupFilter) -> Vec { let collection = if let Some(collection) = self.torrents.get(&info_hash) { collection } else { return vec![]; }; let mut buffer = vec![]; if filter != LookupFilter::IPv6 { buffer.extend(&collection.peers); } if filter != LookupFilter::IPv4 { buffer.extend(&collection.peers6); } buffer } } #[derive(Debug, Clone)] pub struct PeerStorageHolder { inner: Arc>, } impl PeerStorageHolder { pub fn new(peer_storage: PeerStorage) -> PeerStorageHolder { PeerStorageHolder { inner: Arc::new(RwLock::new(peer_storage)), } } pub fn get_peers(&self, info_hash: U160, filter: LookupFilter) -> Vec { self.inner.read().unwrap().get_peers(info_hash, filter) } pub fn add_peers( &mut self, info_hash: U160, peers: Vec, source: PeerStorageSource, ) { self.inner .write() .unwrap() .add_peers(info_hash, peers, source) } }