From 02949be245606137a4311700c8becbd6cc249d56 Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Thu, 4 Dec 2025 21:48:32 -0500
Subject: [PATCH] Support connection testing in native

---
 clients.go         | 26 --------------------------
 clients/clients.go | 24 +++++++++++++++++-------
 2 files changed, 17 insertions(+), 33 deletions(-)

diff --git a/clients.go b/clients.go
index e95eadb..3f28f4c 100644
--- a/clients.go
+++ b/clients.go
@@ -7,15 +7,11 @@ import (
 	wgnetstack "github.com/fosrl/newt/clients"
 	"github.com/fosrl/newt/clients/permissions"
 	"github.com/fosrl/newt/logger"
-	"github.com/fosrl/newt/netstack2"
 	"github.com/fosrl/newt/websocket"
 	"golang.zx2c4.com/wireguard/tun/netstack"
-
-	"github.com/fosrl/newt/wgtester"
 )
 
 var wgService *clients.WireGuardService
-var wgTesterServer *wgtester.Server
 var ready bool
 
 func setupClients(client *websocket.Client) {
@@ -46,23 +42,6 @@ func setupClients(client *websocket.Client) {
 		logger.Fatal("Failed to create WireGuard service: %v", err)
 	}
 
-	// // Set up callback to restart wgtester with netstack when WireGuard is ready
-	wgService.SetOnNetstackReady(func(tnet *netstack2.Net) {
-
-		wgTesterServer = wgtester.NewServerWithNetstack("0.0.0.0", wgService.Port, id, tnet) // TODO: maybe make this the same ip of the wg server?
-		err := wgTesterServer.Start()
-		if err != nil {
-			logger.Error("Failed to start WireGuard tester server: %v", err)
-		}
-	})
-
-	wgService.SetOnNetstackClose(func() {
-		if wgTesterServer != nil {
-			wgTesterServer.Stop()
-			wgTesterServer = nil
-		}
-	})
-
 	client.OnTokenUpdate(func(token string) {
 		wgService.SetToken(token)
 	})
@@ -82,11 +61,6 @@ func closeClients() {
 		wgService.Close()
 		wgService = nil
 	}
-
-	if wgTesterServer != nil {
-		wgTesterServer.Stop()
-		wgTesterServer = nil
-	}
 }
 
 func clientsHandleNewtConnection(publicKey string, endpoint string) {
diff --git a/clients/clients.go b/clients/clients.go
index 4ce1a83..d5fb5f3 100644
--- a/clients/clients.go
+++ b/clients/clients.go
@@ -20,6 +20,7 @@ import (
 	"github.com/fosrl/newt/network"
 	"github.com/fosrl/newt/util"
 	"github.com/fosrl/newt/websocket"
+	"github.com/fosrl/newt/wgtester"
 	"golang.zx2c4.com/wireguard/device"
 	"golang.zx2c4.com/wireguard/ipc"
 	"golang.zx2c4.com/wireguard/tun"
@@ -100,6 +101,7 @@ type WireGuardService struct {
 	directRelayWg      sync.WaitGroup
 	netstackListener   net.PacketConn
 	netstackListenerMu sync.Mutex
+	wgTesterServer     *wgtester.Server
 }
 
 func NewWireGuardService(interfaceName string, mtu int, host string, newtId string, wsClient *websocket.Client, dns string, useNativeInterface bool) (*WireGuardService, error) {
@@ -221,6 +223,11 @@ func (s *WireGuardService) Close() {
 		s.sharedBind = nil
 		logger.Info("Released shared UDP bind")
 	}
+
+	if s.wgTesterServer != nil {
+		s.wgTesterServer.Stop()
+		s.wgTesterServer = nil
+	}
 }
 
 func (s *WireGuardService) SetToken(token string) {
@@ -565,6 +572,12 @@ func (s *WireGuardService) ensureWireguardInterface(wgconfig WgConfig) error {
 			return fmt.Errorf("failed to configure interface: %v", err)
 		}
 
+		s.wgTesterServer = wgtester.NewServer("0.0.0.0", s.Port, s.newtId) // TODO: maybe make this the same ip of the wg server?
+		err = s.wgTesterServer.Start()
+		if err != nil {
+			logger.Error("Failed to start WireGuard tester server: %v", err)
+		}
+
 		logger.Info("WireGuard native device created and configured on %s", interfaceName)
 
 		s.mu.Unlock()
@@ -612,16 +625,13 @@ func (s *WireGuardService) ensureWireguardInterface(wgconfig WgConfig) error {
 
 	logger.Info("WireGuard netstack device created and configured")
 
-	// Store callback and tnet reference before releasing mutex
-	callback := s.onNetstackReady
-	tnet := s.tnet
-
 	// Release the mutex before calling the callback
 	s.mu.Unlock()
 
-	// Call the callback if it's set to notify that netstack is ready
-	if callback != nil {
-		callback(tnet)
+	s.wgTesterServer = wgtester.NewServerWithNetstack("0.0.0.0", s.Port, s.newtId, s.tnet) // TODO: maybe make this the same ip of the wg server?
+	err = s.wgTesterServer.Start()
+	if err != nil {
+		logger.Error("Failed to start WireGuard tester server: %v", err)
 	}
 
 	// Note: we already unlocked above, so don't use defer unlock