diff --git a/fake-tcp/src/lib.rs b/fake-tcp/src/lib.rs index a66da47..855bf36 100644 --- a/fake-tcp/src/lib.rs +++ b/fake-tcp/src/lib.rs @@ -47,6 +47,7 @@ use log::{error, info, trace, warn}; use packet::*; use pnet::packet::{tcp, Packet}; use rand::prelude::*; +use std::collections::hash_map::Entry::{Occupied, Vacant}; use std::collections::{HashMap, HashSet}; use std::fmt; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; @@ -401,31 +402,47 @@ impl Stack { /// the connection attempt failed. pub async fn connect(&mut self, addr: SocketAddr) -> Option { let mut rng = SmallRng::from_entropy(); - let local_port: u16 = rng.gen_range(1024..65535); - let local_addr = SocketAddr::new( - if addr.is_ipv4() { - IpAddr::V4(self.local_ip) - } else { - IpAddr::V6(self.local_ip6.expect("IPv6 local address undefined")) - }, - local_port, - ); - let tuple = AddrTuple::new(local_addr, addr); - let (mut sock, incoming) = Socket::new( - self.shared.clone(), - self.shared.tun.choose(&mut rng).unwrap().clone(), - local_addr, - addr, - None, - State::Idle, - ); + for _i in 0..5 { + let local_port: u16 = rng.gen_range(1024..=65535); + let local_addr = SocketAddr::new( + if addr.is_ipv4() { + IpAddr::V4(self.local_ip) + } else { + IpAddr::V6(self.local_ip6.expect("IPv6 local address undefined")) + }, + local_port, + ); + let tuple = AddrTuple::new(local_addr, addr); + let (mut sock, incoming) = Socket::new( + self.shared.clone(), + self.shared.tun.choose(&mut rng).unwrap().clone(), + local_addr, + addr, + None, + State::Idle, + ); + + { + let mut tuples = self.shared.tuples.write().unwrap(); + match tuples.entry(tuple) { + Occupied(_) => { + // port conflict, try again + continue; + } + Vacant(v) => { + v.insert(incoming.clone()); + } + } + } - { - let mut tuples = self.shared.tuples.write().unwrap(); - assert!(tuples.insert(tuple, incoming.clone()).is_none()); + return sock.connect().await.map(|_| sock); } - sock.connect().await.map(|_| sock) + error!( + "Fake TCP connection to {} failed, local port number not available after 5 attempts", + addr + ); + None } async fn reader_task(