refactor: [#1495] use SocketAddr as key for peers in Swarm

josecelano · josecelano · commit 2882705fbab8 · 2025-05-05T20:57:27.000+01:00
This change prevents duplicate peers with the same address but different IDs, ensuring more accurate peer tracking.
diff --git a/packages/axum-http-tracker-server/tests/server/asserts.rs b/packages/axum-http-tracker-server/tests/server/asserts.rs
@@ -22,6 +22,7 @@ pub fn assert_bencoded_error(response_text: &String, expected_failure_reason: &s
     );
 }
 
+#[allow(dead_code)]
 pub async fn assert_empty_announce_response(response: Response) {
     assert_eq!(response.status(), 200);
     let announce_response: Announce = serde_bencode::from_str(&response.text().await.unwrap()).unwrap();
diff --git a/packages/axum-http-tracker-server/tests/server/requests/announce.rs b/packages/axum-http-tracker-server/tests/server/requests/announce.rs
@@ -126,6 +126,11 @@ impl QueryBuilder {
         self
     }
 
+    pub fn with_port(mut self, port: u16) -> Self {
+        self.announce_query.port = port;
+        self
+    }
+
     pub fn without_compact(mut self) -> Self {
         self.announce_query.compact = None;
         self
diff --git a/packages/axum-http-tracker-server/tests/server/v1/contract.rs b/packages/axum-http-tracker-server/tests/server/v1/contract.rs
@@ -105,8 +105,8 @@ mod for_all_config_modes {
         use crate::common::fixtures::invalid_info_hashes;
         use crate::server::asserts::{
             assert_announce_response, assert_bad_announce_request_error_response, assert_cannot_parse_query_param_error_response,
-            assert_cannot_parse_query_params_error_response, assert_compact_announce_response, assert_empty_announce_response,
-            assert_is_announce_response, assert_missing_query_params_for_announce_request_error_response,
+            assert_cannot_parse_query_params_error_response, assert_compact_announce_response, assert_is_announce_response,
+            assert_missing_query_params_for_announce_request_error_response,
         };
         use crate::server::client::Client;
         use crate::server::requests::announce::{Compact, QueryBuilder};
@@ -559,27 +559,53 @@ mod for_all_config_modes {
         }
 
         #[tokio::test]
-        async fn should_consider_two_peers_to_be_the_same_when_they_have_the_same_peer_id_even_if_the_ip_is_different() {
+        async fn should_consider_two_peers_to_be_the_same_when_they_have_the_same_socket_address_even_if_the_peer_id_is_different(
+        ) {
             logging::setup();
 
             let env = Started::new(&configuration::ephemeral_public().into()).await;
 
             let info_hash = InfoHash::from_str("9c38422213e30bff212b30c360d26f9a02136422").unwrap(); // DevSkim: ignore DS173237
             let peer = PeerBuilder::default().build();
 
-            // Add a peer
-            env.add_torrent_peer(&info_hash, &peer);
-
-            let announce_query = QueryBuilder::default()
+            let announce_query_1 = QueryBuilder::default()
                 .with_info_hash(&info_hash)
                 .with_peer_id(&peer.peer_id)
+                .with_peer_addr(&peer.peer_addr.ip())
+                .with_port(peer.peer_addr.port())
+                .query();
+
+            let announce_query_2 = QueryBuilder::default()
+                .with_info_hash(&info_hash)
+                .with_peer_id(&PeerId(*b"-qB00000000000000002")) // Different peer ID
+                .with_peer_addr(&peer.peer_addr.ip())
+                .with_port(peer.peer_addr.port())
                 .query();
 
-            assert_ne!(peer.peer_addr.ip(), announce_query.peer_addr);
+            // Same peer socket address
+            assert_eq!(announce_query_1.peer_addr, announce_query_2.peer_addr);
+            assert_eq!(announce_query_1.port, announce_query_2.port);
+
+            // Different peer ID
+            assert_ne!(announce_query_1.peer_id, announce_query_2.peer_id);
 
-            let response = Client::new(*env.bind_address()).announce(&announce_query).await;
+            let _response = Client::new(*env.bind_address()).announce(&announce_query_1).await;
+            let response = Client::new(*env.bind_address()).announce(&announce_query_2).await;
 
-            assert_empty_announce_response(response).await;
+            let announce_policy = env.container.tracker_core_container.core_config.announce_policy;
+
+            // The response should contain only the first peer.
+            assert_announce_response(
+                response,
+                &Announce {
+                    complete: 1,
+                    incomplete: 0,
+                    interval: announce_policy.interval,
+                    min_interval: announce_policy.interval_min,
+                    peers: vec![],
+                },
+            )
+            .await;
 
             env.stop().await;
         }
diff --git a/packages/torrent-repository/src/entry/swarm.rs b/packages/torrent-repository/src/entry/swarm.rs
@@ -3,18 +3,12 @@ use std::collections::BTreeMap;
 use std::net::SocketAddr;
 use std::sync::Arc;
 
-use aquatic_udp_protocol::PeerId;
 use torrust_tracker_primitives::peer::{self, Peer};
 use torrust_tracker_primitives::DurationSinceUnixEpoch;
 
-// code-review: the current implementation uses the peer Id as the ``BTreeMap``
-// key. That would allow adding two identical peers except for the Id.
-// For example, two peers with the same socket address but a different peer Id
-// would be allowed. That would lead to duplicated peers in the tracker responses.
-
 #[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Swarm {
-    peers: BTreeMap<PeerId, Arc<Peer>>,
+    peers: BTreeMap<SocketAddr, Arc<Peer>>,
 }
 
 impl Swarm {
@@ -28,12 +22,12 @@ impl Swarm {
         self.peers.is_empty()
     }
 
-    pub fn upsert(&mut self, value: Arc<peer::Peer>) -> Option<Arc<peer::Peer>> {
-        self.peers.insert(value.peer_id, value)
+    pub fn upsert(&mut self, peer: Arc<Peer>) -> Option<Arc<Peer>> {
+        self.peers.insert(peer.peer_addr, peer)
     }
 
-    pub fn remove(&mut self, key: &PeerId) -> Option<Arc<peer::Peer>> {
-        self.peers.remove(key)
+    pub fn remove(&mut self, peer: &Peer) -> Option<Arc<Peer>> {
+        self.peers.remove(&peer.peer_addr)
     }
 
     pub fn remove_inactive_peers(&mut self, current_cutoff: DurationSinceUnixEpoch) {
@@ -42,12 +36,12 @@ impl Swarm {
     }
 
     #[must_use]
-    pub fn get(&self, peer_id: &PeerId) -> Option<&Arc<peer::Peer>> {
-        self.peers.get(peer_id)
+    pub fn get(&self, peer_addr: &SocketAddr) -> Option<&Arc<Peer>> {
+        self.peers.get(peer_addr)
     }
 
     #[must_use]
-    pub fn get_all(&self, limit: Option<usize>) -> Vec<Arc<peer::Peer>> {
+    pub fn get_all(&self, limit: Option<usize>) -> Vec<Arc<Peer>> {
         match limit {
             Some(limit) => self.peers.values().take(limit).cloned().collect(),
             None => self.peers.values().cloned().collect(),
@@ -151,7 +145,7 @@ mod tests {
 
             swarm.upsert(peer.into());
 
-            assert_eq!(swarm.get(&peer.peer_id), Some(Arc::new(peer)).as_ref());
+            assert_eq!(swarm.get(&peer.peer_addr), Some(Arc::new(peer)).as_ref());
         }
 
         #[test]
@@ -173,7 +167,7 @@ mod tests {
 
             swarm.upsert(peer.into());
 
-            swarm.remove(&peer.peer_id);
+            swarm.remove(&peer);
 
             assert!(swarm.is_empty());
         }
@@ -186,9 +180,9 @@ mod tests {
 
             swarm.upsert(peer.into());
 
-            swarm.remove(&peer.peer_id);
+            swarm.remove(&peer);
 
-            assert_eq!(swarm.get(&peer.peer_id), None);
+            assert_eq!(swarm.get(&peer.peer_addr), None);
         }
 
         #[test]
@@ -273,16 +267,42 @@ mod tests {
         }
 
         #[test]
-        fn allow_inserting_two_identical_peers_except_for_the_id() {
+        fn allow_inserting_two_identical_peers_except_for_the_socket_address() {
             let mut swarm = Swarm::default();
 
-            let peer1 = PeerBuilder::default().with_peer_id(&PeerId(*b"-qB00000000000000001")).build();
+            let peer1 = PeerBuilder::default()
+                .with_peer_addr(&SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 6969))
+                .build();
             swarm.upsert(peer1.into());
 
-            let peer2 = PeerBuilder::default().with_peer_id(&PeerId(*b"-qB00000000000000002")).build();
+            let peer2 = PeerBuilder::default()
+                .with_peer_addr(&SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 6969))
+                .build();
             swarm.upsert(peer2.into());
 
             assert_eq!(swarm.len(), 2);
         }
+
+        #[test]
+        fn not_allow_inserting_two_peers_with_different_peer_id_but_the_same_socket_address() {
+            let mut swarm = Swarm::default();
+
+            // When that happens the peer ID will be changed in the swarm.
+            // In practice, it's like if the peer had changed its ID.
+
+            let peer1 = PeerBuilder::default()
+                .with_peer_id(&PeerId(*b"-qB00000000000000001"))
+                .with_peer_addr(&SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 6969))
+                .build();
+            swarm.upsert(peer1.into());
+
+            let peer2 = PeerBuilder::default()
+                .with_peer_id(&PeerId(*b"-qB00000000000000002"))
+                .with_peer_addr(&SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 6969))
+                .build();
+            swarm.upsert(peer2.into());
+
+            assert_eq!(swarm.len(), 1);
+        }
     }
 }
diff --git a/packages/torrent-repository/src/entry/torrent.rs b/packages/torrent-repository/src/entry/torrent.rs
@@ -75,7 +75,7 @@ impl TrackedTorrent {
 
         match peer::ReadInfo::get_event(peer) {
             AnnounceEvent::Stopped => {
-                drop(self.swarm.remove(&peer::ReadInfo::get_id(peer)));
+                drop(self.swarm.remove(peer));
             }
             AnnounceEvent::Completed => {
                 let previous = self.swarm.upsert(Arc::new(*peer));
diff --git a/packages/torrent-repository/tests/common/torrent_peer_builder.rs b/packages/torrent-repository/tests/common/torrent_peer_builder.rs
@@ -1,4 +1,4 @@
-use std::net::SocketAddr;
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
 use aquatic_udp_protocol::{AnnounceEvent, NumberOfBytes, PeerId};
 use torrust_tracker_clock::clock::Time;
@@ -67,24 +67,40 @@ impl TorrentPeerBuilder {
 
 /// A torrent seeder is a peer with 0 bytes left to download which
 /// has not announced it has stopped
+#[allow(clippy::cast_sign_loss)]
+#[allow(clippy::cast_possible_truncation)]
 #[must_use]
 pub fn a_completed_peer(id: i32) -> peer::Peer {
     let peer_id = peer::Id::new(id);
+    let peer_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), id as u16);
+
     TorrentPeerBuilder::new()
         .with_number_of_bytes_left(0)
         .with_event_completed()
         .with_peer_id(*peer_id)
+        .with_peer_address(peer_addr)
         .into()
 }
 
 /// A torrent leecher is a peer that is not a seeder.
 /// Leecher: left > 0 OR event = Stopped
+///
+/// # Panics
+///
+/// This function panics if proved id can't be converted into a valid socket address port.
+///
+/// The `id` argument is used to identify the peer in both the `peer_id` and the `peer_addr`.
+#[allow(clippy::cast_sign_loss)]
+#[allow(clippy::cast_possible_truncation)]
 #[must_use]
 pub fn a_started_peer(id: i32) -> peer::Peer {
     let peer_id = peer::Id::new(id);
+    let peer_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), id as u16);
+
     TorrentPeerBuilder::new()
         .with_number_of_bytes_left(1)
         .with_event_started()
         .with_peer_id(*peer_id)
+        .with_peer_address(peer_addr)
         .into()
 }
diff --git a/packages/torrent-repository/tests/entry/mod.rs b/packages/torrent-repository/tests/entry/mod.rs
@@ -370,8 +370,7 @@ async fn it_should_limit_the_number_of_peers_returned(
 
     // We add one more peer than the scrape limit
     for peer_number in 1..=74 + 1 {
-        let mut peer = a_started_peer(1);
-        peer.peer_id = *peer::Id::new(peer_number);
+        let peer = a_started_peer(peer_number);
         torrent.upsert_peer(&peer);
     }
 
diff --git a/packages/tracker-core/src/lib.rs b/packages/tracker-core/src/lib.rs
@@ -224,14 +224,14 @@ mod tests {
                     // Scrape
                     let scrape_data = scrape_handler.scrape(&vec![info_hash]).await.unwrap();
 
-                    // The expected swarm metadata for the file
+                    // The expected swarm metadata for the torrent
                     let mut expected_scrape_data = ScrapeData::empty();
                     expected_scrape_data.add_file(
                         &info_hash,
                         SwarmMetadata {
-                            complete: 0, // the "complete" peer does not count because it was not previously known
-                            downloaded: 0,
-                            incomplete: 1, // the "incomplete" peer we have just announced
+                            complete: 1,   // the "incomplete" announced
+                            downloaded: 0, // the "complete" peer download does not count because it was not previously known
+                            incomplete: 1, // the "incomplete" peer announced
                         },
                     );
 
diff --git a/packages/tracker-core/src/test_helpers.rs b/packages/tracker-core/src/test_helpers.rs
@@ -104,7 +104,7 @@ pub(crate) mod tests {
     #[must_use]
     pub fn complete_peer() -> Peer {
         Peer {
-            peer_id: PeerId(*b"-qB00000000000000000"),
+            peer_id: PeerId(*b"-qB00000000000000001"),
             peer_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(126, 0, 0, 1)), 8080),
             updated: DurationSinceUnixEpoch::new(1_669_397_478_934, 0),
             uploaded: NumberOfBytes::new(0),
@@ -118,8 +118,8 @@ pub(crate) mod tests {
     #[must_use]
     pub fn incomplete_peer() -> Peer {
         Peer {
-            peer_id: PeerId(*b"-qB00000000000000000"),
-            peer_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(126, 0, 0, 1)), 8080),
+            peer_id: PeerId(*b"-qB00000000000000002"),
+            peer_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(126, 0, 0, 2)), 8080),
             updated: DurationSinceUnixEpoch::new(1_669_397_478_934, 0),
             uploaded: NumberOfBytes::new(0),
             downloaded: NumberOfBytes::new(0),

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,7 @@ pub fn assert_bencoded_error(response_text: &String, expected_failure_reason: &s`
`22`	`22`	`);`
`23`	`23`	`}`
`24`	`24`
	`25`	`+#[allow(dead_code)]`
`25`	`26`	`pub async fn assert_empty_announce_response(response: Response) {`
`26`	`27`	`assert_eq!(response.status(), 200);`
`27`	`28`	`let announce_response: Announce = serde_bencode::from_str(&response.text().await.unwrap()).unwrap();`
Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ impl TrackedTorrent {`
`75`	`75`
`76`	`76`	`match peer::ReadInfo::get_event(peer) {`
`77`	`77`	`AnnounceEvent::Stopped => {`
`78`		`- drop(self.swarm.remove(&peer::ReadInfo::get_id(peer)));`
	`78`	`+ drop(self.swarm.remove(peer));`
`79`	`79`	`}`
`80`	`80`	`AnnounceEvent::Completed => {`
`81`	`81`	`let previous = self.swarm.upsert(Arc::new(*peer));`
Original file line number	Diff line number	Diff line change
`@@ -370,8 +370,7 @@ async fn it_should_limit_the_number_of_peers_returned(`
`370`	`370`
`371`	`371`	`// We add one more peer than the scrape limit`
`372`	`372`	`for peer_number in 1..=74 + 1 {`
`373`		`- let mut peer = a_started_peer(1);`
`374`		`- peer.peer_id = *peer::Id::new(peer_number);`
	`373`	`+ let peer = a_started_peer(peer_number);`
`375`	`374`	`torrent.upsert_peer(&peer);`
`376`	`375`	`}`
`377`	`376`