poll, noise: tell a deleted node to re-authenticate

A bare 404 is indistinguishable from any other map-path error to a Tailscale
client: it retries forever, still logged in. Only a self node with a past
KeyExpiry reaches NeedsLogin. Also skip the reconnect grace wait, which a
deleted node can never satisfy.

Fixes #3410

hscontrol/noise.go

@@ -708,7 +708,32 @@ func (ns *noiseServer) PollNetMapHandler(
 
 	nv, err := ns.getAndValidateNode(mapRequest)
 	if err != nil {
+		// The node is gone, but the client does not know that. Tailscale
+		// clients treat every non-200 on the map path the same way and retry
+		// forever with loggedIn still set; only a self node whose KeyExpiry is
+		// in the past drives them to NeedsLogin. There is no MapResponse field
+		// that says "deleted", so reuse the expiry signal headscale already
+		// sends for expired nodes.
+		// See: https://github.com/juanfont/headscale/issues/3410
+		if errors.Is(err, errNodeNotInStore) && mapRequest.Stream {
+			expired := &tailcfg.MapResponse{
+				Node: &tailcfg.Node{
+					Key:       mapRequest.NodeKey,
+					KeyExpiry: time.Now().Add(-time.Hour).UTC(),
+					Expired:   true,
+				},
+			}
+
+			err = writeMapResponse(writer, mapRequest.Compress, true, expired)
+			if err != nil {
+				log.Error().Caller().Err(err).Msg("noise map handler: failed to write expired response for deleted node")
+			}
+
+			return
+		}
+
 		httpError(writer, err)
+
 		return
 	}
 
@@ -783,12 +808,17 @@ func (ns *noiseServer) RegistrationHandler(
 	}
 }
 
