udp: make handle_scrape_request take protocol-specific TorrentMap

This commit is contained in:
Joakim Frostegård 2022-04-16 02:21:19 +02:00
parent 1025391e4f
commit b8a74f0724

View file

@ -1,6 +1,5 @@
mod storage; mod storage;
use std::collections::BTreeMap;
use std::net::IpAddr; use std::net::IpAddr;
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
use std::time::Duration; use std::time::Duration;
@ -67,8 +66,11 @@ pub fn run_request_worker(
ConnectedResponse::AnnounceIpv6(response) ConnectedResponse::AnnounceIpv6(response)
} }
(ConnectedRequest::Scrape(request), ip) => { (ConnectedRequest::Scrape(request), IpAddr::V4(_)) => {
ConnectedResponse::Scrape(handle_scrape_request(&mut torrents, ip, request)) ConnectedResponse::Scrape(handle_scrape_request(&mut torrents.ipv4, request))
}
(ConnectedRequest::Scrape(request), IpAddr::V6(_)) => {
ConnectedResponse::Scrape(handle_scrape_request(&mut torrents.ipv6, request))
} }
}; };
@ -181,18 +183,17 @@ fn calc_max_num_peers_to_take(config: &Config, peers_wanted: i32) -> usize {
} }
} }
fn handle_scrape_request( fn handle_scrape_request<I: Ip>(
torrents: &mut TorrentMaps, torrents: &mut TorrentMap<I>,
src: IpAddr,
request: PendingScrapeRequest, request: PendingScrapeRequest,
) -> PendingScrapeResponse { ) -> PendingScrapeResponse {
const EMPTY_STATS: TorrentScrapeStatistics = create_torrent_scrape_statistics(0, 0); const EMPTY_STATS: TorrentScrapeStatistics = create_torrent_scrape_statistics(0, 0);
let mut torrent_stats: BTreeMap<usize, TorrentScrapeStatistics> = BTreeMap::new(); let torrent_stats = request
.info_hashes
if src.is_ipv4() { .into_iter()
torrent_stats.extend(request.info_hashes.into_iter().map(|(i, info_hash)| { .map(|(i, info_hash)| {
let s = if let Some(torrent_data) = torrents.ipv4.0.get(&info_hash) { let s = if let Some(torrent_data) = torrents.0.get(&info_hash) {
create_torrent_scrape_statistics( create_torrent_scrape_statistics(
torrent_data.num_seeders as i32, torrent_data.num_seeders as i32,
torrent_data.num_leechers as i32, torrent_data.num_leechers as i32,
@ -202,21 +203,8 @@ fn handle_scrape_request(
}; };
(i, s) (i, s)
})); })
} else { .collect();
torrent_stats.extend(request.info_hashes.into_iter().map(|(i, info_hash)| {
let s = if let Some(torrent_data) = torrents.ipv6.0.get(&info_hash) {
create_torrent_scrape_statistics(
torrent_data.num_seeders as i32,
torrent_data.num_leechers as i32,
)
} else {
EMPTY_STATS
};
(i, s)
}));
}
PendingScrapeResponse { PendingScrapeResponse {
slab_key: request.slab_key, slab_key: request.slab_key,