From 8795c57b2e806a2e4e71ea026f920a7e0006b00a Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Sat, 22 Feb 2025 12:53:23 -0500
Subject: [PATCH] HP works!

---
 go.mod   |  3 ++-
 go.sum   |  4 ++++
 wg/wg.go | 32 +++++++++++++++++++++++++-------
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/go.mod b/go.mod
index 08b4716..c6931ef 100644
--- a/go.mod
+++ b/go.mod
@@ -22,7 +22,8 @@ require (
 	github.com/mdlayher/socket v0.5.1 // indirect
 	github.com/vishvananda/netns v0.0.4 // indirect
 	golang.org/x/crypto v0.31.0 // indirect
-	golang.org/x/sync v0.10.0 // indirect
+	golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa // indirect
+	golang.org/x/sync v0.11.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect
 	golang.org/x/time v0.7.0 // indirect
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
diff --git a/go.sum b/go.sum
index 2f56ede..a0deda0 100644
--- a/go.sum
+++ b/go.sum
@@ -20,10 +20,14 @@ github.com/vishvananda/netns v0.0.4 h1:Oeaw1EM2JMxD51g9uhtC0D7erkIjgmj8+JZc26m1Y
 github.com/vishvananda/netns v0.0.4/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
 golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U=
 golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
+golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa h1:t2QcU6V556bFjYgu4L6C+6VrCPyJZ+eyRsABUPs1mz4=
+golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa/go.mod h1:BHOTPb3L19zxehTsLoJXVaTktb06DFgmdW6Wb9s8jqk=
 golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
 golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
 golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ=
 golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
+golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w=
+golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
diff --git a/wg/wg.go b/wg/wg.go
index 58bf02a..6883ca9 100644
--- a/wg/wg.go
+++ b/wg/wg.go
@@ -11,6 +11,7 @@ import (
 	"github.com/fosrl/newt/logger"
 	"github.com/fosrl/newt/websocket"
 	"github.com/vishvananda/netlink"
+	"golang.org/x/exp/rand"
 	"golang.zx2c4.com/wireguard/conn"
 	"golang.zx2c4.com/wireguard/wgctrl"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
@@ -78,23 +79,31 @@ func FindAvailableUDPPort(minPort, maxPort uint16) (uint16, error) {
 		return 0, fmt.Errorf("invalid port range: min=%d, max=%d", minPort, maxPort)
 	}
 
-	for port := minPort; port <= maxPort; port++ {
-		// Create the UDP address to test
+	// Create a slice of all ports in the range
+	portRange := make([]uint16, maxPort-minPort+1)
+	for i := range portRange {
+		portRange[i] = minPort + uint16(i)
+	}
+
+	// Fisher-Yates shuffle to randomize the port order
+	rand.Seed(uint64(time.Now().UnixNano()))
+	for i := len(portRange) - 1; i > 0; i-- {
+		j := rand.Intn(i + 1)
+		portRange[i], portRange[j] = portRange[j], portRange[i]
+	}
+
+	// Try each port in the randomized order
+	for _, port := range portRange {
 		addr := &net.UDPAddr{
 			IP:   net.ParseIP("127.0.0.1"),
 			Port: int(port),
 		}
-
-		// Attempt to create a UDP listener
 		conn, err := net.ListenUDP("udp", addr)
 		if err != nil {
 			continue // Port is in use or there was an error, try next port
 		}
-
-		// Close the connection immediately
 		_ = conn.SetDeadline(time.Now())
 		conn.Close()
-
 		return port, nil
 	}
 
@@ -150,6 +159,10 @@ func NewWireGuardService(interfaceName string, mtu int, reachableAt string, gene
 		stopHolepunch: make(chan struct{}),
 	}
 
+	if err := service.sendUDPHolePunch(host + ":21820"); err != nil {
+		logger.Error("Failed to send UDP hole punch: %v", err)
+	}
+
 	// start the UDP holepunch
 	go service.keepSendingUDPHolePunch(host)
 
@@ -200,6 +213,9 @@ func (s *WireGuardService) handleConfig(msg websocket.WSMessage) {
 	}
 	s.config = config
 
+	// stop the holepunch
+	close(s.stopHolepunch)
+
 	// Ensure the WireGuard interface and peers are configured
 	if err := s.ensureWireguardInterface(config); err != nil {
 		logger.Error("Failed to ensure WireGuard interface: %v", err)
@@ -698,6 +714,8 @@ func (s *WireGuardService) sendUDPHolePunch(serverAddr string) error {
 		return fmt.Errorf("failed to send UDP packet: %v", err)
 	}
 
+	logger.Info("Sent UDP hole punch to %s", serverAddr)
+
 	return nil
 }