From 79f9f476d9c8d990ef2aec17889ee20c62177c17 Mon Sep 17 00:00:00 2001
From: Datong Sun <dndx@idndx.com>
Date: Mon, 26 Aug 2024 05:43:19 -0700
Subject: [PATCH] retry until succeed

---
 fake-tcp/src/lib.rs | 47 ++++++++++++++++++++++++---------------------
 1 file changed, 25 insertions(+), 22 deletions(-)

diff --git a/fake-tcp/src/lib.rs b/fake-tcp/src/lib.rs
index 79d988b..bf6647a 100644
--- a/fake-tcp/src/lib.rs
+++ b/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,33 +413,36 @@ impl Stack {
                 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 sock;
 
             {
                 let mut tuples = self.shared.tuples.write().unwrap();
-                match tuples.entry(tuple) {
-                    Occupied(_) => {
-                        // port conflict, try again
-                        continue;
-                    }
-                    Vacant(v) => {
-                        v.insert(incoming.clone());
-                    }
+                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,
+                    addr,
+                    None,
+                    State::Idle,
+                );
+
+                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