From ee7948f3d55c9e399e528a87281ede9c94eb5935 Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Thu, 24 Jul 2025 20:45:17 -0700
Subject: [PATCH] Reconnect to newt

---
 common.go                  |  2 ++
 peermonitor/peermonitor.go | 18 ++++++++++++------
 wgtester/wgtester.go       | 20 ++++++++++++++++++++
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/common.go b/common.go
index 0274395..b3d1ce4 100644
--- a/common.go
+++ b/common.go
@@ -488,6 +488,8 @@ func ConfigurePeer(dev *device.Device, siteConfig SiteConfig, privateKey wgtypes
 		monitorAddress := strings.Split(siteConfig.ServerIP, "/")[0]
 		monitorPeer := fmt.Sprintf("%s:%d", monitorAddress, siteConfig.ServerPort+1) // +1 for the monitor port
 
+		logger.Debug("Setting up peer monitor for site %d at %s", siteConfig.SiteId, monitorPeer)
+
 		primaryRelay, err := resolveDomain(endpoint) // Using global endpoint variable
 		if err != nil {
 			logger.Warn("Failed to resolve primary relay endpoint: %v", err)
diff --git a/peermonitor/peermonitor.go b/peermonitor/peermonitor.go
index 684d767..df90de2 100644
--- a/peermonitor/peermonitor.go
+++ b/peermonitor/peermonitor.go
@@ -103,7 +103,7 @@ func (pm *PeerMonitor) AddPeer(siteID int, endpoint string, wgConfig *WireGuardC
 	// Check if we're already monitoring this peer
 	if _, exists := pm.monitors[siteID]; exists {
 		// Update the endpoint instead of creating a new monitor
-		pm.RemovePeer(siteID)
+		pm.removePeerUnlocked(siteID)
 	}
 
 	client, err := wgtester.NewClient(endpoint)
@@ -131,11 +131,9 @@ func (pm *PeerMonitor) AddPeer(siteID int, endpoint string, wgConfig *WireGuardC
 	return err
 }
 
-// RemovePeer stops monitoring a peer and removes it from the monitor
-func (pm *PeerMonitor) RemovePeer(siteID int) {
-	pm.mutex.Lock()
-	defer pm.mutex.Unlock()
-
+// removePeerUnlocked stops monitoring a peer and removes it from the monitor
+// This function assumes the mutex is already held by the caller
+func (pm *PeerMonitor) removePeerUnlocked(siteID int) {
 	client, exists := pm.monitors[siteID]
 	if !exists {
 		return
@@ -147,6 +145,14 @@ func (pm *PeerMonitor) RemovePeer(siteID int) {
 	delete(pm.configs, siteID)
 }
 
+// RemovePeer stops monitoring a peer and removes it from the monitor
+func (pm *PeerMonitor) RemovePeer(siteID int) {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	pm.removePeerUnlocked(siteID)
+}
+
 // Start begins monitoring all peers
 func (pm *PeerMonitor) Start() {
 	pm.mutex.Lock()
diff --git a/wgtester/wgtester.go b/wgtester/wgtester.go
index d63fc8d..28ffdba 100644
--- a/wgtester/wgtester.go
+++ b/wgtester/wgtester.go
@@ -30,6 +30,7 @@ type Client struct {
 	serverAddr     string
 	monitorRunning bool
 	monitorLock    sync.Mutex
+	connLock       sync.Mutex // Protects connection operations
 	shutdownCh     chan struct{}
 	packetInterval time.Duration
 	timeout        time.Duration
@@ -71,6 +72,10 @@ func (c *Client) SetMaxAttempts(attempts int) {
 // Close cleans up client resources
 func (c *Client) Close() {
 	c.StopMonitor()
+
+	c.connLock.Lock()
+	defer c.connLock.Unlock()
+
 	if c.conn != nil {
 		c.conn.Close()
 		c.conn = nil
@@ -79,6 +84,9 @@ func (c *Client) Close() {
 
 // ensureConnection makes sure we have an active UDP connection
 func (c *Client) ensureConnection() error {
+	c.connLock.Lock()
+	defer c.connLock.Unlock()
+
 	if c.conn != nil {
 		return nil
 	}
@@ -119,9 +127,19 @@ func (c *Client) TestConnection(ctx context.Context) (bool, time.Duration) {
 			timestamp := time.Now().UnixNano()
 			binary.BigEndian.PutUint64(packet[5:13], uint64(timestamp))
 
+			// Lock the connection for the entire send/receive operation
+			c.connLock.Lock()
+
+			// Check if connection is still valid after acquiring lock
+			if c.conn == nil {
+				c.connLock.Unlock()
+				return false, 0
+			}
+
 			logger.Debug("Attempting to send monitor packet to %s", c.serverAddr)
 			_, err := c.conn.Write(packet)
 			if err != nil {
+				c.connLock.Unlock()
 				logger.Info("Error sending packet: %v", err)
 				continue
 			}
@@ -133,6 +151,8 @@ func (c *Client) TestConnection(ctx context.Context) (bool, time.Duration) {
 			// Wait for response
 			responseBuffer := make([]byte, packetSize)
 			n, err := c.conn.Read(responseBuffer)
+			c.connLock.Unlock()
+
 			if err != nil {
 				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
 					// Timeout, try next attempt