poll: expire the write deadline only for a write in flight

An expired deadline resets the HTTP/2 stream with INTERNAL_ERROR, so
every batcher stop (delete, logout) looked like a server error.

hscontrol/poll.go

@@ -26,6 +26,10 @@ const (
 	keepAliveInterval = 50 * time.Second
 )
 
+// errMapSessionStopped is a write refused because the batcher stopped the
+// session; the poll loop ends as if it had seen cancelCh.
+var errMapSessionStopped = errors.New("map session stopped")
+
 type contextKey string
 
 const nodeNameContextKey = contextKey("nodeName")
@@ -40,6 +44,10 @@ type mapSession struct {
 	cancelCh       chan struct{}
 	cancelChClosed atomic.Bool
 
+	// writing is set while [mapSession.writeMap] runs, so
+	// [mapSession.stopFromBatcher] interrupts only a write that may be stuck.
+	writing atomic.Bool
+
 	keepAlive       time.Duration
 	keepAliveTicker *time.Ticker
 
@@ -97,9 +105,16 @@ func (m *mapSession) stopFromBatcher() {
 		close(m.cancelCh)
 
 		// A channel signal cannot interrupt a response write that is blocked on
-		// HTTP/2 flow control. Expire the stream's write deadline as well so a
+		// HTTP/2 flow control, so a client that stopped reading would keep the
-		// client that stopped reading cannot keep the map session, and therefore
+		// map session, and therefore server shutdown, alive indefinitely.
-		// server shutdown, alive indefinitely.
+		// Expiring the write deadline resets the stream with INTERNAL_ERROR,
+		// so only do it to a write in flight; an idle session sees cancelCh
+		// and ends the stream cleanly. writeMap checks cancelChClosed after
+		// raising writing, so one side always sees the other.
+		if !m.writing.Load() {
+			return
+		}
+
 		err := http.NewResponseController(m.w).SetWriteDeadline(time.Now())
 		if err != nil && !errors.Is(err, http.ErrNotSupported) {
 			m.log.Debug().Caller().Err(err).Msg("failed to interrupt map response write")
@@ -323,6 +338,12 @@ func (m *mapSession) serveLongPoll() {
 			}
 
 			err := m.writeMap(update)
+			if errors.Is(err, errMapSessionStopped) {
+				mapResponseEnded.WithLabelValues("cancelled").Inc()
+
+				return
+			}
+
 			if err != nil {
 				m.log.Error().Caller().Err(err).Msg("cannot write update to client")
 				return
@@ -333,6 +354,12 @@ func (m *mapSession) serveLongPoll() {
 
 		case <-m.keepAliveTicker.C:
 			err := m.writeMap(&keepAlive)
+			if errors.Is(err, errMapSessionStopped) {
+				mapResponseEnded.WithLabelValues("cancelled").Inc()
+
+				return
+			}
+
 			if err != nil {
 				m.log.Error().Caller().Err(err).Msg("cannot write keep alive")
 				return
@@ -353,6 +380,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 {
+	m.writing.Store(true)
+	defer m.writing.Store(false)
+
+	if m.cancelChClosed.Load() {
+		return errMapSessionStopped
+	}
+
 	startWrite := time.Now()
 
 	err := writeMapResponse(m.w, m.req.Compress, m.isStreaming(), msg)

hscontrol/poll_test.go

@@ -2,7 +2,9 @@ package hscontrol
 
 import (
 	"context"
+	"io"
 	"net/http"
+	"net/http/httptest"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -661,6 +663,70 @@ func TestDeletedNodeEndsLongPoll(t *testing.T) {
 	}
 }
 
+// TestDeletedNodeEndsStreamCleanly serves a long poll over HTTP/2, as the
+// Noise server does, and deletes the node while the session is idle. Nothing
+// is left to interrupt, so the stream must end with END_STREAM: an expired
+// write deadline resets it with INTERNAL_ERROR, which the client reports as a
+// server error.
+func TestDeletedNodeEndsStreamCleanly(t *testing.T) {
+	t.Parallel()
+
+	app := createTestApp(t)
+	app.cfg.Tuning.NodeMapSessionBufferedChanSize = 30
+
+	user := app.state.CreateUserForTest("poll-delete-clean-user")
+	pak, err := app.state.CreatePreAuthKey(user.TypedID(), false, false, nil, nil)
+	require.NoError(t, err)
+
+	nodeKey := key.NewNode()
+
+	_, err = app.handleRegister(context.Background(), tailcfg.RegisterRequest{
+		Auth:     &tailcfg.RegisterResponseAuth{AuthKey: pak.Key},
+		NodeKey:  nodeKey.Public(),
+		Hostinfo: &tailcfg.Hostinfo{Hostname: "poll-delete-clean-node"},
+		Expiry:   time.Now().Add(24 * time.Hour),
+	}, key.NewMachine().Public())
+	require.NoError(t, err)
+
+	nodeView, ok := app.state.GetNodeByNodeKey(nodeKey.Public())
+	require.True(t, ok)
+
+	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		app.newMapSession(r.Context(), tailcfg.MapRequest{
+			Stream:  true,
+			Version: tailcfg.CapabilityVersion(100),
+		}, w, nodeView.AsStruct()).serveLongPoll()
+	}))
+	srv.Config.Protocols = new(http.Protocols)
+	srv.Config.Protocols.SetUnencryptedHTTP2(true)
+	srv.Start()
+	t.Cleanup(srv.Close)
+
+	h2c := new(http.Protocols)
+	h2c.SetUnencryptedHTTP2(true)
+	client := &http.Client{Transport: &http.Transport{Protocols: h2c}}
+
+	req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, srv.URL, nil)
+	require.NoError(t, err)
+
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+
+	defer resp.Body.Close()
+
+	var size [4]byte
+
+	_, err = io.ReadFull(resp.Body, size[:])
+	require.NoError(t, err, "expected the initial map")
+
+	cs, err := app.state.DeleteNode(nodeView)
+	require.NoError(t, err)
+	app.Change(cs...)
+
+	_, err = io.Copy(io.Discard, resp.Body)
+	require.NoError(t, err, "stopping an idle session must end the stream, not reset it")
+}
+
 func TestDeletedNodeInterruptsBlockedWrite(t *testing.T) {
 	t.Parallel()