From ddea0011709f091acd902d871822187f75fb091d Mon Sep 17 00:00:00 2001
From: Maycon Santos <mlsmaycon@gmail.com>
Date: Wed, 21 Aug 2024 19:24:40 +0200
Subject: [PATCH] [client] Refactor free port function (#2455)

Rely on net.ListenUDP to get an available port for wireguard in case the configured one is in use

---------

Co-authored-by: Viktor Liu <17948409+lixmal@users.noreply.github.com>
---
 client/internal/connect.go      | 48 ++++++++++++++++++------
 client/internal/connect_test.go | 66 +++++++++++++++++----------------
 2 files changed, 71 insertions(+), 43 deletions(-)

diff --git a/client/internal/connect.go b/client/internal/connect.go
index 1cfabe910de..3937b7846ec 100644
--- a/client/internal/connect.go
+++ b/client/internal/connect.go
@@ -397,19 +397,43 @@ func statusRecorderToSignalConnStateNotifier(statusRecorder *peer.Status) signal
 	return notifier
 }
 
-func freePort(start int) (int, error) {
+// freePort attempts to determine if the provided port is available, if not it will ask the system for a free port.
+func freePort(initPort int) (int, error) {
 	addr := net.UDPAddr{}
-	if start == 0 {
-		start = iface.DefaultWgPort
+	if initPort == 0 {
+		initPort = iface.DefaultWgPort
 	}
-	for x := start; x <= 65535; x++ {
-		addr.Port = x
-		conn, err := net.ListenUDP("udp", &addr)
-		if err != nil {
-			continue
-		}
-		conn.Close()
-		return x, nil
+
+	addr.Port = initPort
+
+	conn, err := net.ListenUDP("udp", &addr)
+	if err == nil {
+		closeConnWithLog(conn)
+		return initPort, nil
+	}
+
+	// if the port is already in use, ask the system for a free port
+	addr.Port = 0
+	conn, err = net.ListenUDP("udp", &addr)
+	if err != nil {
+		return 0, fmt.Errorf("unable to get a free port: %v", err)
+	}
+
+	udpAddr, ok := conn.LocalAddr().(*net.UDPAddr)
+	if !ok {
+		return 0, errors.New("wrong address type when getting a free port")
+	}
+	closeConnWithLog(conn)
+	return udpAddr.Port, nil
+}
+
+func closeConnWithLog(conn *net.UDPConn) {
+	startClosing := time.Now()
+	err := conn.Close()
+	if err != nil {
+		log.Warnf("closing probe port %d failed: %v. NetBird will still attempt to use this port for connection.", conn.LocalAddr().(*net.UDPAddr).Port, err)
+	}
+	if time.Since(startClosing) > time.Second {
+		log.Warnf("closing the testing port %d took %s. Usually it is safe to ignore, but continuous warnings may indicate a problem.", conn.LocalAddr().(*net.UDPAddr).Port, time.Since(startClosing))
 	}
-	return 0, errors.New("no free ports")
 }
diff --git a/client/internal/connect_test.go b/client/internal/connect_test.go
index 6f4a6bbb7c7..78b4b06e891 100644
--- a/client/internal/connect_test.go
+++ b/client/internal/connect_test.go
@@ -7,51 +7,55 @@ import (
 
 func Test_freePort(t *testing.T) {
 	tests := []struct {
-		name    string
-		port    int
-		want    int
-		wantErr bool
+		name        string
+		port        int
+		want        int
+		shouldMatch bool
 	}{
 		{
-			name:    "available",
-			port:    51820,
-			want:    51820,
-			wantErr: false,
+			name:        "not provided, fallback to default",
+			port:        0,
+			want:        51820,
+			shouldMatch: true,
 		},
 		{
-			name:    "notavailable",
-			port:    51830,
-			want:    51831,
-			wantErr: false,
+			name:        "provided and available",
+			port:        51821,
+			want:        51821,
+			shouldMatch: true,
 		},
 		{
-			name:    "noports",
-			port:    65535,
-			want:    0,
-			wantErr: true,
+			name:        "provided and not available",
+			port:        51830,
+			want:        51830,
+			shouldMatch: false,
 		},
 	}
+	c1, err := net.ListenUDP("udp", &net.UDPAddr{Port: 51830})
+	if err != nil {
+		t.Errorf("freePort error = %v", err)
+	}
+	defer func(c1 *net.UDPConn) {
+		_ = c1.Close()
+	}(c1)
+
 	for _, tt := range tests {
 
-		c1, err := net.ListenUDP("udp", &net.UDPAddr{Port: 51830})
-		if err != nil {
-			t.Errorf("freePort error = %v", err)
-		}
-		c2, err := net.ListenUDP("udp", &net.UDPAddr{Port: 65535})
-		if err != nil {
-			t.Errorf("freePort error = %v", err)
-		}
 		t.Run(tt.name, func(t *testing.T) {
 			got, err := freePort(tt.port)
-			if (err != nil) != tt.wantErr {
-				t.Errorf("freePort() error = %v, wantErr %v", err, tt.wantErr)
-				return
+
+			if err != nil {
+				t.Errorf("got an error while getting free port: %v", err)
 			}
-			if got != tt.want {
-				t.Errorf("freePort() = %v, want %v", got, tt.want)
+
+			if tt.shouldMatch && got != tt.want {
+				t.Errorf("got a different port %v, want %v", got, tt.want)
+			}
+
+			if !tt.shouldMatch && got == tt.want {
+				t.Errorf("got the same port %v, want a different port", tt.want)
 			}
 		})
-		c1.Close()
-		c2.Close()
+
 	}
 }