From bad5108baf9bbeb603be2e881d3bd1263abdc6de Mon Sep 17 00:00:00 2001
From: Datong Sun <dndx@idndx.com>
Date: Mon, 19 Aug 2024 07:04:09 -0700
Subject: [PATCH] fix(fake-tcp): when `connect()`-ing, attempt to get ephemeral
 port 5 times to reduce the chance of paniking

Fixes #79
---
 fake-tcp/src/lib.rs | 61 +++++++++++++++++++++++++++++----------------
 1 file changed, 39 insertions(+), 22 deletions(-)

diff --git a/fake-tcp/src/lib.rs b/fake-tcp/src/lib.rs
index a66da47..79d988b 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<Socket> {
         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 _ 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();
-            assert!(tuples.insert(tuple, incoming.clone()).is_none());
+            {
+                let mut tuples = self.shared.tuples.write().unwrap();
+                match tuples.entry(tuple) {
+                    Occupied(_) => {
+                        // port conflict, try again
+                        continue;
+                    }
+                    Vacant(v) => {
+                        v.insert(incoming.clone());
+                    }
+                }
+            }
+
+            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(