diff --git a/crates/udp/src/common.rs b/crates/udp/src/common.rs index 5536e72..4dd3496 100644 --- a/crates/udp/src/common.rs +++ b/crates/udp/src/common.rs @@ -244,7 +244,7 @@ impl Default for State { #[cfg(test)] mod tests { - use std::net::Ipv6Addr; + use std::{net::Ipv6Addr, num::NonZeroU16}; use crate::config::Config; @@ -260,7 +260,7 @@ mod tests { let peers = ::std::iter::repeat(ResponsePeer { ip_address: Ipv6AddrBytes(Ipv6Addr::new(1, 1, 1, 1, 1, 1, 1, 1).octets()), - port: Port::new(1), + port: Port::new(NonZeroU16::new(1).unwrap()), }) .take(config.protocol.max_response_peers) .collect(); diff --git a/crates/udp/tests/access_list.rs b/crates/udp/tests/access_list.rs index cd0ec79..bc3dacf 100644 --- a/crates/udp/tests/access_list.rs +++ b/crates/udp/tests/access_list.rs @@ -6,6 +6,7 @@ use std::{ fs::File, io::Write, net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket}, + num::NonZeroU16, time::Duration, }; @@ -78,7 +79,7 @@ fn test_access_list( &socket, tracker_addr, connection_id, - 1, + NonZeroU16::new(1).unwrap(), info_hash_fail, 10, false, @@ -95,7 +96,7 @@ fn test_access_list( &socket, tracker_addr, connection_id, - 1, + NonZeroU16::new(1).unwrap(), info_hash_success, 10, false, diff --git a/crates/udp/tests/common/mod.rs b/crates/udp/tests/common/mod.rs index 0656b73..7fbec29 100644 --- a/crates/udp/tests/common/mod.rs +++ b/crates/udp/tests/common/mod.rs @@ -3,6 +3,7 @@ use std::{ io::Cursor, net::{SocketAddr, UdpSocket}, + num::NonZeroU16, time::Duration, }; @@ -42,7 +43,7 @@ pub fn announce( socket: &UdpSocket, tracker_addr: SocketAddr, connection_id: ConnectionId, - peer_port: u16, + peer_port: NonZeroU16, info_hash: InfoHash, peers_wanted: usize, seeder: bool, @@ -50,7 +51,7 @@ pub fn announce( let mut peer_id = PeerId([0; 20]); for chunk in peer_id.0.chunks_exact_mut(2) { - chunk.copy_from_slice(&peer_port.to_ne_bytes()); + chunk.copy_from_slice(&peer_port.get().to_ne_bytes()); } let request = Request::Announce(AnnounceRequest { diff --git a/crates/udp/tests/invalid_connection_id.rs b/crates/udp/tests/invalid_connection_id.rs index 7506854..96ecf67 100644 --- a/crates/udp/tests/invalid_connection_id.rs +++ b/crates/udp/tests/invalid_connection_id.rs @@ -5,6 +5,7 @@ use common::*; use std::{ io::{Cursor, ErrorKind}, net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket}, + num::NonZeroU16, time::Duration, }; @@ -51,7 +52,7 @@ fn test_invalid_connection_id() -> anyhow::Result<()> { ip_address: Ipv4AddrBytes([0; 4]), key: PeerKey::new(0), peers_wanted: NumberOfPeers::new(10), - port: Port::new(1), + port: Port::new(NonZeroU16::new(1).unwrap()), }); let scrape_request = Request::Scrape(ScrapeRequest { diff --git a/crates/udp/tests/requests_responses.rs b/crates/udp/tests/requests_responses.rs index 0ce345d..dc55aa0 100644 --- a/crates/udp/tests/requests_responses.rs +++ b/crates/udp/tests/requests_responses.rs @@ -5,6 +5,7 @@ use common::*; use std::{ collections::{hash_map::RandomState, HashSet}, net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket}, + num::NonZeroU16, time::Duration, }; @@ -45,7 +46,7 @@ fn test_multiple_connect_announce_scrape() -> anyhow::Result<()> { &socket, tracker_addr, connection_id, - PEER_PORT_START + i as u16, + NonZeroU16::new(PEER_PORT_START + i as u16).unwrap(), info_hash, PEERS_WANTED, is_seeder, diff --git a/crates/udp_load_test/src/worker/mod.rs b/crates/udp_load_test/src/worker/mod.rs index 40b0f3c..8c7c606 100644 --- a/crates/udp_load_test/src/worker/mod.rs +++ b/crates/udp_load_test/src/worker/mod.rs @@ -1,6 +1,5 @@ use std::io::Cursor; use std::net::SocketAddr; -use std::num::NonZeroU16; use std::sync::atomic::Ordering; use std::time::Duration; @@ -193,7 +192,7 @@ impl Worker { scrape_hash_indices, connection_id, peer_id: generate_peer_id(), - port: Port::new(self.rng.gen::().into()), + port: Port::new(self.rng.gen()), } } diff --git a/crates/udp_protocol/src/common.rs b/crates/udp_protocol/src/common.rs index 6c54cb2..51fbfaf 100644 --- a/crates/udp_protocol/src/common.rs +++ b/crates/udp_protocol/src/common.rs @@ -1,5 +1,6 @@ use std::fmt::Debug; use std::net::{Ipv4Addr, Ipv6Addr}; +use std::num::NonZeroU16; pub use aquatic_peer_id::{PeerClient, PeerId}; use zerocopy::network_endian::{I32, I64, U16, U32}; @@ -76,8 +77,8 @@ impl NumberOfDownloads { pub struct Port(pub U16); impl Port { - pub fn new(v: u16) -> Self { - Self(U16::new(v)) + pub fn new(v: NonZeroU16) -> Self { + Self(U16::new(v.into())) } } diff --git a/crates/udp_protocol/src/request.rs b/crates/udp_protocol/src/request.rs index cea19c4..1b8cb43 100644 --- a/crates/udp_protocol/src/request.rs +++ b/crates/udp_protocol/src/request.rs @@ -281,7 +281,7 @@ impl RequestParseError { mod tests { use quickcheck::TestResult; use quickcheck_macros::quickcheck; - use zerocopy::network_endian::{I32, I64, U16}; + use zerocopy::network_endian::{I32, I64}; use super::*; @@ -319,7 +319,7 @@ mod tests { ip_address: Ipv4AddrBytes::arbitrary(g), key: PeerKey::new(i32::arbitrary(g)), peers_wanted: NumberOfPeers(I32::new(i32::arbitrary(g))), - port: Port(U16::new(u16::arbitrary(g))), + port: Port::new(quickcheck::Arbitrary::arbitrary(g)), } } }