From 75b60f0bdb128fb3da8abfa5314fdfcb5a5bab7a Mon Sep 17 00:00:00 2001
From: Devon Bautista <17506592+synackd@users.noreply.github.com>
Date: Mon, 31 Aug 2026 10:16:18 -0600
Subject: [PATCH] fix: prevent stale peer removal in PeerRemovalQueue

Resolve a race condition in the PeerRemovalQueue where a removal job
could be processed after a node had reconnected and its public key
changed. The queue now stores the peer's public key at enqueue time and
validates it against the current key before invoking RemovePeer. If the
keys differ, the stale removal is skipped and a warning is logged. Added
a new removalJob struct to carry the key, updated the queue logic, and
introduced comprehensive unit tests
(TestPeerRemovalQueueStaleRemovalSkipped and
TestPeerRemovalQueueConcurrentEnqueue) to verify correct behavior. This
improves reliability of WireGuard peer management during rapid reconnect
scenarios.

Signed-off-by: Devon Bautista <17506592+synackd@users.noreply.github.com>
---
 cmd/cloud-init-server/peer_removal_queue.go   | 29 +++++++++++++++----
 .../peer_removal_queue_test.go                |  5 ++++
 pkg/wgtunnel/tunnels.go                       | 10 +++++++
 3 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/cmd/cloud-init-server/peer_removal_queue.go b/cmd/cloud-init-server/peer_removal_queue.go
index 71cf1fcf..1ff10518 100644
--- a/cmd/cloud-init-server/peer_removal_queue.go
+++ b/cmd/cloud-init-server/peer_removal_queue.go
@@ -13,11 +13,17 @@ const (
 
 type peerRemover interface {
 	RemovePeer(peerName string) error
+	GetPublicKey(peerName string) (string, bool)
 }
 
 type PeerRemovalQueue struct {
 	remover peerRemover
-	jobs    chan string
+	jobs    chan removalJob
+}
+
+type removalJob struct {
+	name string
+	key  string
 }
 
 func NewPeerRemovalQueue(remover peerRemover) *PeerRemovalQueue {
@@ -27,7 +33,7 @@ func NewPeerRemovalQueue(remover peerRemover) *PeerRemovalQueue {
 func newPeerRemovalQueue(remover peerRemover, workers int, buffer int) *PeerRemovalQueue {
 	queue := &PeerRemovalQueue{
 		remover: remover,
-		jobs:    make(chan string, buffer),
+		jobs:    make(chan removalJob, buffer),
 	}
 	for range workers {
 		go queue.work()
@@ -39,8 +45,11 @@ func (q *PeerRemovalQueue) TryEnqueue(peerName string) bool {
 	if q == nil || q.remover == nil {
 		return true
 	}
+	// Capture the current public key for the peer (if any).
+	key, _ := q.remover.GetPublicKey(peerName)
+	job := removalJob{name: peerName, key: key}
 	select {
-	case q.jobs <- peerName:
+	case q.jobs <- job:
 		return true
 	default:
 		return false
@@ -48,9 +57,17 @@ func (q *PeerRemovalQueue) TryEnqueue(peerName string) bool {
 }
 
 func (q *PeerRemovalQueue) work() {
-	for peerName := range q.jobs {
-		if err := q.remover.RemovePeer(peerName); err != nil {
-			log.Error().Err(err).Str("peer", peerName).Msg("failed to remove WireGuard peer")
+	for job := range q.jobs {
+		// Verify that the public key hasn't changed since enqueue time.
+		currentKey, ok := q.remover.GetPublicKey(job.name)
+		if ok && job.key != "" && currentKey != job.key {
+			// Stale removal – skip and log.
+			log.Warn().Str("peer", job.name).Str("queuedKey", job.key).Str("currentKey", currentKey).
+				Msg("stale removal job skipped")
+			continue
+		}
+		if err := q.remover.RemovePeer(job.name); err != nil {
+			log.Error().Err(err).Str("peer", job.name).Msg("failed to remove WireGuard peer")
 		}
 	}
 }
diff --git a/cmd/cloud-init-server/peer_removal_queue_test.go b/cmd/cloud-init-server/peer_removal_queue_test.go
index e5ab9559..f5083780 100644
--- a/cmd/cloud-init-server/peer_removal_queue_test.go
+++ b/cmd/cloud-init-server/peer_removal_queue_test.go
@@ -37,6 +37,11 @@ func (r *blockingPeerRemover) RemovePeer(peerName string) error {
 	return nil
 }
 
+// GetPublicKey implements peerRemover but returns no key for tests.
+func (r *blockingPeerRemover) GetPublicKey(peerName string) (string, bool) {
+	return "", false
+}
+
 type phoneHomeSMDClient struct {
 	smdclient.FakeSMDClient
 }
diff --git a/pkg/wgtunnel/tunnels.go b/pkg/wgtunnel/tunnels.go
index e1d42a5f..b23a05fa 100644
--- a/pkg/wgtunnel/tunnels.go
+++ b/pkg/wgtunnel/tunnels.go
@@ -178,6 +178,16 @@ func (m *InterfaceManager) RemovePeer(peerName string) error {
 	return nil
 }
 
+func (m *InterfaceManager) GetPublicKey(peerName string) (string, bool) {
+	m.peersMutex.RLock()
+	defer m.peersMutex.RUnlock()
+	peer, ok := m.peers[peerName]
+	if !ok {
+		return "", false
+	}
+	return peer.PublicKey, true
+}
+
 func (m *InterfaceManager) GetPeers() map[string]PeerConfig {
 	m.peersMutex.RLock()
 	defer m.peersMutex.RUnlock()