udp scrape improvements (#43)

* udp_protocol: forbid full scrapes

* udp: improve PendingScrapeResponseMap logging

* udp: PendingScrapeResponseMap: store less data, improve logging

* udp: PendingScrapeResponseMap: log if replacing entry on insert

* udp: PendingScrapeResponseMap: use remote addr in key

* Run cargo fmt

* README: update copyright end year

* udp: move scrape request splitting logic into PendingScrapeResponseMap

* udp: add quickcheck test test_pending_scrape_response_map

* udp protocol: fix failing test_scrape_request_convert_identity

README.md

@@ -174,7 +174,7 @@ To fairly compare HTTP performance to opentracker, set keepalive to false in
 
 ## Copyright and license
 
-Copyright (c) 2020-2021 Joakim Frostegård
+Copyright (c) 2020-2022 Joakim Frostegård
 
 Distributed under Apache 2.0 license (details in `LICENSE` file.)
 

aquatic_udp/src/workers/socket.rs

@@ -48,58 +48,101 @@ impl ConnectionMap {
     }
 }
 
-pub struct PendingScrapeResponseMeta {
+#[derive(Debug)]
+pub struct PendingScrapeResponseMapEntry {
     num_pending: usize,
     valid_until: ValidUntil,
+    torrent_stats: BTreeMap<usize, TorrentScrapeStatistics>,
 }
 
 #[derive(Default)]
 pub struct PendingScrapeResponseMap(
-    AHashIndexMap<(ConnectionId, TransactionId), (PendingScrapeResponseMeta, PendingScrapeResponse)>,
+    AHashIndexMap<(ConnectionId, TransactionId, SocketAddr), PendingScrapeResponseMapEntry>,
 );
 
 impl PendingScrapeResponseMap {
-    pub fn prepare(
+    pub fn prepare_split_requests(
         &mut self,
-        connection_id: ConnectionId,
+        config: &Config,
-        transaction_id: TransactionId,
+        request: ScrapeRequest,
-        num_pending: usize,
+        addr: SocketAddr,
         valid_until: ValidUntil,
-    ) {
+    ) -> impl IntoIterator<Item = (RequestWorkerIndex, PendingScrapeRequest)> {
-        let meta = PendingScrapeResponseMeta {
+        let mut split_requests: AHashIndexMap<RequestWorkerIndex, PendingScrapeRequest> =
-            num_pending,
+            Default::default();
+
+        if request.info_hashes.is_empty() {
+            ::log::warn!(
+                "Attempted to prepare PendingScrapeResponseMap entry with zero info hashes"
+            );
+
+            return split_requests;
+        }
+
+        let connection_id = request.connection_id;
+        let transaction_id = request.transaction_id;
+
+        for (i, info_hash) in request.info_hashes.into_iter().enumerate() {
+            let split_request = split_requests
+                .entry(RequestWorkerIndex::from_info_hash(&config, info_hash))
+                .or_insert_with(|| PendingScrapeRequest {
+                    connection_id,
+                    transaction_id,
+                    info_hashes: BTreeMap::new(),
+                });
+
+            split_request.info_hashes.insert(i, info_hash);
+        }
+
+        let key = (connection_id, transaction_id, addr);
+
+        let entry = PendingScrapeResponseMapEntry {
+            num_pending: split_requests.len(),
             valid_until,
-        };
+            torrent_stats: Default::default(),
-        let response = PendingScrapeResponse {
-            connection_id,
-            transaction_id,
-            torrent_stats: BTreeMap::new(),
         };
 
-        self.0.insert((connection_id, transaction_id), (meta, response));
+        if let Some(previous_entry) = self.0.insert(key, entry) {
+            ::log::warn!(
+                "PendingScrapeResponseMap.prepare replaced previous entry {:?} for key {:?}",
+                previous_entry,
+                key
+            );
+        }
+
+        split_requests
     }
 
-    pub fn add_and_get_finished(&mut self, response: PendingScrapeResponse) -> Option<Response> {
+    pub fn add_and_get_finished(
-        let key = (response.connection_id, response.transaction_id);
+        &mut self,
+        response: PendingScrapeResponse,
+        addr: SocketAddr,
+    ) -> Option<Response> {
+        let key = (response.connection_id, response.transaction_id, addr);
 
-        let finished = if let Some(r) = self.0.get_mut(&key) {
+        let finished = if let Some(entry) = self.0.get_mut(&key) {
-            r.0.num_pending -= 1;
+            entry.num_pending -= 1;
 
-            r.1.torrent_stats.extend(response.torrent_stats.into_iter());
+            entry
+                .torrent_stats
+                .extend(response.torrent_stats.into_iter());
 
-            r.0.num_pending == 0
+            entry.num_pending == 0
         } else {
-            ::log::warn!("PendingScrapeResponses.add didn't find PendingScrapeResponse in map");
+            ::log::warn!(
+                "PendingScrapeResponseMap.add didn't find entry for key {:?}",
+                key
+            );
 
             false
         };
 
         if finished {
-            let response = self.0.remove(&key).unwrap().1;
+            let entry = self.0.remove(&key).unwrap();
 
             Some(Response::Scrape(ScrapeResponse {
                 transaction_id: response.transaction_id,
-                torrent_stats: response.torrent_stats.into_values().collect(),
+                torrent_stats: entry.torrent_stats.into_values().collect(),
             }))
         } else {
             None
@@ -109,7 +152,19 @@ impl PendingScrapeResponseMap {
     pub fn clean(&mut self) {
         let now = Instant::now();
 
-        self.0.retain(|_, v| v.0.valid_until.0 > now);
+        self.0.retain(|k, v| {
+            let keep = v.valid_until.0 > now;
+
+            if !keep {
+                ::log::warn!(
+                    "Removing PendingScrapeResponseMap entry while cleaning. {:?}: {:?}",
+                    k,
+                    v
+                );
+            }
+
+            keep
+        });
         self.0.shrink_to_fit();
     }
 }
@@ -371,32 +426,14 @@ pub fn handle_request(
         }
         Ok(Request::Scrape(request)) => {
             if connections.contains(request.connection_id, src) {
-                let mut requests: AHashIndexMap<RequestWorkerIndex, PendingScrapeRequest> =
+                let split_requests = pending_scrape_responses.prepare_split_requests(
-                    Default::default();
+                    config,
-
+                    request,
-                let connection_id = request.connection_id;
+                    src,
-                let transaction_id = request.transaction_id;
-
-                for (i, info_hash) in request.info_hashes.into_iter().enumerate() {
-                    let pending = requests
-                        .entry(RequestWorkerIndex::from_info_hash(&config, info_hash))
-                        .or_insert_with(|| PendingScrapeRequest {
-                            connection_id,
-                            transaction_id,
-                            info_hashes: BTreeMap::new(),
-                        });
-
-                    pending.info_hashes.insert(i, info_hash);
-                }
-
-                pending_scrape_responses.prepare(
-                    connection_id,
-                    transaction_id,
-                    requests.len(),
                     pending_scrape_valid_until,
                 );
 
-                for (request_worker_index, request) in requests {
+                for (request_worker_index, request) in split_requests {
                     request_sender.try_send_to(
                         request_worker_index,
                         ConnectedRequest::Scrape(request),
@@ -458,7 +495,7 @@ fn send_responses(
 
     for (response, addr) in response_receiver.try_iter() {
         let opt_response = match response {
-            ConnectedResponse::Scrape(r) => pending_scrape_responses.add_and_get_finished(r),
+            ConnectedResponse::Scrape(r) => pending_scrape_responses.add_and_get_finished(r, addr),
             ConnectedResponse::AnnounceIpv4(r) => Some(Response::AnnounceIpv4(r)),
             ConnectedResponse::AnnounceIpv6(r) => Some(Response::AnnounceIpv6(r)),
         };
@@ -591,3 +628,111 @@ pub fn create_socket(config: &Config) -> ::std::net::UdpSocket {
 
     socket.into()
 }
+
+#[cfg(test)]
+mod tests {
+    use quickcheck::TestResult;
+    use quickcheck_macros::quickcheck;
+
+    use super::*;
+
+    #[quickcheck]
+    fn test_pending_scrape_response_map(
+        request_data: Vec<(i32, i64, SocketAddr, u8)>,
+        request_workers: u8,
+    ) -> TestResult {
+        if request_workers == 0 {
+            return TestResult::discard();
+        }
+
+        let mut config = Config::default();
+
+        config.request_workers = request_workers as usize;
+
+        let valid_until = ValidUntil::new(1);
+
+        let mut map = PendingScrapeResponseMap::default();
+
+        let mut requests = Vec::new();
+
+        for (t, c, a, b) in request_data {
+            if b == 0 {
+                return TestResult::discard();
+            }
+
+            let mut info_hashes = Vec::new();
+
+            for i in 0..b {
+                let info_hash = InfoHash([i; 20]);
+
+                info_hashes.push(info_hash);
+            }
+
+            let request = ScrapeRequest {
+                transaction_id: TransactionId(t),
+                connection_id: ConnectionId(c),
+                info_hashes,
+            };
+
+            requests.push((request, a));
+        }
+
+        let mut all_split_requests = Vec::new();
+
+        for (request, addr) in requests.iter() {
+            let split_requests = map.prepare_split_requests(
+                &config,
+                request.to_owned(),
+                addr.to_owned(),
+                valid_until,
+            );
+
+            all_split_requests.push((
+                addr,
+                split_requests
+                    .into_iter()
+                    .collect::<Vec<(RequestWorkerIndex, PendingScrapeRequest)>>(),
+            ));
+        }
+
+        assert_eq!(map.0.len(), requests.len());
+
+        let mut responses = Vec::new();
+
+        for (addr, split_requests) in all_split_requests {
+            for (worker_index, split_request) in split_requests {
+                assert!(worker_index.0 < request_workers as usize);
+
+                let torrent_stats = split_request
+                    .info_hashes
+                    .into_iter()
+                    .map(|(i, info_hash)| {
+                        (
+                            i,
+                            TorrentScrapeStatistics {
+                                seeders: NumberOfPeers((info_hash.0[0]) as i32),
+                                leechers: NumberOfPeers(0),
+                                completed: NumberOfDownloads(0),
+                            },
+                        )
+                    })
+                    .collect();
+
+                let response = PendingScrapeResponse {
+                    transaction_id: split_request.transaction_id,
+                    connection_id: split_request.connection_id,
+                    torrent_stats,
+                };
+
+                if let Some(response) = map.add_and_get_finished(response, addr.to_owned()) {
+                    responses.push(response);
+                }
+            }
+        }
+
+        assert!(map.0.is_empty());
+        assert_eq!(responses.len(), requests.len());
+
+        TestResult::from_bool(true)
+    }
+}

aquatic_udp_protocol/src/request.rs

@@ -271,18 +271,26 @@ impl Request {
                 let position = cursor.position() as usize;
                 let inner = cursor.into_inner();
 
-                let info_hashes = (&inner[position..])
+                let info_hashes: Vec<InfoHash> = (&inner[position..])
                     .chunks_exact(20)
                     .take(max_scrape_torrents as usize)
                     .map(|chunk| InfoHash(chunk.try_into().unwrap()))
                     .collect();
 
-                Ok((ScrapeRequest {
+                if info_hashes.is_empty() {
-                    connection_id: ConnectionId(connection_id),
+                    Err(RequestParseError::sendable_text(
-                    transaction_id: TransactionId(transaction_id),
+                        "Full scrapes are not allowed",
-                    info_hashes,
+                        connection_id,
-                })
+                        transaction_id,
-                .into())
+                    ))
+                } else {
+                    Ok((ScrapeRequest {
+                        connection_id: ConnectionId(connection_id),
+                        transaction_id: TransactionId(transaction_id),
+                        info_hashes,
+                    })
+                    .into())
+                }
             }
 
             _ => Err(RequestParseError::sendable_text(
@@ -296,6 +304,7 @@ impl Request {
 
 #[cfg(test)]
 mod tests {
+    use quickcheck::TestResult;
     use quickcheck_macros::quickcheck;
 
     use super::*;
@@ -378,7 +387,11 @@ mod tests {
     }
 
     #[quickcheck]
-    fn test_scrape_request_convert_identity(request: ScrapeRequest) -> bool {
+    fn test_scrape_request_convert_identity(request: ScrapeRequest) -> TestResult {
-        same_after_conversion(request.into())
+        if request.info_hashes.is_empty() {
+            return TestResult::discard();
+        }
+
+        TestResult::from_bool(same_after_conversion(request.into()))
     }
 }