+// errNodeNotInStore distinguishes an unknown NodeKey from a NodeKey presented
+// by the wrong machine key. Both are answered with 404, but only the former
+// means the node is gone and its client should re-authenticate.
+var errNodeNotInStore = errors.New("node not found")
+
 // getAndValidateNode retrieves the node from the database using the NodeKey
 // and validates that it matches the MachineKey from the Noise session.
 func (ns *noiseServer) getAndValidateNode(mapRequest tailcfg.MapRequest) (types.NodeView, error) {
 	nv, ok := ns.headscale.state.GetNodeByNodeKey(mapRequest.NodeKey)
 	if !ok {
-		return types.NodeView{}, NewHTTPError(http.StatusNotFound, "node not found", nil)
+		return types.NodeView{}, NewHTTPError(http.StatusNotFound, "node not found", errNodeNotInStore)
 	}
 
 	// Validate that the MachineKey in the Noise session matches the one associated with the NodeKey.

hscontrol/noise_test.go

@@ -3,6 +3,7 @@ package hscontrol
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -11,13 +12,16 @@ import (
 	"net/url"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
+	"tailscale.com/util/zstdframe"
 )
 
 // newNoiseRouterWithBodyLimit builds a chi router with the same body-limit
@@ -537,3 +541,109 @@ func newSSHActionFollowUpRequest(t *testing.T, src, dst types.NodeID, authID typ
 
 	return req
 }
+
+// newMapRequest builds a streaming [tailcfg.MapRequest] POST for
+// /machine/map. Version is mandatory: [rejectUnsupported] runs before the
+// handler looks the node up, and a zero version is rejected with 400.
+func newMapRequest(t *testing.T, req tailcfg.MapRequest) *http.Request {
+	t.Helper()
+
+	body, err := json.Marshal(req)
+	require.NoError(t, err)
+
+	return httptest.NewRequestWithContext(context.Background(), http.MethodPost, "/machine/map", bytes.NewReader(body))
+}
+
+// decodeMapResponse reads a map response frame the way a Tailscale client
+// does: a little-endian length prefix followed by a body that is zstd-framed
+// whenever the request asked for compression.
+func decodeMapResponse(t *testing.T, compress string, body []byte) tailcfg.MapResponse {
+	t.Helper()
+
+	require.GreaterOrEqual(t, len(body), reservedResponseHeaderSize, "response too short to carry a length prefix")
+
+	size := binary.LittleEndian.Uint32(body[:reservedResponseHeaderSize])
+	payload := body[reservedResponseHeaderSize:]
+	require.Len(t, payload, int(size), "length prefix must match the body it precedes")
+
+	if compress == util.ZstdCompression {
+		decoded, err := zstdframe.AppendDecode(nil, payload)
+		require.NoError(t, err, "client decodes every frame as zstd when it asked for zstd")
+
+		payload = decoded
+	}
+
+	var resp tailcfg.MapResponse
+	require.NoError(t, json.Unmarshal(payload, &resp))
+
+	return resp
+}
+
+// TestPollNetMapHandler_DeletedNodeGetsExpiredSelf verifies that a streaming
+// map request for a node that no longer exists is answered with an expired
+// self node instead of a bare 404. A Tailscale client treats every non-200 on
+// the map path identically and retries forever with loggedIn still set; only a
+// self node whose KeyExpiry is in the past moves it to NeedsLogin.
+//
+// See: https://github.com/juanfont/headscale/issues/3410
+func TestPollNetMapHandler_DeletedNodeGetsExpiredSelf(t *testing.T) {
+	t.Parallel()
+
+	for _, compress := range []string{"", util.ZstdCompression} {
+		t.Run("compress="+compress, func(t *testing.T) {
+			t.Parallel()
+
+			app := createTestApp(t)
+			user := app.state.CreateUserForTest("deleted-node-user")
+			node := putTestNodeInStore(t, app, user, "deleted-node")
+
+			nodeView, ok := app.state.GetNodeByID(node.ID)
+			require.True(t, ok)
+
+			_, err := app.state.DeleteNode(nodeView)
+			require.NoError(t, err)
+
+			ns := &noiseServer{headscale: app, machineKey: node.MachineKey}
+
+			rec := httptest.NewRecorder()
+			ns.PollNetMapHandler(rec, newMapRequest(t, tailcfg.MapRequest{
+				Version:  tailcfg.CurrentCapabilityVersion,
+				NodeKey:  node.NodeKey,
+				Stream:   true,
+				Compress: compress,
+			}))
+
+			require.Equal(t, http.StatusOK, rec.Code, "body=%q", rec.Body.String())
+
+			resp := decodeMapResponse(t, compress, rec.Body.Bytes())
+			require.NotNil(t, resp.Node, "clients reject an initial map response without a node")
+			assert.Equal(t, node.NodeKey, resp.Node.Key)
+			assert.True(t, resp.Node.KeyExpiry.Before(time.Now()),
+				"a past KeyExpiry is what drives the client to NeedsLogin, got %v", resp.Node.KeyExpiry)
+		})
+	}
+}
+
+// TestPollNetMapHandler_ForeignMachineKeyStillRejected pins that the
+// expired-self response is limited to a genuinely unknown node. A known
+// NodeKey presented by the wrong machine key is an impostor, and answering it
+// with "your key expired" would wipe the real client's persisted node ID.
+func TestPollNetMapHandler_ForeignMachineKeyStillRejected(t *testing.T) {
+	t.Parallel()
+
+	app := createTestApp(t)
+	user := app.state.CreateUserForTest("impostor-user")
+	victim := putTestNodeInStore(t, app, user, "victim-node")
+	impostor := putTestNodeInStore(t, app, user, "impostor-node")
+
+	ns := &noiseServer{headscale: app, machineKey: impostor.MachineKey}
+
+	rec := httptest.NewRecorder()
+	ns.PollNetMapHandler(rec, newMapRequest(t, tailcfg.MapRequest{
+		Version: tailcfg.CurrentCapabilityVersion,
+		NodeKey: victim.NodeKey,
+		Stream:  true,
+	}))
+
+	assert.Equal(t, http.StatusNotFound, rec.Code, "body=%q", rec.Body.String())
+}

hscontrol/poll.go

@@ -4,12 +4,14 @@ import (
 	"context"
 	"encoding/binary"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/rand/v2"
 	"net/http"
 	"sync/atomic"
 	"time"
 
+	"github.com/juanfont/headscale/hscontrol/state"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/types/change"
 	"github.com/juanfont/headscale/hscontrol/util"
@@ -171,7 +173,12 @@ func (m *mapSession) serveLongPoll() {
 		// handler ran late is exactly such a session: if it kept its session
 		// acquired on this path, the surviving session's release could never
 		// take the node offline (the relogin flake).
-		if !stillConnected {
+		// A deleted node cannot reconnect, so waiting for it only delays the
+		// client's next map request, and with it the re-authentication signal
+		// it needs. See: https://github.com/juanfont/headscale/issues/3410
+		_, nodeExists := m.h.state.GetNodeByID(m.node.ID)
+
+		if !stillConnected && nodeExists {
 			// Wait up to 10 seconds for the node to reconnect.
 			// 10 seconds was arbitrary chosen as a reasonable time to reconnect.
 			ticker := time.NewTicker(time.Second)
@@ -191,7 +198,13 @@ func (m *mapSession) serveLongPoll() {
 		// sessions are harmless regardless of the order they run in.
 		disconnectChanges, err := m.h.state.Disconnect(m.node.ID, connectGen)
 		if err != nil {
-			m.log.Error().Caller().Err(err).Msg("failed to disconnect node")
+			// A node deleted mid-session is gone by the time its own session
+			// releases; that is the expected order, not a failure.
+			if errors.Is(err, state.ErrNodeNotFound) {
+				m.log.Debug().Caller().Err(err).Msg("node deleted before its session was released")
+			} else {
+				m.log.Error().Caller().Err(err).Msg("failed to disconnect node")
+			}
 		}
 
 		if len(disconnectChanges) == 0 {
@@ -326,35 +339,13 @@ func (m *mapSession) serveLongPoll() {
 // It also handles flushing the response if the [http.ResponseWriter]
 // implements [http.Flusher].
 func (m *mapSession) writeMap(msg *tailcfg.MapResponse) error {
-	jsonBody, err := json.Marshal(msg)
-	if err != nil {
-		return fmt.Errorf("marshalling map response: %w", err)
-	}
-
-	if m.req.Compress == util.ZstdCompression {
-		jsonBody = zstdframe.AppendEncode(nil, jsonBody, zstdframe.FastestCompression)
-	}
-
-	data := make([]byte, reservedResponseHeaderSize, reservedResponseHeaderSize+len(jsonBody))
-	//nolint:gosec // G115: JSON response size will not exceed uint32 max
-	binary.LittleEndian.PutUint32(data, uint32(len(jsonBody)))
-	data = append(data, jsonBody...)
-
 	startWrite := time.Now()
 
-	_, err = m.w.Write(data)
+	err := writeMapResponse(m.w, m.req.Compress, m.isStreaming(), msg)
 	if err != nil {
 		return err
 	}
 
-	if m.isStreaming() {
-		if f, ok := m.w.(http.Flusher); ok {
-			f.Flush()
-		} else {
-			m.log.Error().Caller().Msg("responseWriter does not implement http.Flusher, cannot flush")
-		}
-	}
-
 	m.log.Trace().
 		Caller().
 		Str(zf.Chan, fmt.Sprintf("%p", m.ch)).
@@ -366,6 +357,44 @@ func (m *mapSession) writeMap(msg *tailcfg.MapResponse) error {
 	return nil
 }
 
+// writeMapResponse writes a single map response frame: the JSON body,
+// zstd-framed when the client asked for compression, behind a little-endian
+// length prefix. Tailscale clients request zstd unconditionally and decode
+// every frame with it, so the compression step is not optional.
+//
+// It is shared with the deleted-node path in [noiseServer.PollNetMapHandler],
+// which has no [mapSession] to write through.
+func writeMapResponse(w http.ResponseWriter, compress string, flush bool, msg *tailcfg.MapResponse) error {
+	jsonBody, err := json.Marshal(msg)
+	if err != nil {
+		return fmt.Errorf("marshalling map response: %w", err)
+	}
+
+	if compress == util.ZstdCompression {
+		jsonBody = zstdframe.AppendEncode(nil, jsonBody, zstdframe.FastestCompression)
+	}
+
+	data := make([]byte, reservedResponseHeaderSize, reservedResponseHeaderSize+len(jsonBody))
+	//nolint:gosec // G115: JSON response size will not exceed uint32 max
+	binary.LittleEndian.PutUint32(data, uint32(len(jsonBody)))
+	data = append(data, jsonBody...)
+
+	_, err = w.Write(data)
+	if err != nil {
+		return err
+	}
+
+	if flush {
+		if f, ok := w.(http.Flusher); ok {
+			f.Flush()
+		} else {
+			log.Error().Caller().Msg("responseWriter does not implement http.Flusher, cannot flush")
+		}
+	}
+
+	return nil
+}
+
 var keepAlive = tailcfg.MapResponse{
 	KeepAlive: true,
 }

hscontrol/poll_test.go

@@ -386,3 +386,92 @@ func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t *
 		}
 	}, time.Second, 20*time.Millisecond, "after stale-send cleanup, the stale session should exit")
 }
+
+// TestDeletedNodeEndsLongPoll proves that deleting a node ends its long-poll
+// session, rather than leaving the goroutine streaming to a node that no longer
+// exists. An orphaned session blocks server shutdown on clientStreamsOpen and
+// keeps the client polling instead of re-authenticating.
+//
+// It also pins the teardown latency: a deleted node cannot reconnect, so the
+// session must not spend the reconnect grace period waiting for one.
+//
+// See: https://github.com/juanfont/headscale/issues/3410
+func TestDeletedNodeEndsLongPoll(t *testing.T) {
+	t.Parallel()
+
+	app := createTestApp(t)
+	user := app.state.CreateUserForTest("poll-delete-user")
+	createdNode := app.state.CreateRegisteredNodeForTest(user, "poll-delete-node")
+	require.NoError(t, app.state.UpdatePolicyManagerUsersForTest())
+
+	app.cfg.Tuning.NodeMapSessionBufferedChanSize = 1
+
+	// Reload so the NodeStore is populated from the database; the create
+	// helpers only write to the database.
+	app.mapBatcher.Close()
+	require.NoError(t, app.state.Close())
+
+	reloadedState, err := state.NewState(app.cfg)
+	require.NoError(t, err)
+
+	app.state = reloadedState
+	app.mapBatcher = mapper.NewBatcherAndMapper(app.cfg, app.state)
+	app.mapBatcher.Start()
+
+	t.Cleanup(func() {
+		app.mapBatcher.Close()
+		require.NoError(t, app.state.Close())
+	})
+
+	nodeView, ok := app.state.GetNodeByID(createdNode.ID)
+	require.True(t, ok)
+
+	node := nodeView.AsStruct()
+
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+
+	writer := newDelayedSuccessResponseWriter(0)
+	session := app.newMapSession(ctx, tailcfg.MapRequest{
+		Stream:  true,
+		Version: tailcfg.CapabilityVersion(100),
+	}, writer, node)
+
+	serveDone := make(chan struct{})
+
+	go func() {
+		session.serveLongPoll()
+		close(serveDone)
+	}()
+
+	select {
+	case <-writer.FirstWriteStarted():
+	case <-time.After(2 * time.Second):
+		t.Fatal("expected the initial map write to start")
+	}
+
+	c, err := app.state.DeleteNode(nodeView)
+	require.NoError(t, err)
+	app.Change(c)
+
+	// The reconnect grace period is 10s, so a generous bound here still fails
+	// if teardown waits for a node that can never come back.
+	select {
+	case <-serveDone:
+	case <-time.After(3 * time.Second):
+		t.Fatal("deleting a node must promptly end its long-poll session")
+	}
+
+	streamsClosed := make(chan struct{})
+
+	go func() {
+		app.clientStreamsOpen.Wait()
+		close(streamsClosed)
+	}()
+
+	select {
+	case <-streamsClosed:
+	case <-time.After(2 * time.Second):
+		t.Fatal("an orphaned session would block server shutdown on clientStreamsOpen")
+	}
+}