refactor: extract duplicate code from announce request handlers

Author: josecelano

src/http/warp_implementation/handlers.rs

@@ -47,13 +47,9 @@ pub async fn handle_announce(
 
     authenticate(&info_hash, &auth_key, tracker.clone()).await?;
 
-    let peer_ip = tracker.assign_ip_address_to_peer(&remote_client_ip);
+    let mut peer = peer_builder::from_request(&announce_request, &remote_client_ip);
 
-    let peer = peer_builder::from_request(&announce_request, &peer_ip);
-
-    let torrent_stats = tracker.update_torrent_with_peer_and_get_stats(&info_hash, &peer).await;
-
-    let peers = tracker.get_other_peers(&info_hash, &peer.peer_addr).await;
+    let response = tracker.announce(&info_hash, &mut peer, &remote_client_ip).await;
 
     match remote_client_ip {
         IpAddr::V4(_) => {
@@ -66,8 +62,8 @@ pub async fn handle_announce(
 
     send_announce_response(
         &announce_request,
-        &torrent_stats,
-        &peers,
+        &response.swam_stats,
+        &response.peers,
         tracker.config.announce_interval,
         tracker.config.min_announce_interval,
     )

src/tracker/mod.rs

@@ -17,6 +17,8 @@ use tokio::sync::mpsc::error::SendError;
 use tokio::sync::{RwLock, RwLockReadGuard};
 
 use self::error::Error;
+use self::peer::Peer;
+use self::torrent::SwamStats;
 use crate::config::Configuration;
 use crate::databases::driver::Driver;
 use crate::databases::{self, Database};
@@ -41,6 +43,11 @@ pub struct TorrentsMetrics {
     pub torrents: u64,
 }
 
+pub struct AnnounceResponse {
+    pub peers: Vec<Peer>,
+    pub swam_stats: SwamStats,
+}
+
 impl Tracker {
     /// # Errors
     ///
@@ -76,7 +83,18 @@ impl Tracker {
         self.mode == mode::Mode::Listed || self.mode == mode::Mode::PrivateListed
     }
 
-    /// It assigns a socket address to the peer
+    /// It handles an announce request
+    pub async fn announce(&self, info_hash: &InfoHash, peer: &mut Peer, remote_client_ip: &IpAddr) -> AnnounceResponse {
+        peer.change_ip(&self.assign_ip_address_to_peer(remote_client_ip));
+
+        let swam_stats = self.update_torrent_with_peer_and_get_stats(info_hash, peer).await;
+
+        let peers = self.get_other_peers(info_hash, &peer.peer_addr).await;
+
+        AnnounceResponse { peers, swam_stats }
+    }
+
+    /// It assigns an IP address to the peer
     #[must_use]
     pub fn assign_ip_address_to_peer(&self, remote_client_ip: &IpAddr) -> IpAddr {
         assign_ip_address_to_peer(remote_client_ip, self.config.get_ext_ip())

src/tracker/peer.rs

@@ -1,4 +1,4 @@
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 use std::panic::Location;
 
 use aquatic_udp_protocol::{AnnounceEvent, NumberOfBytes};
@@ -31,6 +31,10 @@ impl Peer {
     pub fn is_seeder(&self) -> bool {
         self.left.0 <= 0 && self.event != AnnounceEvent::Stopped
     }
+
+    pub fn change_ip(&mut self, new_ip: &IpAddr) {
+        self.peer_addr = SocketAddr::new(*new_ip, self.peer_addr.port());
+    }
 }
 
 #[derive(PartialEq, Eq, Hash, Clone, Debug, PartialOrd, Ord, Copy)]

src/udp/handlers.rs

@@ -118,13 +118,9 @@ pub async fn handle_announce(
 
     authenticate(&info_hash, tracker.clone()).await?;
 
-    let peer_ip = tracker.assign_ip_address_to_peer(&remote_client_ip);
+    let mut peer = peer_builder::from_request(&wrapped_announce_request, &remote_client_ip);
 
-    let peer = peer_builder::from_request(&wrapped_announce_request, &peer_ip);
-
-    let torrent_stats = tracker.update_torrent_with_peer_and_get_stats(&info_hash, &peer).await;
-
-    let peers = tracker.get_other_peers(&info_hash, &peer.peer_addr).await;
+    let response = tracker.announce(&info_hash, &mut peer, &remote_client_ip).await;
 
     match remote_client_ip {
         IpAddr::V4(_) => {
@@ -140,9 +136,10 @@ pub async fn handle_announce(
         Response::from(AnnounceResponse {
             transaction_id: wrapped_announce_request.announce_request.transaction_id,
             announce_interval: AnnounceInterval(i64::from(tracker.config.announce_interval) as i32),
-            leechers: NumberOfPeers(i64::from(torrent_stats.leechers) as i32),
-            seeders: NumberOfPeers(i64::from(torrent_stats.seeders) as i32),
-            peers: peers
+            leechers: NumberOfPeers(i64::from(response.swam_stats.leechers) as i32),
+            seeders: NumberOfPeers(i64::from(response.swam_stats.seeders) as i32),
+            peers: response
+                .peers
                 .iter()
                 .filter_map(|peer| {
                     if let IpAddr::V4(ip) = peer.peer_addr.ip() {
@@ -160,9 +157,10 @@ pub async fn handle_announce(
         Response::from(AnnounceResponse {
             transaction_id: wrapped_announce_request.announce_request.transaction_id,
             announce_interval: AnnounceInterval(i64::from(tracker.config.announce_interval) as i32),
-            leechers: NumberOfPeers(i64::from(torrent_stats.leechers) as i32),
-            seeders: NumberOfPeers(i64::from(torrent_stats.seeders) as i32),
-            peers: peers
+            leechers: NumberOfPeers(i64::from(response.swam_stats.leechers) as i32),
+            seeders: NumberOfPeers(i64::from(response.swam_stats.seeders) as i32),
+            peers: response
+                .peers
                 .iter()
                 .filter_map(|peer| {
                     if let IpAddr::V6(ip) = peer.peer_addr.ip() {

src/udp/mod.rs

@@ -1,9 +1,9 @@
 pub mod connection_cookie;
 pub mod error;
 pub mod handlers;
+pub mod peer_builder;
 pub mod request;
 pub mod server;
-pub mod peer_builder;
 
 pub type Bytes = u64;
 pub type Port = u16;