retry until succeed

fake-tcp/src/lib.rs

@@ -47,12 +47,13 @@ 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};
-use std::sync::atomic::{AtomicU32, Ordering};
-use std::sync::{Arc, RwLock};
+use std::sync::{
+    atomic::{AtomicU32, Ordering},
+    Arc, RwLock,
+};
 use tokio::sync::broadcast;
 use tokio::sync::mpsc;
 use tokio::time;
@@ -402,8 +403,7 @@ impl Stack {
     /// the connection attempt failed.
     pub async fn connect(&mut self, addr: SocketAddr) -> Option<Socket> {
         let mut rng = SmallRng::from_entropy();
-        for _ in 0..5 {
-            let local_port: u16 = rng.gen_range(1024..=65535);
+        for local_port in rng.gen_range(32768..=60999)..=60999 {
             let local_addr = SocketAddr::new(
                 if addr.is_ipv4() {
                     IpAddr::V4(self.local_ip)
@@ -413,7 +413,20 @@ impl Stack {
                 local_port,
             );
             let tuple = AddrTuple::new(local_addr, addr);
-            let (mut sock, incoming) = Socket::new(
+            let mut sock;
+
+            {
+                let mut tuples = self.shared.tuples.write().unwrap();
+                if tuples.contains_key(&tuple) {
+                    trace!(
+                        "Fake TCP connection to {}, local port number {} already in use, trying another one",
+                        addr, local_port
+                    );
+                    continue;
+                }
+
+                let incoming;
+                (sock, incoming) = Socket::new(
                     self.shared.clone(),
                     self.shared.tun.choose(&mut rng).unwrap().clone(),
                     local_addr,
@@ -422,24 +435,14 @@ impl Stack {
                     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());
-                    }
-                }
+                assert!(tuples.insert(tuple, incoming).is_none());
             }
 
             return sock.connect().await.map(|_| sock);
         }
 
         error!(
-            "Fake TCP connection to {} failed, local port number not available after 5 attempts",
+            "Fake TCP connection to {} failed, emphemeral port number exhausted",
             addr
         );
         None