Skip to content

Commit

Permalink
dev: use the connection cookie implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
@da2ce7
da2ce7 committed Oct 18, 2022
1 parent 43685b6 commit 4949424
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 41 deletions.
14 changes: 1 addition & 13 deletions src/protocol/utils.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,4 @@
use std::net::SocketAddr;

use aquatic_udp_protocol::ConnectionId;

use super::clock::{DefaultClock, DurationSinceUnixEpoch, Time};

pub fn get_connection_id(remote_address: &SocketAddr) -> ConnectionId {
ConnectionId(((current_time() / 3600) | ((remote_address.port() as u64) << 36)) as i64)
}

pub fn current_time() -> u64 {
DefaultClock::now().as_secs()
}
use super::clock::DurationSinceUnixEpoch;

pub fn ser_unix_time_value<S: serde::Serializer>(unix_time_value: &DurationSinceUnixEpoch, ser: S) -> Result<S::Ok, S::Error> {
ser.serialize_u64(unix_time_value.as_millis() as u64)
Expand Down
4 changes: 2 additions & 2 deletions src/tracker/peer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ mod test {
};

use crate::peer::TorrentPeer;
use crate::protocol::utils::get_connection_id;
use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
// todo: duplicate functions is PR 82. Remove duplication once both PR are merged.

fn sample_ipv4_remote_addr() -> SocketAddr {
Expand All @@ -152,7 +152,7 @@ mod test {
let info_hash_aquatic = aquatic_udp_protocol::InfoHash([0u8; 20]);

let default_request = AnnounceRequest {
connection_id: get_connection_id(&sample_ipv4_remote_addr()),
connection_id: into_connection_id(&make_connection_cookie(&sample_ipv4_remote_addr())),
transaction_id: TransactionId(0i32),
info_hash: info_hash_aquatic,
peer_id: AquaticPeerId(*b"-qB00000000000000000"),
Expand Down
60 changes: 34 additions & 26 deletions src/udp/handlers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@ use aquatic_udp_protocol::{
NumberOfPeers, Port, Request, Response, ResponsePeer, ScrapeRequest, ScrapeResponse, TorrentScrapeStatistics, TransactionId,
};

use super::connection_cookie::{check_connection_cookie, from_connection_id, into_connection_id, make_connection_cookie};
use crate::peer::TorrentPeer;
use crate::protocol::utils::get_connection_id;
use crate::tracker::statistics::TrackerStatisticsEvent;
use crate::tracker::torrent::TorrentError;
use crate::tracker::tracker::TorrentTracker;
Expand Down Expand Up @@ -69,7 +69,8 @@ pub async fn handle_connect(
request: &ConnectRequest,
tracker: Arc<TorrentTracker>,
) -> Result<Response, ServerError> {
let connection_id = get_connection_id(&remote_addr);
let connection_cookie = make_connection_cookie(&remote_addr);
let connection_id = into_connection_id(&connection_cookie);

let response = Response::from(ConnectResponse {
transaction_id: request.transaction_id,
Expand All @@ -94,6 +95,13 @@ pub async fn handle_announce(
announce_request: &AnnounceRequest,
tracker: Arc<TorrentTracker>,
) -> Result<Response, ServerError> {
match check_connection_cookie(&remote_addr, &from_connection_id(&announce_request.connection_id)) {
Ok(_) => {}
Err(e) => {
return Err(e);
}
}

let wrapped_announce_request = AnnounceRequestWrapper::new(announce_request.clone());

authenticate(&wrapped_announce_request.info_hash, tracker.clone()).await?;
Expand Down Expand Up @@ -401,9 +409,9 @@ mod tests {
use aquatic_udp_protocol::{ConnectRequest, ConnectResponse, Response, TransactionId};

use super::{default_tracker_config, sample_ipv4_socket_address, sample_ipv6_remote_addr, TrackerStatsServiceMock};
use crate::protocol::utils::get_connection_id;
use crate::statistics::TrackerStatisticsEvent;
use crate::tracker::tracker::TorrentTracker;
use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
use crate::udp::handle_connect;
use crate::udp::handlers::tests::{initialized_public_tracker, sample_ipv4_remote_addr};

Expand All @@ -426,7 +434,7 @@ mod tests {
assert_eq!(
response,
Response::Connect(ConnectResponse {
connection_id: get_connection_id(&sample_ipv4_remote_addr()),
connection_id: into_connection_id(&make_connection_cookie(&sample_ipv4_remote_addr())),
transaction_id: request.transaction_id
})
);
Expand All @@ -445,7 +453,7 @@ mod tests {
assert_eq!(
response,
Response::Connect(ConnectResponse {
connection_id: get_connection_id(&sample_ipv4_remote_addr()),
connection_id: into_connection_id(&make_connection_cookie(&sample_ipv4_remote_addr())),
transaction_id: request.transaction_id
})
);
Expand Down Expand Up @@ -486,7 +494,7 @@ mod tests {
TransactionId,
};

use crate::protocol::utils::get_connection_id;
use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
use crate::udp::handlers::tests::sample_ipv4_remote_addr;

struct AnnounceRequestBuilder {
Expand All @@ -500,7 +508,7 @@ mod tests {
let info_hash_aquatic = aquatic_udp_protocol::InfoHash([0u8; 20]);

let default_request = AnnounceRequest {
connection_id: get_connection_id(&sample_ipv4_remote_addr()),
connection_id: into_connection_id(&make_connection_cookie(&sample_ipv4_remote_addr())),
transaction_id: TransactionId(0i32),
info_hash: info_hash_aquatic,
peer_id: AquaticPeerId([255u8; 20]),
Expand Down Expand Up @@ -558,9 +566,9 @@ mod tests {
Response, ResponsePeer,
};

use crate::protocol::utils::get_connection_id;
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::{
Expand All @@ -581,7 +589,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V4(client_ip), client_port);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.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)
Expand All @@ -605,7 +613,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(126, 0, 0, 1)), 8080);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
.into();

let response = handle_announce(remote_addr, &request, initialized_public_tracker())
Expand Down Expand Up @@ -644,7 +652,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V4(remote_client_ip), remote_client_port);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.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)
Expand Down Expand Up @@ -679,7 +687,7 @@ mod tests {
async fn announce_a_new_peer_using_ipv4(tracker: Arc<TorrentTracker>) -> Response {
let remote_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(126, 0, 0, 1)), 8080);
let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
.into();
let response = handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
response
Expand Down Expand Up @@ -723,7 +731,7 @@ mod tests {

use aquatic_udp_protocol::{InfoHash as AquaticInfoHash, PeerId as AquaticPeerId};

use crate::protocol::utils::get_connection_id;
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};
Expand All @@ -741,7 +749,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V4(client_ip), client_port);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.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)
Expand Down Expand Up @@ -775,9 +783,9 @@ mod tests {
Response, ResponsePeer,
};

use crate::protocol::utils::get_connection_id;
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::{
Expand All @@ -799,7 +807,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.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)
Expand All @@ -826,7 +834,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), 8080);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
.into();

let response = handle_announce(remote_addr, &request, initialized_public_tracker())
Expand Down Expand Up @@ -865,7 +873,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V6(remote_client_ip), remote_client_port);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.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)
Expand Down Expand Up @@ -903,7 +911,7 @@ mod tests {
let client_port = 8080;
let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);
let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
.into();
let response = handle_announce(remote_addr, &request, tracker.clone()).await.unwrap();
response
Expand Down Expand Up @@ -937,7 +945,7 @@ mod tests {
let remote_addr = sample_ipv6_remote_addr();

let announce_request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.with_connection_id(into_connection_id(&make_connection_cookie(&remote_addr)))
.into();

handle_announce(remote_addr, &announce_request, tracker.clone())
Expand All @@ -951,9 +959,9 @@ mod tests {

use aquatic_udp_protocol::{InfoHash as AquaticInfoHash, PeerId as AquaticPeerId};

use crate::protocol::utils::get_connection_id;
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;
Expand All @@ -977,7 +985,7 @@ mod tests {
let remote_addr = SocketAddr::new(IpAddr::V6(client_ip_v6), client_port);

let request = AnnounceRequestBuilder::default()
.with_connection_id(get_connection_id(&remote_addr))
.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)
Expand Down Expand Up @@ -1012,8 +1020,8 @@ mod tests {
};

use super::TorrentPeerBuilder;
use crate::protocol::utils::get_connection_id;
use crate::tracker::tracker::TorrentTracker;
use crate::udp::connection_cookie::{into_connection_id, make_connection_cookie};
use crate::udp::handle_scrape;
use crate::udp::handlers::tests::{initialized_public_tracker, sample_ipv4_remote_addr};
use crate::PeerId;
Expand All @@ -1034,7 +1042,7 @@ mod tests {
let info_hashes = vec![info_hash];

let request = ScrapeRequest {
connection_id: get_connection_id(&remote_addr),
connection_id: into_connection_id(&make_connection_cookie(&remote_addr)),
transaction_id: TransactionId(0i32),
info_hashes,
};
Expand Down Expand Up @@ -1072,7 +1080,7 @@ mod tests {
let info_hashes = vec![*info_hash];

ScrapeRequest {
connection_id: get_connection_id(&remote_addr),
connection_id: into_connection_id(&make_connection_cookie(&remote_addr)),
transaction_id: TransactionId(0i32),
info_hashes,
}
Expand Down Expand Up @@ -1217,7 +1225,7 @@ mod tests {
let info_hashes = vec![info_hash];

ScrapeRequest {
connection_id: get_connection_id(&remote_addr),
connection_id: into_connection_id(&make_connection_cookie(&remote_addr)),
transaction_id: TransactionId(0i32),
info_hashes,
}
Expand Down

0 comments on commit 4949424

Please sign in to comment.