From 43fdabfe05d0c75dde0c57e752c4534beb56aec5 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 ++++
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/cmd/cloud-init-server/peer_removal_queue.go b/cmd/cloud-init-server/peer_removal_queue.go
index 077b9f18..39b26e8a 100644
--- a/cmd/cloud-init-server/peer_removal_queue.go
+++ b/cmd/cloud-init-server/peer_removal_queue.go
@@ -9,11 +9,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 {
@@ -23,7 +29,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()
@@ -35,8 +41,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
@@ -44,9 +53,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 91490191..5477e375 100644
--- a/cmd/cloud-init-server/peer_removal_queue_test.go
+++ b/cmd/cloud-init-server/peer_removal_queue_test.go
@@ -32,6 +32,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
 }