From fe197f0a0bb03672b1de9143f3e983197a6afd39 Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Thu, 18 Dec 2025 15:04:20 -0500
Subject: [PATCH] Remove exit nodes from HPing if peers are removed

Former-commit-id: a4365988ebabd7d238760e86b2690e29853fe03a
---
 go.mod              |  2 ++
 olm/olm.go          | 11 +++++++++++
 websocket/client.go |  1 +
 3 files changed, 14 insertions(+)

diff --git a/go.mod b/go.mod
index 5e3ca07..baf9a13 100644
--- a/go.mod
+++ b/go.mod
@@ -75,3 +75,5 @@ require (
 	google.golang.org/protobuf v1.36.8 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+replace github.com/fosrl/newt => ../newt
diff --git a/olm/olm.go b/olm/olm.go
index 2d9b42a..a85b4c0 100644
--- a/olm/olm.go
+++ b/olm/olm.go
@@ -566,6 +566,14 @@ func StartTunnel(config TunnelConfig) {
 			return
 		}
 
+		// Remove any exit nodes associated with this peer from hole punching
+		if holePunchManager != nil {
+			removed := holePunchManager.RemoveExitNodesByPeer(removeData.SiteId)
+			if removed > 0 {
+				logger.Info("Removed %d exit nodes associated with peer %d from hole punch rotation", removed, removeData.SiteId)
+			}
+		}
+
 		// Remove successful
 		logger.Info("Successfully removed peer for site %d", removeData.SiteId)
 	})
@@ -798,10 +806,12 @@ func StartTunnel(config TunnelConfig) {
 			relayPort = 21820 // default relay port
 		}
 
+		siteId := handshakeData.SiteId
 		exitNode := holepunch.ExitNode{
 			Endpoint:  handshakeData.ExitNode.Endpoint,
 			RelayPort: relayPort,
 			PublicKey: handshakeData.ExitNode.PublicKey,
+			SiteIds:  []int{siteId},
 		}
 
 		added := holePunchManager.AddExitNode(exitNode)
@@ -894,6 +904,7 @@ func StartTunnel(config TunnelConfig) {
 				Endpoint:  node.Endpoint,
 				RelayPort: relayPort,
 				PublicKey: node.PublicKey,
+				SiteIds:   node.SiteIds,
 			}
 		}
 
diff --git a/websocket/client.go b/websocket/client.go
index faede03..1c5afaf 100644
--- a/websocket/client.go
+++ b/websocket/client.go
@@ -50,6 +50,7 @@ type ExitNode struct {
 	Endpoint  string `json:"endpoint"`
 	RelayPort uint16 `json:"relayPort"`
 	PublicKey string `json:"publicKey"`
+	SiteIds   []int  `json:"siteIds"`
 }
 
 type WSMessage struct {