diff --git a/rust/connlib/snownet/tests/lib.rs b/rust/connlib/snownet/tests/lib.rs index 22d64e2c4..b5c204dfa 100644 --- a/rust/connlib/snownet/tests/lib.rs +++ b/rust/connlib/snownet/tests/lib.rs @@ -351,7 +351,7 @@ impl Firewall { impl TestRelay { fn new(local: impl Into, span: Span) -> Self { let local = local.into(); - let inner = firezone_relay::Server::new(to_ip_stack(local), OsRng, 3478, 49152, 65535); + let inner = firezone_relay::Server::new(to_ip_stack(local), OsRng, 3478, 49152..=65535); Self { inner, diff --git a/rust/connlib/tunnel/src/tests/sut.rs b/rust/connlib/tunnel/src/tests/sut.rs index 95b262c68..6169218d0 100644 --- a/rust/connlib/tunnel/src/tests/sut.rs +++ b/rust/connlib/tunnel/src/tests/sut.rs @@ -91,8 +91,7 @@ impl StateMachineTest for TunnelTest { IpStack::from((ip4, ip6)), rand::rngs::StdRng::seed_from_u64(seed), 3478, - 49152, - 65535, + 49152..=65535, )) }, debug_span!("relay", rid = %id), diff --git a/rust/relay/src/main.rs b/rust/relay/src/main.rs index 7c7545f71..2fe23209c 100644 --- a/rust/relay/src/main.rs +++ b/rust/relay/src/main.rs @@ -114,8 +114,7 @@ async fn main() -> Result<()> { public_addr, make_rng(args.rng_seed), args.listen_port, - args.lowest_port, - args.highest_port, + args.lowest_port..=args.highest_port, ); let last_heartbeat_sent = Arc::new(Mutex::new(Option::::None)); diff --git a/rust/relay/src/server.rs b/rust/relay/src/server.rs index 9d7f03a7f..48d77a77e 100644 --- a/rust/relay/src/server.rs +++ b/rust/relay/src/server.rs @@ -20,6 +20,7 @@ use secrecy::SecretString; use std::collections::{HashMap, VecDeque}; use std::hash::Hash; use std::net::{IpAddr, SocketAddr}; +use std::ops::RangeInclusive; use std::time::{Duration, Instant, SystemTime}; use stun_codec::rfc5389::attributes::{ ErrorCode, MessageIntegrity, Nonce, Realm, Username, XorMappedAddress, @@ -62,8 +63,8 @@ pub struct Server { HashMap<(AllocationPort, PeerSocket), (ClientSocket, ChannelNumber)>, listen_port: u16, - lowest_port: u16, - highest_port: u16, + + ports: RangeInclusive, /// Channel numbers are unique by client, thus indexed by both. channels_by_client_and_number: HashMap<(ClientSocket, ChannelNumber), Channel>, @@ -157,8 +158,7 @@ where public_address: impl Into, mut rng: R, listen_port: u16, - lowest_port: u16, - highest_port: u16, + ports: RangeInclusive, ) -> Self { // TODO: Validate that local IP isn't multicast / loopback etc. @@ -185,8 +185,7 @@ where allocations: Default::default(), clients_by_allocation: Default::default(), listen_port, - lowest_port, - highest_port, + ports, channels_by_client_and_number: Default::default(), channel_numbers_by_client_and_peer: Default::default(), pending_commands: Default::default(), @@ -807,7 +806,7 @@ where ); let port = loop { - let candidate = AllocationPort(self.rng.gen_range(self.lowest_port..self.highest_port)); + let candidate = AllocationPort(self.rng.gen_range(self.ports.clone())); if !self.clients_by_allocation.contains_key(&candidate) { break candidate; @@ -823,7 +822,7 @@ where } fn max_available_ports(&self) -> u16 { - self.highest_port - self.lowest_port + self.ports.clone().count() as u16 } fn create_channel_binding( diff --git a/rust/relay/tests/regression.rs b/rust/relay/tests/regression.rs index 8ef8d8dbd..e2d5d44a8 100644 --- a/rust/relay/tests/regression.rs +++ b/rust/relay/tests/regression.rs @@ -696,7 +696,7 @@ struct TestServer { impl TestServer { fn new(relay_public_addr: impl Into) -> Self { Self { - server: Server::new(relay_public_addr, StepRng::new(0, 0), 3478, 49152, 65535), + server: Server::new(relay_public_addr, StepRng::new(0, 0), 3478, 49152..=65535), } }