fix(fake-tcp): when connect()-ing, attempt to get ephemeral port 5 times to reduce the chance

of paniking

Fixes #79
This commit is contained in:
Datong Sun 2024-08-19 07:04:09 -07:00
parent 90b93370ce
commit bad5108baf

View File

@ -47,6 +47,7 @@ use log::{error, info, trace, warn};
use packet::*; use packet::*;
use pnet::packet::{tcp, Packet}; use pnet::packet::{tcp, Packet};
use rand::prelude::*; use rand::prelude::*;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
use std::fmt; use std::fmt;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@ -401,7 +402,8 @@ impl Stack {
/// the connection attempt failed. /// the connection attempt failed.
pub async fn connect(&mut self, addr: SocketAddr) -> Option<Socket> { pub async fn connect(&mut self, addr: SocketAddr) -> Option<Socket> {
let mut rng = SmallRng::from_entropy(); let mut rng = SmallRng::from_entropy();
let local_port: u16 = rng.gen_range(1024..65535); for _ in 0..5 {
let local_port: u16 = rng.gen_range(1024..=65535);
let local_addr = SocketAddr::new( let local_addr = SocketAddr::new(
if addr.is_ipv4() { if addr.is_ipv4() {
IpAddr::V4(self.local_ip) IpAddr::V4(self.local_ip)
@ -422,10 +424,25 @@ impl Stack {
{ {
let mut tuples = self.shared.tuples.write().unwrap(); let mut tuples = self.shared.tuples.write().unwrap();
assert!(tuples.insert(tuple, incoming.clone()).is_none()); match tuples.entry(tuple) {
Occupied(_) => {
// port conflict, try again
continue;
}
Vacant(v) => {
v.insert(incoming.clone());
}
}
} }
sock.connect().await.map(|_| sock) return 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( async fn reader_task(