hscontrol: preserve ephemeral nodes with live sessions

Updates #3535, #3539

(cherry picked from commit 8798c9af83)

hscontrol/app.go

@@ -162,6 +162,20 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
 			return
 		}
 
+		// Schedule and Cancel run outside the session transition, so a
+		// disconnecting session can arm a timer after a reconnect already
+		// cancelled it. The session count is authoritative; a stale timer
+		// is dropped and the next disconnect arms a fresh one. Not Online():
+		// an expired node keeps polling while offline.
+		// ponytail: a Connect between this check and DeleteNode still loses
+		// the node; only reachable after a full inactivity timeout. Needs a
+		// delete-if-idle in State if that edge matters.
+		if node.ActiveSessions() > 0 {
+			log.Debug().Caller().EmbedObject(node).Msg("ephemeral node has a live session, skipping garbage collection")
+
+			return
+		}
+
 		policyChanged, err := app.state.DeleteNode(node)
 		if !policyChanged.IsEmpty() {
 			app.Change(policyChanged)

hscontrol/poll_test.go

@@ -342,6 +342,88 @@ func TestFailedReconnectDoesNotCancelEphemeralGC(t *testing.T) {
 		"failed reconnect must not cancel the ephemeral GC timer (issue #3382)")
 }
 
+// TestEphemeralGCDoesNotDeleteReconnectedNode replays, in order, the
+// interleaving of two sessions from
+// https://github.com/juanfont/headscale/issues/3535: the old session takes the
+// node offline, a reconnect brings it back online and cancels the GC before
+// the old session arms its timer. The armed timer is stale; firing it must not
+// delete a node that has a live session.
+func TestEphemeralGCDoesNotDeleteReconnectedNode(t *testing.T) {
+	t.Parallel()
+
+	app := createTestApp(t)
+	app.cfg.Node.Ephemeral.InactivityTimeout = 50 * time.Millisecond
+	app.StartEphemeralGCForTest(t)
+
+	user := app.state.CreateUserForTest("eph-gc-reconnect-user")
+	pak, err := app.state.CreatePreAuthKey(user.TypedID(), false, true, 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: "eph-gc-reconnect-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)
+	require.True(t, nodeView.IsEphemeral(), "test sanity: node must be ephemeral")
+
+	node := nodeView.AsStruct()
+	session := app.newMapSession(context.Background(), tailcfg.MapRequest{
+		Stream:  true,
+		Version: tailcfg.CapabilityVersion(100),
+	}, &recordingResponseWriter{}, node)
+
+	_, oldGen := app.state.Connect(node.ID)
+
+	// Old session's cleanup releases the last session.
+	offline, err := app.state.Disconnect(node.ID, oldGen)
+	require.NoError(t, err)
+	require.NotEmpty(t, offline, "test sanity: old session must take the node offline")
+
+	// Reconnect: what serveLongPoll does at Connect.
+	_, newGen := app.state.Connect(node.ID)
+	app.ephemeralGC.Cancel(node.ID)
+
+	// Old session's cleanup continues and arms the GC.
+	session.afterServeLongPoll()
+
+	online, ok := app.state.GetNodeByID(node.ID)
+	require.True(t, ok)
+	require.True(t, online.Online(), "test sanity: reconnected node must be online")
+
+	// The GC's delete lands only after a NodeStore batch flush, so watch
+	// well past the inactivity timeout.
+	assert.Never(t, func() bool {
+		_, ok := app.state.GetNodeByID(node.ID)
+
+		return !ok
+	}, 2*time.Second, 10*time.Millisecond,
+		"ephemeral GC must not delete a node with a live session (issue #3535)")
+
+	// Once the reconnected session ends, the node is idle and must still be
+	// collected.
+	_, err = app.state.Disconnect(node.ID, newGen)
+	require.NoError(t, err)
+	session.afterServeLongPoll()
+
+	assert.Eventually(t, func() bool {
+		_, ok := app.state.GetNodeByID(node.ID)
+
+		return !ok
+	}, 5*time.Second, 10*time.Millisecond,
+		"an idle ephemeral node must still be garbage collected")
+}
+
 // TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession
 // tests the scenario reported in
 // https://github.com/juanfont/headscale/issues/3129.