From 56377ec87e1cfcf020f67a15f99a85e81f750ea0 Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Thu, 24 Jul 2025 20:46:33 -0700
Subject: [PATCH] Exit well

---
 main.go              |  7 +++++--
 wg/wg.go             | 14 +++++++++-----
 wgtester/wgtester.go |  8 +++++++-
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/main.go b/main.go
index daa7964..7557d6e 100644
--- a/main.go
+++ b/main.go
@@ -798,10 +798,13 @@ persistent_keepalive_interval=5`, fixKey(privateKey.String()), fixKey(wgData.Pub
 	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
 	<-sigCh
 
-	dev.Close()
-
+	// Close clients first (including WGTester)
 	closeClients()
 
+	if dev != nil {
+		dev.Close()
+	}
+
 	if pm != nil {
 		pm.Stop()
 	}
diff --git a/wg/wg.go b/wg/wg.go
index a09cab8..364b3c8 100644
--- a/wg/wg.go
+++ b/wg/wg.go
@@ -55,6 +55,7 @@ type WireGuardService struct {
 	wgClient          *wgctrl.Client
 	config            WgConfig
 	key               wgtypes.Key
+	keyFilePath       string
 	newtId            string
 	lastReadings      map[string]PeerReading
 	mu                sync.Mutex
@@ -179,6 +180,7 @@ func NewWireGuardService(interfaceName string, mtu int, generateAndSaveKeyTo str
 		client:        wsClient,
 		wgClient:      wgClient,
 		key:           key,
+		keyFilePath:   generateAndSaveKeyTo,
 		newtId:        newtId,
 		host:          host,
 		lastReadings:  make(map[string]PeerReading),
@@ -229,9 +231,11 @@ func (s *WireGuardService) Close(rm bool) {
 		}
 
 		// Remove the private key file
-		if err := os.Remove(s.key.String()); err != nil {
-			logger.Error("Failed to remove private key file: %v", err)
-		}
+		// if s.keyFilePath != "" {
+		// 	if err := os.Remove(s.keyFilePath); err != nil {
+		// 		logger.Error("Failed to remove private key file: %v", err)
+		// 	}
+		// }
 	}
 }
 
@@ -251,7 +255,7 @@ func (s *WireGuardService) SetToken(token string) {
 
 func (s *WireGuardService) LoadRemoteConfig() error {
 	s.stopGetConfig = s.client.SendMessageInterval("newt/wg/get-config", map[string]interface{}{
-		"publicKey": fmt.Sprintf("%s", s.key.PublicKey().String()),
+		"publicKey": s.key.PublicKey().String(),
 		"port":      s.Port,
 	}, 2*time.Second)
 
@@ -638,7 +642,7 @@ func (s *WireGuardService) handleUpdatePeer(msg websocket.WSMessage) {
 	}
 
 	// Only update AllowedIPs if provided in the request
-	if request.AllowedIPs != nil && len(request.AllowedIPs) > 0 {
+	if len(request.AllowedIPs) > 0 {
 		var allowedIPs []net.IPNet
 		for _, ipStr := range request.AllowedIPs {
 			_, ipNet, err := net.ParseCIDR(ipStr)
diff --git a/wgtester/wgtester.go b/wgtester/wgtester.go
index b302fd4..4495e77 100644
--- a/wgtester/wgtester.go
+++ b/wgtester/wgtester.go
@@ -119,7 +119,13 @@ func (s *Server) handleConnections() {
 					// Just a timeout, keep going
 					continue
 				}
-				logger.Error(s.outputPrefix+"Error reading from UDP: %v", err)
+				// Check if we're shutting down and the connection was closed
+				select {
+				case <-s.shutdownCh:
+					return // Don't log error if we're shutting down
+				default:
+					logger.Error(s.outputPrefix+"Error reading from UDP: %v", err)
+				}
 				continue
 			}