fix: tests after enabling connection id check

Author: josecelano

src/udp/handlers.rs

@@ -490,7 +490,8 @@ mod tests {
         use std::net::Ipv4Addr;
 
         use aquatic_udp_protocol::{
-            AnnounceEvent, AnnounceRequest, NumberOfBytes, NumberOfPeers, PeerId as AquaticPeerId, PeerKey, Port, TransactionId,
+            AnnounceEvent, AnnounceRequest, ConnectionId, NumberOfBytes, NumberOfPeers, PeerId as AquaticPeerId, PeerKey, Port,
+            TransactionId,
         };
 
         use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
@@ -525,6 +526,11 @@ mod tests {
                 }
             }
 
+            pub fn with_connection_id(mut self, connection_id: ConnectionId) -> Self {
+                self.request.connection_id = connection_id;
+                self
+            }
+
             pub fn with_info_hash(mut self, info_hash: aquatic_udp_protocol::InfoHash) -> Self {
                 self.request.info_hash = info_hash;
                 self
@@ -562,6 +568,7 @@ mod tests {
 
             use crate::statistics::TrackerStatisticsEvent;
             use crate::tracker::tracker::TorrentTracker;
+            use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
             use crate::udp::handle_announce;
             use crate::udp::handlers::tests::announce_request::AnnounceRequestBuilder;
             use crate::udp::handlers::tests::{
@@ -579,14 +586,16 @@ mod tests {
                 let info_hash = AquaticInfoHash([0u8; 20]);
                 let peer_id = AquaticPeerId([255u8; 20]);
 
+                let remote_addr = SocketAddr::new(IpAddr::V4(client_ip), client_port);
+
                 let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
                     .with_info_hash(info_hash)
                     .with_peer_id(peer_id)
                     .with_ip_address(client_ip)
                     .with_port(client_port)
                     .into();
 
-                let remote_addr = SocketAddr::new(IpAddr::V4(client_ip), client_port);
                 handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
 
                 let peers = tracker.get_all_torrent_peers(&info_hash.0.into()).await;
@@ -601,9 +610,12 @@ mod tests {
 
             #[tokio::test]
             async fn the_announced_peer_should_not_be_included_in_the_response() {
-                let request = AnnounceRequestBuilder::default().into();
                 let remote_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(126, 0, 0, 1)), 8080);
 
+                let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
+                    .into();
+
                 let response = handle_announce(remote_addr, &request, initialized_public_tracker())
                     .await
                     .unwrap();
@@ -637,14 +649,16 @@ mod tests {
                 let remote_client_port = 8081;
                 let peer_address = Ipv4Addr::new(126, 0, 0, 2);
 
+                let remote_addr = SocketAddr::new(IpAddr::V4(remote_client_ip), remote_client_port);
+
                 let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
                     .with_info_hash(info_hash)
                     .with_peer_id(peer_id)
                     .with_ip_address(peer_address)
                     .with_port(client_port)
                     .into();
 
-                let remote_addr = SocketAddr::new(IpAddr::V4(remote_client_ip), remote_client_port);
                 handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
 
                 let peers = tracker.get_all_torrent_peers(&info_hash.0.into()).await;
@@ -671,8 +685,10 @@ mod tests {
             }
 
             async fn announce_a_new_peer_using_ipv4(tracker: Arc<TorrentTracker>) -> Response {
-                let request = AnnounceRequestBuilder::default().into();
                 let remote_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(126, 0, 0, 1)), 8080);
+                let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
+                    .into();
                 let response = handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
                 response
             }
@@ -715,6 +731,7 @@ mod tests {
 
                 use aquatic_udp_protocol::{InfoHash as AquaticInfoHash, PeerId as AquaticPeerId};
 
+                use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
                 use crate::udp::handle_announce;
                 use crate::udp::handlers::tests::announce_request::AnnounceRequestBuilder;
                 use crate::udp::handlers::tests::{initialized_public_tracker, TorrentPeerBuilder};
@@ -729,14 +746,16 @@ mod tests {
                     let info_hash = AquaticInfoHash([0u8; 20]);
                     let peer_id = AquaticPeerId([255u8; 20]);
 
+                    let remote_addr = SocketAddr::new(IpAddr::V4(client_ip), client_port);
+
                     let request = AnnounceRequestBuilder::default()
+                        .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
                         .with_info_hash(info_hash)
                         .with_peer_id(peer_id)
                         .with_ip_address(client_ip)
                         .with_port(client_port)
                         .into();
 
-                    let remote_addr = SocketAddr::new(IpAddr::V4(client_ip), client_port);
                     handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
 
                     let peers = tracker.get_all_torrent_peers(&info_hash.0.into()).await;
@@ -766,6 +785,7 @@ mod tests {
 
             use crate::statistics::TrackerStatisticsEvent;
             use crate::tracker::tracker::TorrentTracker;
+            use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
             use crate::udp::handle_announce;
             use crate::udp::handlers::tests::announce_request::AnnounceRequestBuilder;
             use crate::udp::handlers::tests::{
@@ -784,14 +804,16 @@ mod tests {
                 let info_hash = AquaticInfoHash([0u8; 20]);
                 let peer_id = AquaticPeerId([255u8; 20]);
 
+                let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);
+
                 let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
                     .with_info_hash(info_hash)
                     .with_peer_id(peer_id)
                     .with_ip_address(client_ip_v4)
                     .with_port(client_port)
                     .into();
 
-                let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);
                 handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
 
                 let peers = tracker.get_all_torrent_peers(&info_hash.0.into()).await;
@@ -806,11 +828,15 @@ mod tests {
 
             #[tokio::test]
             async fn the_announced_peer_should_not_be_included_in_the_response() {
-                let request = AnnounceRequestBuilder::default().into();
                 let client_ip_v4 = Ipv4Addr::new(126, 0, 0, 1);
                 let client_ip_v6 = client_ip_v4.to_ipv6_compatible();
+
                 let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), 8080);
 
+                let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
+                    .into();
+
                 let response = handle_announce(remote_addr, &request, initialized_public_tracker())
                     .await
                     .unwrap();
@@ -844,14 +870,16 @@ mod tests {
                 let remote_client_port = 8081;
                 let peer_address = "126.0.0.1".parse().unwrap();
 
+                let remote_addr = SocketAddr::new(IpAddr::V6(remote_client_ip), remote_client_port);
+
                 let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
                     .with_info_hash(info_hash)
                     .with_peer_id(peer_id)
                     .with_ip_address(peer_address)
                     .with_port(client_port)
                     .into();
 
-                let remote_addr = SocketAddr::new(IpAddr::V6(remote_client_ip), remote_client_port);
                 handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
 
                 let peers = tracker.get_all_torrent_peers(&info_hash.0.into()).await;
@@ -882,7 +910,9 @@ mod tests {
                 let client_ip_v6 = client_ip_v4.to_ipv6_compatible();
                 let client_port = 8080;
                 let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);
-                let request = AnnounceRequestBuilder::default().into();
+                let request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
+                    .into();
                 let response = handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
                 response
             }
@@ -911,13 +941,16 @@ mod tests {
                 tracker_stats_service.should_throw_event(TrackerStatisticsEvent::Udp6Announce);
 
                 let tracker = Arc::new(TorrentTracker::new(default_tracker_config(), tracker_stats_service).unwrap());
-                handle_announce(
-                    sample_ipv6_remote_addr(),
-                    &AnnounceRequestBuilder::default().into(),
-                    tracker.clone(),
-                )
-                .await
-                .unwrap();
+
+                let remote_addr = sample_ipv6_remote_addr();
+
+                let announce_request = AnnounceRequestBuilder::default()
+                    .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
+                    .into();
+
+                handle_announce(remote_addr, &announce_request, tracker.clone())
+                    .await
+                    .unwrap();
             }
 
             mod from_a_loopback_ip {
@@ -928,6 +961,7 @@ mod tests {
 
                 use crate::statistics::StatsTracker;
                 use crate::tracker::tracker::TorrentTracker;
+                use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
                 use crate::udp::handle_announce;
                 use crate::udp::handlers::tests::announce_request::AnnounceRequestBuilder;
                 use crate::udp::handlers::tests::TrackerConfigurationBuilder;
@@ -948,14 +982,16 @@ mod tests {
                     let info_hash = AquaticInfoHash([0u8; 20]);
                     let peer_id = AquaticPeerId([255u8; 20]);
 
+                    let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);
+
                     let request = AnnounceRequestBuilder::default()
+                        .with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
                         .with_info_hash(info_hash)
                         .with_peer_id(peer_id)
                         .with_ip_address(client_ip_v4)
                         .with_port(client_port)
                         .into();
 
-                    let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);
                     handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
 
                     let peers = tracker.get_all_torrent_peers(&info_hash.0.into()).await;