From 405053edeebab6343741081117e7ceaedd5b3340 Mon Sep 17 00:00:00 2001 From: Jose Celano Date: Tue, 18 Oct 2022 16:46:46 +0100 Subject: [PATCH] fix: tests after enabling connection id check --- src/udp/handlers.rs | 72 +++++++++++++++++++++++++++++++++------------ 1 file changed, 54 insertions(+), 18 deletions(-) diff --git a/src/udp/handlers.rs b/src/udp/handlers.rs index 801923542..d46cd9231 100644 --- a/src/udp/handlers.rs +++ b/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) -> 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;