diff --git a/rust/connlib/clients/shared/src/eventloop.rs b/rust/connlib/clients/shared/src/eventloop.rs index 9f2a55f7d..159fc7796 100644 --- a/rust/connlib/clients/shared/src/eventloop.rs +++ b/rust/connlib/clients/shared/src/eventloop.rs @@ -13,7 +13,7 @@ use connlib_shared::{ use firezone_tunnel::{ClientTunnel, Tun}; use phoenix_channel::{ErrorReply, OutboundRequestId, PhoenixChannel}; use std::{ - collections::{HashMap, HashSet}, + collections::{BTreeSet, HashMap}, net::IpAddr, task::{Context, Poll}, }; @@ -226,7 +226,7 @@ where tracing::info!("Firezone Started!"); self.tunnel.set_resources(resources); - self.tunnel.update_relays(HashSet::default(), relays) + self.tunnel.update_relays(BTreeSet::default(), relays) } IngressMessages::ResourceCreatedOrUpdated(resource) => { self.tunnel.add_resource(resource); @@ -239,7 +239,7 @@ where connected, }) => self .tunnel - .update_relays(HashSet::from_iter(disconnected_ids), connected), + .update_relays(BTreeSet::from_iter(disconnected_ids), connected), IngressMessages::InvalidateIceCandidates(GatewayIceCandidates { gateway_id, candidates, diff --git a/rust/connlib/clients/shared/src/messages.rs b/rust/connlib/clients/shared/src/messages.rs index 6068c1487..139fe66c7 100644 --- a/rust/connlib/clients/shared/src/messages.rs +++ b/rust/connlib/clients/shared/src/messages.rs @@ -4,7 +4,10 @@ use connlib_shared::messages::{ ResourceId, ReuseConnection, }; use serde::{Deserialize, Serialize}; -use std::{collections::HashSet, net::IpAddr}; +use std::{ + collections::{BTreeSet, HashSet}, + net::IpAddr, +}; #[derive(Debug, PartialEq, Eq, Deserialize, Clone)] pub struct InitClient { @@ -61,7 +64,7 @@ pub struct GatewaysIceCandidates { /// The list of gateway IDs these candidates will be broadcast to. pub gateway_ids: Vec, /// Actual RTC ice candidates - pub candidates: HashSet, + pub candidates: BTreeSet, } #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)] @@ -117,7 +120,7 @@ mod test { "client", EgressMessages::BroadcastIceCandidates(GatewaysIceCandidates { gateway_ids: vec!["b3d34a15-55ab-40df-994b-a838e75d65d7".parse().unwrap()], - candidates: HashSet::from([ + candidates: BTreeSet::from([ "candidate:7031633958891736544 1 udp 50331391 35.244.108.190 53909 typ relay" .to_owned(), ]), diff --git a/rust/connlib/snownet/src/allocation.rs b/rust/connlib/snownet/src/allocation.rs index e231a5ba3..e948d8dd8 100644 --- a/rust/connlib/snownet/src/allocation.rs +++ b/rust/connlib/snownet/src/allocation.rs @@ -93,7 +93,7 @@ struct Credentials { } /// Describes the socket address(es) we know about the relay. -#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] pub enum RelaySocket { /// The relay is only reachable via IPv4. V4(SocketAddrV4), diff --git a/rust/connlib/snownet/src/node.rs b/rust/connlib/snownet/src/node.rs index 4b242751e..772e76f4b 100644 --- a/rust/connlib/snownet/src/node.rs +++ b/rust/connlib/snownet/src/node.rs @@ -14,16 +14,13 @@ use ip_packet::{ use rand::random; use secrecy::{ExposeSecret, Secret}; use std::borrow::Cow; +use std::collections::{BTreeMap, BTreeSet, HashSet}; use std::hash::Hash; use std::marker::PhantomData; use std::mem; use std::ops::ControlFlow; use std::time::{Duration, Instant}; -use std::{ - collections::{HashMap, HashSet, VecDeque}, - net::SocketAddr, - sync::Arc, -}; +use std::{collections::VecDeque, net::SocketAddr, sync::Arc}; use str0m::ice::{IceAgent, IceAgentEvent, IceCreds, StunMessage, StunPacket}; use str0m::net::Protocol; use str0m::{Candidate, CandidateKind, IceConnectionState}; @@ -89,7 +86,7 @@ pub struct Node { next_rate_limiter_reset: Option, - allocations: HashMap, + allocations: BTreeMap, connections: Connections, pending_events: VecDeque>, @@ -121,8 +118,8 @@ pub enum Error { impl Node where - TId: Eq + Hash + Copy + fmt::Display, - RId: Copy + Eq + Hash + PartialEq + fmt::Debug + fmt::Display, + TId: Eq + Hash + Copy + Ord + fmt::Display, + RId: Copy + Eq + Hash + PartialEq + Ord + fmt::Debug + fmt::Display, { pub fn new(private_key: StaticSecret) -> Self { let public_key = &(&private_key).into(); @@ -132,12 +129,12 @@ where marker: Default::default(), index: IndexLfsr::default(), rate_limiter: Arc::new(RateLimiter::new(public_key, HANDSHAKE_RATE_LIMIT)), - host_candidates: HashSet::default(), + host_candidates: Default::default(), buffered_transmits: VecDeque::default(), next_rate_limiter_reset: None, pending_events: VecDeque::default(), buffer: Box::new([0u8; MAX_UDP_SIZE]), - allocations: HashMap::default(), + allocations: Default::default(), connections: Default::default(), stats: Default::default(), } @@ -491,8 +488,8 @@ where pub fn update_relays( &mut self, - to_remove: HashSet, - to_add: &HashSet<(RId, RelaySocket, String, String, String)>, + to_remove: BTreeSet, + to_add: &BTreeSet<(RId, RelaySocket, String, String, String)>, now: Instant, ) { // First, invalidate all candidates from relays that we should stop using. @@ -576,7 +573,7 @@ where signalling_completed_at: now, remote_pub_key: remote, state: ConnectionState::Connecting { - possible_sockets: HashSet::default(), + possible_sockets: BTreeSet::default(), buffered: RingBuffer::new(10), }, last_outgoing: now, @@ -609,7 +606,7 @@ where /// Tries to handle the packet using one of our [`Allocation`]s. /// /// This function is in the hot-path of packet processing and thus must be as efficient as possible. - /// Even look-ups in [`HashMap`]s and linear searches across small lists are expensive at this point. + /// Even look-ups in [`BTreeMap`]s and linear searches across small lists are expensive at this point. /// Thus, we use the first byte of the message as a heuristic for whether we should attempt to handle it here. /// /// See for details on de-multiplexing. @@ -783,8 +780,8 @@ where impl Node where - TId: Eq + Hash + Copy + fmt::Display, - RId: Copy + Eq + Hash + PartialEq + fmt::Debug + fmt::Display, + TId: Eq + Hash + Copy + Ord + fmt::Display, + RId: Copy + Eq + Hash + PartialEq + Ord + fmt::Debug + fmt::Display, { /// Create a new connection indexed by the given ID. /// @@ -874,8 +871,8 @@ where impl Node where - TId: Eq + Hash + Copy + fmt::Display, - RId: Copy + Eq + Hash + PartialEq + fmt::Debug + fmt::Display, + TId: Eq + Hash + Copy + Ord + fmt::Display, + RId: Copy + Eq + Hash + PartialEq + Ord + fmt::Debug + fmt::Display, { /// Accept a new connection indexed by the given ID. /// @@ -959,8 +956,8 @@ where } struct Connections { - initial: HashMap, - established: HashMap>, + initial: BTreeMap, + established: BTreeMap>, } impl Default for Connections { @@ -974,8 +971,8 @@ impl Default for Connections { impl Connections where - TId: Eq + Hash + Copy + fmt::Display, - RId: Copy + Eq + Hash + PartialEq + fmt::Debug + fmt::Display, + TId: Eq + Hash + Copy + Ord + fmt::Display, + RId: Copy + Eq + Hash + PartialEq + Ord + fmt::Debug + fmt::Display, { fn gc(&mut self, events: &mut VecDeque>) { self.initial.retain(|id, conn| { @@ -1062,11 +1059,11 @@ fn encode_as_channel_data( relay: RId, dest: SocketAddr, contents: &[u8], - allocations: &mut HashMap, + allocations: &mut BTreeMap, now: Instant, ) -> Result, EncodeError> where - RId: Copy + Eq + Hash + PartialEq + fmt::Debug, + RId: Copy + Eq + Hash + PartialEq + Ord + fmt::Debug, { let allocation = allocations .get_mut(&relay) @@ -1310,7 +1307,7 @@ enum ConnectionState { /// We are still running ICE to figure out, which socket to use to send data. Connecting { /// Socket addresses from which we might receive data (even before we are connected). - possible_sockets: HashSet, + possible_sockets: BTreeSet, /// Packets emitted by wireguard whilst are still running ICE. /// /// This can happen if the remote's WG session initiation arrives at our socket before we nominate it. @@ -1322,7 +1319,7 @@ enum ConnectionState { /// Our nominated socket. peer_socket: PeerSocket, /// Other addresses that we might see traffic from (e.g. STUN messages during roaming). - possible_sockets: HashSet, + possible_sockets: BTreeSet, }, /// The connection failed in an unrecoverable way and will be GC'd. Failed, @@ -1361,7 +1358,7 @@ enum PeerSocket { impl Connection where - RId: PartialEq + Eq + Hash + fmt::Debug + Copy, + RId: PartialEq + Eq + Hash + fmt::Debug + Copy + Ord, { /// Checks if we want to accept a packet from a certain address. /// @@ -1428,11 +1425,11 @@ where &mut self, cid: TId, now: Instant, - allocations: &mut HashMap, + allocations: &mut BTreeMap, transmits: &mut VecDeque>, ) where - TId: fmt::Display + Copy, - RId: Copy + fmt::Display, + TId: Copy + Ord + fmt::Display, + RId: Copy + Ord + fmt::Display, { self.agent.handle_timeout(now); @@ -1621,7 +1618,7 @@ where &mut self, packet: &[u8], buffer: &'b mut [u8], - allocations: &mut HashMap, + allocations: &mut BTreeMap, transmits: &mut VecDeque>, now: Instant, ) -> ControlFlow, MutableIpPacket<'b>> { @@ -1705,7 +1702,7 @@ where fn force_handshake( &mut self, - allocations: &mut HashMap, + allocations: &mut BTreeMap, transmits: &mut VecDeque>, now: Instant, ) where @@ -1753,11 +1750,11 @@ where fn make_owned_transmit( socket: PeerSocket, message: &[u8], - allocations: &mut HashMap, + allocations: &mut BTreeMap, now: Instant, ) -> Option> where - RId: Copy + Eq + Hash + PartialEq + fmt::Debug, + RId: Copy + Eq + Hash + PartialEq + Ord + fmt::Debug, { let transmit = match socket { PeerSocket::Direct { diff --git a/rust/connlib/snownet/tests/lib.rs b/rust/connlib/snownet/tests/lib.rs index b5c204dfa..2a3b04e49 100644 --- a/rust/connlib/snownet/tests/lib.rs +++ b/rust/connlib/snownet/tests/lib.rs @@ -4,7 +4,7 @@ use ip_packet::*; use rand::rngs::OsRng; use snownet::{Answer, Client, ClientNode, Event, Node, RelaySocket, Server, ServerNode, Transmit}; use std::{ - collections::{HashSet, VecDeque}, + collections::{BTreeSet, HashSet, VecDeque}, iter, net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4}, time::{Duration, Instant, SystemTime}, @@ -30,13 +30,13 @@ fn migrate_connection_to_new_relay() { )]; let mut alice = TestNode::new(debug_span!("Alice"), alice, "1.1.1.1:80").with_relays( "alice", - HashSet::default(), + BTreeSet::default(), &mut relays, clock.now, ); let mut bob = TestNode::new(debug_span!("Bob"), bob, "2.2.2.2:80").with_relays( "bob", - HashSet::default(), + BTreeSet::default(), &mut relays, clock.now, ); @@ -62,8 +62,8 @@ fn migrate_connection_to_new_relay() { debug_span!("Robert"), ), )]; - alice = alice.with_relays("alice", HashSet::from([1]), &mut relays, clock.now); - bob = bob.with_relays("bob", HashSet::from([1]), &mut relays, clock.now); + alice = alice.with_relays("alice", BTreeSet::from([1]), &mut relays, clock.now); + bob = bob.with_relays("bob", BTreeSet::from([1]), &mut relays, clock.now); // Make some progress. (the fact that we only need 22 clock ticks means we are no relying on timeouts here (22 * 100ms = 2.2s)) for _ in 0..22 { @@ -95,13 +95,13 @@ fn idle_connection_is_closed_after_5_minutes() { )]; let mut alice = TestNode::new(debug_span!("Alice"), alice, "1.1.1.1:80").with_relays( "alice", - HashSet::default(), + BTreeSet::default(), &mut relays, clock.now, ); let mut bob = TestNode::new(debug_span!("Bob"), bob, "2.2.2.2:80").with_relays( "bob", - HashSet::default(), + BTreeSet::default(), &mut relays, clock.now, ); @@ -580,7 +580,7 @@ impl TestNode { fn with_relays( mut self, username: &str, - to_remove: HashSet, + to_remove: BTreeSet, relays: &mut [(u64, TestRelay)], now: Instant, ) -> Self { @@ -597,7 +597,7 @@ impl TestNode { "firezone".to_owned(), ) }) - .collect::>(); + .collect::>(); self.span .in_scope(|| self.node.update_relays(to_remove, &turn_servers, now)); diff --git a/rust/connlib/tunnel/src/client.rs b/rust/connlib/tunnel/src/client.rs index 32987ce65..57fef1578 100644 --- a/rust/connlib/tunnel/src/client.rs +++ b/rust/connlib/tunnel/src/client.rs @@ -27,7 +27,7 @@ use core::fmt; use secrecy::{ExposeSecret as _, Secret}; use snownet::{ClientNode, RelaySocket}; use std::collections::hash_map::Entry; -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; use std::iter; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::str::FromStr; @@ -67,7 +67,7 @@ impl ClientTunnel { self.io.device_mut().set_tun(tun); } - pub fn update_relays(&mut self, to_remove: HashSet, to_add: Vec) { + pub fn update_relays(&mut self, to_remove: BTreeSet, to_add: Vec) { self.role_state .update_relays(to_remove, turn(&to_add), Instant::now()) } @@ -820,8 +820,8 @@ impl ClientState { fn drain_node_events(&mut self) { let mut resources_changed = false; // Track this separately to batch together `ResourcesChanged` events. - let mut added_ice_candidates = HashMap::>::default(); - let mut removed_ice_candidates = HashMap::>::default(); + let mut added_ice_candidates = BTreeMap::>::default(); + let mut removed_ice_candidates = BTreeMap::>::default(); while let Some(event) = self.node.poll_event() { match event { @@ -861,7 +861,7 @@ impl ClientState { }); } - for (conn_id, candidates) in added_ice_candidates.drain() { + for (conn_id, candidates) in added_ice_candidates.into_iter() { self.buffered_events .push_back(ClientEvent::AddedIceCandidates { conn_id, @@ -869,7 +869,7 @@ impl ClientState { }) } - for (conn_id, candidates) in removed_ice_candidates.drain() { + for (conn_id, candidates) in removed_ice_candidates.into_iter() { self.buffered_events .push_back(ClientEvent::RemovedIceCandidates { conn_id, @@ -1053,8 +1053,8 @@ impl ClientState { pub fn update_relays( &mut self, - to_remove: HashSet, - to_add: HashSet<(RelayId, RelaySocket, String, String, String)>, + to_remove: BTreeSet, + to_add: BTreeSet<(RelayId, RelaySocket, String, String, String)>, now: Instant, ) { self.node.update_relays(to_remove, &to_add, now); diff --git a/rust/connlib/tunnel/src/gateway.rs b/rust/connlib/tunnel/src/gateway.rs index 8ffdcb3cb..f2ccb5bc8 100644 --- a/rust/connlib/tunnel/src/gateway.rs +++ b/rust/connlib/tunnel/src/gateway.rs @@ -12,7 +12,7 @@ use connlib_shared::{DomainName, Error, Result, StaticSecret}; use ip_packet::{IpPacket, MutableIpPacket}; use secrecy::{ExposeSecret as _, Secret}; use snownet::{RelaySocket, ServerNode}; -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::time::{Duration, Instant}; @@ -331,8 +331,8 @@ impl GatewayState { Some(_) => {} } - let mut added_ice_candidates = HashMap::>::default(); - let mut removed_ice_candidates = HashMap::>::default(); + let mut added_ice_candidates = BTreeMap::>::default(); + let mut removed_ice_candidates = BTreeMap::>::default(); while let Some(event) = self.node.poll_event() { match event { @@ -361,7 +361,7 @@ impl GatewayState { } } - for (conn_id, candidates) in added_ice_candidates.drain() { + for (conn_id, candidates) in added_ice_candidates.into_iter() { self.buffered_events .push_back(GatewayEvent::AddedIceCandidates { conn_id, @@ -369,7 +369,7 @@ impl GatewayState { }) } - for (conn_id, candidates) in removed_ice_candidates.drain() { + for (conn_id, candidates) in removed_ice_candidates.into_iter() { self.buffered_events .push_back(GatewayEvent::RemovedIceCandidates { conn_id, @@ -398,8 +398,8 @@ impl GatewayState { pub fn update_relays( &mut self, - to_remove: HashSet, - to_add: HashSet<(RelayId, RelaySocket, String, String, String)>, + to_remove: BTreeSet, + to_add: BTreeSet<(RelayId, RelaySocket, String, String, String)>, now: Instant, ) { self.node.update_relays(to_remove, &to_add, now); diff --git a/rust/connlib/tunnel/src/lib.rs b/rust/connlib/tunnel/src/lib.rs index 394c92458..9990dfbe5 100644 --- a/rust/connlib/tunnel/src/lib.rs +++ b/rust/connlib/tunnel/src/lib.rs @@ -12,7 +12,7 @@ use connlib_shared::{ }; use io::Io; use std::{ - collections::{HashMap, HashSet}, + collections::{BTreeSet, HashMap, HashSet}, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}, sync::Arc, task::{Context, Poll}, @@ -181,7 +181,7 @@ impl GatewayTunnel { }) } - pub fn update_relays(&mut self, to_remove: HashSet, to_add: Vec) { + pub fn update_relays(&mut self, to_remove: BTreeSet, to_add: Vec) { self.role_state .update_relays(to_remove, turn(&to_add), Instant::now()) } @@ -255,11 +255,11 @@ impl GatewayTunnel { pub enum ClientEvent { AddedIceCandidates { conn_id: GatewayId, - candidates: HashSet, + candidates: BTreeSet, }, RemovedIceCandidates { conn_id: GatewayId, - candidates: HashSet, + candidates: BTreeSet, }, ConnectionIntent { resource: ResourceId, @@ -294,11 +294,11 @@ pub enum ClientEvent { pub enum GatewayEvent { AddedIceCandidates { conn_id: ClientId, - candidates: HashSet, + candidates: BTreeSet, }, RemovedIceCandidates { conn_id: ClientId, - candidates: HashSet, + candidates: BTreeSet, }, RefreshDns { name: DomainName, diff --git a/rust/connlib/tunnel/src/tests/reference.rs b/rust/connlib/tunnel/src/tests/reference.rs index ee954b962..abe310af1 100644 --- a/rust/connlib/tunnel/src/tests/reference.rs +++ b/rust/connlib/tunnel/src/tests/reference.rs @@ -17,7 +17,7 @@ use prop::collection; use proptest::{prelude::*, sample}; use proptest_state_machine::ReferenceStateMachine; use std::{ - collections::{BTreeMap, HashMap, HashSet}, + collections::{BTreeMap, HashSet}, fmt, iter, net::IpAddr, time::Instant, @@ -31,8 +31,8 @@ pub(crate) struct ReferenceState { pub(crate) now: Instant, pub(crate) utc_now: DateTime, pub(crate) client: Host, - pub(crate) gateways: HashMap>, - pub(crate) relays: HashMap>, + pub(crate) gateways: BTreeMap>, + pub(crate) relays: BTreeMap>, pub(crate) portal: StubPortal, /// All IP addresses a domain resolves to in our test. @@ -66,7 +66,7 @@ impl ReferenceStateMachine for ReferenceState { ( ref_client_host(Just(client_tunnel_ip4), Just(client_tunnel_ip6)), gateways_and_portal(), - collection::hash_map(relay_id(), relay_prototype(), 1..=2), + collection::btree_map(relay_id(), relay_prototype(), 1..=2), global_dns_records(), // Start out with a set of global DNS records so we have something to resolve outside of DNS resources. any::(), Just(Instant::now()), diff --git a/rust/connlib/tunnel/src/tests/strategies.rs b/rust/connlib/tunnel/src/tests/strategies.rs index f20e0ef33..0e4c239fd 100644 --- a/rust/connlib/tunnel/src/tests/strategies.rs +++ b/rust/connlib/tunnel/src/tests/strategies.rs @@ -101,7 +101,7 @@ pub(crate) fn tunnel_ip6s() -> impl Iterator { /// Lastly, we also sample a set of DNS records for the DNS resources that we created. pub(crate) fn gateways_and_portal() -> impl Strategy< Value = ( - HashMap>, + BTreeMap>, StubPortal, HashMap>, ), @@ -132,7 +132,7 @@ pub(crate) fn gateways_and_portal() -> impl Strategy< |(gateways, cidr_resources, dns_resources, gateway_selector)| { let (gateways, gateways_by_site) = gateways.into_iter().fold( ( - HashMap::::default(), + BTreeMap::::default(), HashMap::>::default(), ), |(mut gateways, mut sites), (gid, (gateway, site))| { diff --git a/rust/connlib/tunnel/src/tests/sut.rs b/rust/connlib/tunnel/src/tests/sut.rs index 67014d364..7905b37f1 100644 --- a/rust/connlib/tunnel/src/tests/sut.rs +++ b/rust/connlib/tunnel/src/tests/sut.rs @@ -25,14 +25,8 @@ use proptest_state_machine::{ReferenceStateMachine, StateMachineTest}; use rand::SeedableRng as _; use secrecy::ExposeSecret as _; use snownet::Transmit; -use std::collections::BTreeMap; -use std::{ - collections::{HashMap, HashSet, VecDeque}, - net::IpAddr, - str::FromStr as _, - sync::Arc, - time::Instant, -}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::{collections::HashSet, net::IpAddr, str::FromStr as _, sync::Arc, time::Instant}; use tracing::debug_span; use tracing::subscriber::DefaultGuard; use tracing_subscriber::layer::SubscriberExt; @@ -47,8 +41,8 @@ pub(crate) struct TunnelTest { utc_now: DateTime, pub(crate) client: Host, - pub(crate) gateways: HashMap>, - relays: HashMap>, + pub(crate) gateways: BTreeMap>, + relays: BTreeMap>, drop_direct_client_traffic: bool, network: RoutingTable, @@ -87,7 +81,7 @@ impl StateMachineTest for TunnelTest { (*id, gateway) }) - .collect::>(); + .collect::>(); let relays = ref_state .relays @@ -107,21 +101,21 @@ impl StateMachineTest for TunnelTest { (*id, relay) }) - .collect::>(); + .collect::>(); // Configure client and gateway with the relays. client.exec_mut(|c| { c.sut.update_relays( - HashSet::default(), - HashSet::from_iter(map_explode(relays.iter(), "client")), + BTreeSet::default(), + BTreeSet::from_iter(map_explode(relays.iter(), "client")), ref_state.now, ) }); for (id, gateway) in &mut gateways { gateway.exec_mut(|g| { g.sut.update_relays( - HashSet::default(), - HashSet::from_iter(map_explode(relays.iter(), &format!("gateway_{id}"))), + BTreeSet::default(), + BTreeSet::from_iter(map_explode(relays.iter(), &format!("gateway_{id}"))), ref_state.now, ) }); @@ -267,8 +261,8 @@ impl StateMachineTest for TunnelTest { // In prod, we reconnect to the portal and receive a new `init` message. c.sut.update_relays( - HashSet::default(), - HashSet::from_iter(map_explode(state.relays.iter(), "client")), + BTreeSet::default(), + BTreeSet::from_iter(map_explode(state.relays.iter(), "client")), ref_state.now, ); c.sut @@ -279,7 +273,7 @@ impl StateMachineTest for TunnelTest { let ipv4 = state.client.inner().sut.tunnel_ip4().unwrap(); let ipv6 = state.client.inner().sut.tunnel_ip6().unwrap(); let upstream_dns = ref_state.client.inner().upstream_dns_resolvers.clone(); - let relays = HashSet::from_iter(map_explode(state.relays.iter(), "client")); + let relays = BTreeSet::from_iter(map_explode(state.relays.iter(), "client")); let all_resources = ref_state.client.inner().all_resources(); // Simulate receiving `init`. @@ -290,7 +284,7 @@ impl StateMachineTest for TunnelTest { upstream_dns, }); c.sut - .update_relays(HashSet::default(), relays, ref_state.now); + .update_relays(BTreeSet::default(), relays, ref_state.now); c.sut.set_resources(all_resources); }); } @@ -434,7 +428,6 @@ impl TunnelTest { if self.handle_timeout(self.now, self.utc_now) { continue; } - break; } } diff --git a/rust/connlib/tunnel/src/utils.rs b/rust/connlib/tunnel/src/utils.rs index 52c627a76..72bce5139 100644 --- a/rust/connlib/tunnel/src/utils.rs +++ b/rust/connlib/tunnel/src/utils.rs @@ -3,9 +3,9 @@ use connlib_shared::messages::{Relay, RelayId}; use ip_network::{IpNetwork, Ipv4Network, Ipv6Network}; use itertools::Itertools; use snownet::RelaySocket; -use std::{collections::HashSet, net::SocketAddr, time::Instant}; +use std::{collections::BTreeSet, net::SocketAddr, time::Instant}; -pub fn turn(relays: &[Relay]) -> HashSet<(RelayId, RelaySocket, String, String, String)> { +pub fn turn(relays: &[Relay]) -> BTreeSet<(RelayId, RelaySocket, String, String, String)> { relays .iter() .filter_map(|r| { diff --git a/rust/gateway/src/eventloop.rs b/rust/gateway/src/eventloop.rs index 98e9ab581..7110dab08 100644 --- a/rust/gateway/src/eventloop.rs +++ b/rust/gateway/src/eventloop.rs @@ -14,7 +14,7 @@ use firezone_tunnel::GatewayTunnel; use futures::channel::mpsc; use futures_bounded::Timeout; use phoenix_channel::PhoenixChannel; -use std::collections::HashSet; +use std::collections::BTreeSet; use std::convert::Infallible; use std::net::IpAddr; use std::task::{Context, Poll}; @@ -224,12 +224,12 @@ impl Eventloop { .. } => self .tunnel - .update_relays(HashSet::from_iter(disconnected_ids), connected), + .update_relays(BTreeSet::from_iter(disconnected_ids), connected), phoenix_channel::Event::InboundMessage { msg: IngressMessages::Init(init), .. } => { - self.tunnel.update_relays(HashSet::default(), init.relays); + self.tunnel.update_relays(BTreeSet::default(), init.relays); // FIXME(tech-debt): Currently, the `Tunnel` creates the TUN device as part of `set_interface`. // For the gateway, it doesn't do anything else so in an ideal world, we would cause the side-effect out here and just pass an opaque `Device` to the `Tunnel`. diff --git a/rust/gateway/src/messages.rs b/rust/gateway/src/messages.rs index fc8a7e42f..a0f64dac7 100644 --- a/rust/gateway/src/messages.rs +++ b/rust/gateway/src/messages.rs @@ -1,4 +1,4 @@ -use std::{collections::HashSet, net::IpAddr}; +use std::{collections::BTreeSet, net::IpAddr}; use chrono::{serde::ts_seconds_option, DateTime, Utc}; use connlib_shared::{ @@ -121,7 +121,7 @@ pub struct ClientsIceCandidates { /// Client's id the ice candidates are meant for pub client_ids: Vec, /// Actual RTC ice candidates - pub candidates: HashSet, + pub candidates: BTreeSet, } /// A client's ice candidate message. diff --git a/rust/snownet-tests/src/main.rs b/rust/snownet-tests/src/main.rs index b9337ce73..8522659c0 100644 --- a/rust/snownet-tests/src/main.rs +++ b/rust/snownet-tests/src/main.rs @@ -1,5 +1,5 @@ use std::{ - collections::HashSet, + collections::BTreeSet, future::poll_fn, net::{Ipv4Addr, SocketAddrV4}, str::FromStr, @@ -69,7 +69,7 @@ async fn main() -> Result<()> { "firezone".to_owned(), ) }); - let relays = HashSet::from_iter(relay_stun_only.into_iter().chain(relay_valid_turn)); + let relays = BTreeSet::from_iter(relay_stun_only.into_iter().chain(relay_valid_turn)); tracing::info!(%listen_addr); @@ -89,7 +89,7 @@ async fn main() -> Result<()> { match role { Role::Dialer => { let mut pool = ClientNode::::new(private_key); - pool.update_relays(HashSet::new(), &relays, Instant::now()); + pool.update_relays(BTreeSet::new(), &relays, Instant::now()); let offer = pool.new_connection(1, Instant::now(), Instant::now()); @@ -169,7 +169,7 @@ async fn main() -> Result<()> { } Role::Listener => { let mut pool = ServerNode::::new(private_key); - pool.update_relays(HashSet::new(), &relays, Instant::now()); + pool.update_relays(BTreeSet::new(), &relays, Instant::now()); let offer = redis_connection .blpop::<_, (String, wire::Offer)>("offers", 10.0)