From 0b69e844f58b61e06f5d38c41d2041f0d002ca5c Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@dalby.cc>
Date: Wed, 2 Sep 2026 05:44:42 +0000
Subject: [PATCH] mapper: stop a deleted node's map session

Dropping the batcher entry left serveLongPoll streaming to a node that no
longer exists: Close ranges b.nodes and can no longer reach it, so shutdown
blocks and the client keeps polling instead of re-authenticating.

Updates #3410
---
 hscontrol/mapper/batcher.go           | 16 +++++++++++++++-
 hscontrol/mapper/batcher_unit_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/hscontrol/mapper/batcher.go b/hscontrol/mapper/batcher.go
index e21d86465..de3da87e0 100644
--- a/hscontrol/mapper/batcher.go
+++ b/hscontrol/mapper/batcher.go
@@ -596,12 +596,26 @@ func (b *Batcher) addToBatch(changes ...change.Change) {
 	// putting them in [tailcfg.MapResponse.PeersRemoved] (a different struct).
 	// Therefore, this cleanup only removes nodes that are truly being deleted,
 	// not nodes that are still connected but have lost visibility of certain peers.
+	// This loop now also terminates the node's map sessions, so a future
+	// [change.Change.PeersRemoved] producer that is not a deletion would kill a
+	// live node's long poll, not merely evict its batcher entry.
 	//
 	// See: https://github.com/juanfont/headscale/issues/2924
 	for _, ch := range changes {
 		for _, removedID := range ch.PeersRemoved {
-			if _, existed := b.nodes.LoadAndDelete(removedID); existed {
+			if nc, existed := b.nodes.LoadAndDelete(removedID); existed {
 				b.totalNodes.Add(-1)
+
+				// Tear the node's map sessions down. Dropping the entry alone
+				// leaves [mapSession.serveLongPoll] streaming to a node that no
+				// longer exists: [Batcher.Close] ranges b.nodes and can no
+				// longer reach it, so shutdown blocks on clientStreamsOpen, and
+				// the client keeps polling a node it should be re-authenticating.
+				// See: https://github.com/juanfont/headscale/issues/3410
+				if nc != nil {
+					nc.close()
+				}
+
 				log.Debug().
 					Uint64(zf.NodeID, removedID.Uint64()).
 					Msg("removed deleted node from batcher")
diff --git a/hscontrol/mapper/batcher_unit_test.go b/hscontrol/mapper/batcher_unit_test.go
index 08d515cdd..378247261 100644
--- a/hscontrol/mapper/batcher_unit_test.go
+++ b/hscontrol/mapper/batcher_unit_test.go
@@ -428,6 +428,33 @@ func TestMultiChannelClose_PreventsSendPanic(t *testing.T) {
 		"send after close should return errConnectionClosed, not panic")
 }
 
+func TestAddToBatch_NodeRemovedStopsSession(t *testing.T) {
+	lb := setupLightweightBatcher(t, 1, 1)
+	defer lb.cleanup()
+
+	mc, ok := lb.b.nodes.Load(1)
+	require.True(t, ok)
+
+	stopped := make(chan struct{})
+
+	mc.mutex.Lock()
+	mc.connections[0].stop = func() { close(stopped) }
+	mc.mutex.Unlock()
+
+	lb.b.AddWork(change.NodeRemoved(1))
+
+	// addToBatch runs synchronously, so the session must already be stopped.
+	select {
+	case <-stopped:
+	default:
+		t.Fatal("deleting a node must stop its map session, otherwise the long poll is orphaned")
+	}
+
+	_, stillTracked := lb.b.nodes.Load(1)
+	assert.False(t, stillTracked, "deleted node should no longer be tracked by the batcher")
+	assert.Equal(t, int64(0), lb.b.totalNodes.Load())
+}
+
 // ============================================================================
 // multiChannelNodeConn connection management Tests
 // ============================================================================