From 407145ab845d646cbebdd989d87ec02e99061b41 Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Sun, 30 Nov 2025 20:42:22 -0500
Subject: [PATCH] Fix bind issue when switching orgs

---
 main.go      |  1 +
 olm/olm.go   | 80 ++++++++++++++++++++++++++++++++--------------------
 olm/types.go |  2 ++
 3 files changed, 52 insertions(+), 31 deletions(-)

diff --git a/main.go b/main.go
index 1282469..5e4e1d9 100644
--- a/main.go
+++ b/main.go
@@ -235,6 +235,7 @@ func runOlmMainWithArgs(ctx context.Context, args []string) {
 			OrgID:                config.OrgID,
 			OverrideDNS:          config.OverrideDNS,
 			EnableUAPI:           true,
+			DisableRelay:         true,
 		}
 		go olm.StartTunnel(tunnelConfig)
 	} else {
diff --git a/olm/olm.go b/olm/olm.go
index 3444a94..b1ffb12 100644
--- a/olm/olm.go
+++ b/olm/olm.go
@@ -52,6 +52,41 @@ var (
 	peerManager      *peers.PeerManager
 )
 
+// initSharedBindAndHolepunch creates the shared UDP socket and holepunch manager.
+// This is used during initial tunnel setup and when switching organizations.
+func initSharedBindAndHolepunch(clientID string) error {
+	sourcePort, err := util.FindAvailableUDPPort(49152, 65535)
+	if err != nil {
+		return fmt.Errorf("failed to find available UDP port: %w", err)
+	}
+
+	localAddr := &net.UDPAddr{
+		Port: int(sourcePort),
+		IP:   net.IPv4zero,
+	}
+
+	udpConn, err := net.ListenUDP("udp", localAddr)
+	if err != nil {
+		return fmt.Errorf("failed to create UDP socket: %w", err)
+	}
+
+	sharedBind, err = bind.New(udpConn)
+	if err != nil {
+		udpConn.Close()
+		return fmt.Errorf("failed to create shared bind: %w", err)
+	}
+
+	// Add a reference for the hole punch senders (creator already has one reference for WireGuard)
+	sharedBind.AddRef()
+
+	logger.Info("Created shared UDP socket on port %d (refcount: %d)", sourcePort, sharedBind.GetRefCount())
+
+	// Create the holepunch manager
+	holePunchManager = holepunch.NewManager(sharedBind, clientID, "olm")
+
+	return nil
+}
+
 func Init(ctx context.Context, config GlobalConfig) {
 	globalConfig = config
 	globalCtx = ctx
@@ -220,39 +255,12 @@ func StartTunnel(config TunnelConfig) {
 		return
 	}
 
-	// Create shared UDP socket for both holepunch and WireGuard
-	sourcePort, err := util.FindAvailableUDPPort(49152, 65535)
-	if err != nil {
-		logger.Error("Error finding available port: %v", err)
+	// Create shared UDP socket and holepunch manager
+	if err := initSharedBindAndHolepunch(id); err != nil {
+		logger.Error("%v", err)
 		return
 	}
 
-	localAddr := &net.UDPAddr{
-		Port: int(sourcePort),
-		IP:   net.IPv4zero,
-	}
-
-	udpConn, err := net.ListenUDP("udp", localAddr)
-	if err != nil {
-		logger.Error("Failed to create shared UDP socket: %v", err)
-		return
-	}
-
-	sharedBind, err = bind.New(udpConn)
-	if err != nil {
-		logger.Error("Failed to create shared bind: %v", err)
-		udpConn.Close()
-		return
-	}
-
-	// Add a reference for the hole punch senders (creator already has one reference for WireGuard)
-	sharedBind.AddRef()
-
-	logger.Info("Created shared UDP socket on port %d (refcount: %d)", sourcePort, sharedBind.GetRefCount())
-
-	// Create the holepunch manager
-	holePunchManager = holepunch.NewManager(sharedBind, id, "olm")
-
 	olm.RegisterHandler("olm/wg/holepunch/all", func(msg websocket.WSMessage) {
 		logger.Debug("Received message: %v", msg.Data)
 
@@ -467,7 +475,7 @@ func StartTunnel(config TunnelConfig) {
 			util.FixKey(privateKey.String()),
 			olm,
 			dev,
-			config.Holepunch,
+			config.Holepunch && !config.DisableRelay, // Enable relay only if holepunching is enabled and DisableRelay is false
 			middleDev,
 			interfaceIP,
 		)
@@ -861,6 +869,10 @@ func Close() {
 		peerMonitor = nil
 	}
 
+	if peerManager != nil {
+		peerManager = nil
+	}
+
 	if uapiListener != nil {
 		uapiListener.Close()
 		uapiListener = nil
@@ -976,8 +988,14 @@ func SwitchOrg(orgID string) error {
 	// Mark as not connected to trigger re-registration
 	connected = false
 
+	// Close existing tunnel resources (but keep websocket alive)
 	Close()
 
+	// Recreate sharedBind and holepunch manager - needed because Close() releases them
+	if err := initSharedBindAndHolepunch(olmClient.GetConfig().ID); err != nil {
+		return err
+	}
+
 	// Clear peer statuses in API
 	apiServer.SetRegistered(false)
 
diff --git a/olm/types.go b/olm/types.go
index cae876b..39fef25 100644
--- a/olm/types.go
+++ b/olm/types.go
@@ -81,4 +81,6 @@ type TunnelConfig struct {
 	EnableUAPI bool
 
 	OverrideDNS bool
+
+	DisableRelay bool
 }