diff --git a/src/tracker.rs b/src/tracker.rs index daa5e00..c207899 100644 --- a/src/tracker.rs +++ b/src/tracker.rs @@ -45,7 +45,7 @@ pub const GARBAGE_COLLECTION_INTERVAL: Duration = Duration::from_secs(20); pub const CONNECTION_EXPIRE_TIME: Duration = Duration::from_secs(90); pub const DEFAULT_ANNOUNCE_INTERVAL: Duration = Duration::from_secs(60); -pub const DEFAULT_ANNOUNCE_WANT: u32 = 80; +pub const DEFAULT_ANNOUNCE_WANT: usize = 80; type ConnectionIdMap = HashMap; type InfoHashMap = HashMap>; @@ -122,7 +122,7 @@ impl Tracker { } fn handle_request(&mut self, request: &RequestMessage) -> Option { - return match &request.request { + match &request.request { UdpRequest::Connect(connect) => { let new_id: i64 = self.rng.random(); @@ -152,9 +152,9 @@ impl Tracker { // Ensure we honour the desired number of peers, within our boundaries - let n_announce_want: u32 = if let Some(n) = announce.num_want { - if n < DEFAULT_ANNOUNCE_WANT { - n + let n_announce_want: usize = if let Some(n) = announce.num_want { + if (n as usize) < DEFAULT_ANNOUNCE_WANT { + n as usize } else { DEFAULT_ANNOUNCE_WANT } @@ -162,14 +162,11 @@ impl Tracker { DEFAULT_ANNOUNCE_WANT }; - let mut n_announce_entries: u32 = 0; - let mut n_seeders: u32 = 0; - let mut v4_peers: Vec = Vec::new(); - let mut v6_peers: Vec = Vec::new(); - + let mut n_swarm_peers: usize = 0; + let mut n_seeders: usize = 0; let info_hashes = &mut self.info_hashes; - match info_hashes.get_mut(&announce.info_hash) { + let swarm_addrs = match info_hashes.get_mut(&announce.info_hash) { None => { // Info hash isn't currently tracked // No relevant peers in the swarm @@ -185,8 +182,13 @@ impl Tracker { remaining: announce.left as u64, }], ); + + None } Some(swarm) => { + n_swarm_peers = swarm.len(); + n_seeders = Self::count_seeders(swarm); + // Insert into swarm if not already present // TODO: sort (?) @@ -222,60 +224,45 @@ impl Tracker { } }; - if announce.event == Event::Stopped { - if let Some(idx) = existing_swarm_idx { - swarm.remove(idx); - } + if let Some(idx) = existing_swarm_idx + && announce.event == Event::Stopped + { + swarm.remove(idx); + return None; } // Iterate over all peers in the swarm for announce response - for peer_status in swarm { - // Respect number of peers requested + let swarm_addrs: Vec = swarm + .iter() + .filter_map(|status| { + let peer_invalid: bool = (status.last_event == Event::Stopped) + || (status.socket_addr == request.src_addr); - if n_announce_entries >= n_announce_want { - break; - } - - // Don't provide useless peers - // (peers who are no longer seeding or are from the source address) - - let peer_invalid: bool = (peer_status.last_event == Event::Stopped) - || (peer_status.socket_addr == request.src_addr); - - if peer_invalid { - log::trace!( - "(src: {}) {} with status \"{:?}\" deemed invalid", - request.src_addr, - peer_status.socket_addr, - peer_status.last_event - ); - continue; - } - - // Collect and add the relevant peer's address to the appropriate vector - - let is_seed: bool = peer_status.last_event == Event::Completed - || peer_status.remaining == 0; - - match &peer_status.socket_addr { - SocketAddr::V4(v4) => { - if request.src_addr.is_ipv4() { - v4_peers.push(*v4); - n_seeders += is_seed as u32; + if !peer_invalid { + if !status.socket_addr.is_ipv4() == request.src_addr.is_ipv4() { + None + } else { + Some(status.socket_addr) } - } - SocketAddr::V6(v6) => { - if request.src_addr.is_ipv6() { - v6_peers.push(*v6); - n_seeders += is_seed as u32; - } - } - }; + } else { + log::trace!( + "(src: {}) {} with status \"{:?}\" deemed invalid", + request.src_addr, + status.socket_addr, + status.last_event + ); - // Keep track of the announce entries gathered + None + } + }) + .take(n_announce_want as usize) + .collect(); - n_announce_entries += 1; + if !swarm_addrs.is_empty() { + Some(swarm_addrs) + } else { + None } } }; @@ -283,20 +270,40 @@ impl Tracker { let resp_common = AnnounceResponseCommon { transaction_id: announce.transaction_id, interval: DEFAULT_ANNOUNCE_INTERVAL.as_secs() as i32, - leechers: (n_announce_entries - n_seeders) as i32, + leechers: (n_swarm_peers - n_seeders) as i32, seeders: n_seeders as i32, }; let response = if request.src_addr.is_ipv4() { // IPv4 AnnounceResponse::V4(AnnounceResponseV4 { common: resp_common, - peers: v4_peers, + peers: swarm_addrs.map_or_else( + || Vec::new(), + |v| { + v.iter() + .filter_map(|a| match a { + SocketAddr::V4(v4) => Some(*v4), + _ => None, + }) + .collect() + }, + ), }) } else { // IPv6 AnnounceResponse::V6(AnnounceResponseV6 { common: resp_common, - peers: v6_peers, + peers: swarm_addrs.map_or_else( + || Vec::new(), + |v| { + v.iter() + .filter_map(|a| match a { + SocketAddr::V6(v6) => Some(*v6), + _ => None, + }) + .collect() + }, + ), }) }; @@ -318,12 +325,7 @@ impl Tracker { .iter() .filter_map(|info_hash| { if let Some(swarm) = info_hashes.get(info_hash) { - let n_seeders = swarm - .iter() - .filter(|status| { - status.last_event == Event::Completed || status.remaining == 0 - }) - .count(); + let n_seeders = Self::count_seeders(swarm); let n_leechers = swarm.len() - n_seeders; Some(ScrapeStats { @@ -345,7 +347,7 @@ impl Tracker { dst_addr: request.src_addr, }) } - }; + } } fn garbage_collect(&mut self) { @@ -405,4 +407,11 @@ impl Tracker { return n_purged; } + + fn count_seeders(swarm: &[PeerStatus]) -> usize { + swarm + .iter() + .filter(|status| status.last_event == Event::Completed || status.remaining == 0) + .count() + } }