noise: re-delegate SSH check when the auth session is missing (#3306)

.github/workflows/test-integration.yaml

@@ -332,6 +332,7 @@ jobs:
           - TestSSHCheckModeUnapprovedTimeout
           - TestSSHCheckModeCheckPeriodCLI
           - TestSSHCheckModeAutoApprove
+          - TestSSHCheckModeSessionLossReDelegates
           - TestSSHCheckModeNegativeCLI
           - TestSSHLocalpart
           - TestTagsAuthKeyWithTagRequestDifferentTag

hscontrol/noise.go

@@ -610,6 +610,19 @@ func (ns *noiseServer) sshActionFollowUp(
 
 	auth, ok := ns.headscale.state.GetAuthCacheEntry(authID)
 	if !ok {
+		// The session is gone (expired, evicted, or lost on a control-plane
+		// restart). A bare error dead-ends the client: it keeps polling this
+		// now-defunct auth_id until the SSH connection times out. Re-delegate
+		// so a still-required check can complete instead.
+		if checkFound {
+			reqLog.Info().Caller().
+				Msg("SSH check auth session missing; re-delegating")
+
+			return ns.sshActionHoldAndDelegate(
+				reqLog, action, srcNodeID, dstNodeID,
+			)
+		}
+
 		return nil, NewHTTPError(
 			http.StatusBadRequest,
 			"Invalid auth_id",

hscontrol/noise_test.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strconv"
 	"testing"
 
@@ -396,3 +397,78 @@ func TestOverrideRemoteAddr(t *testing.T) {
 
 	assert.Equal(t, clientAddr, observed)
 }
+
+// TestSSHActionHoldAndDelegate_PersistsAuthSession guards the happy path: the
+// initial SSH-check poll returns a HoldAndDelegate URL carrying an auth_id, and
+// that auth session must remain in the cache for the follow-up poll to find.
+func TestSSHActionHoldAndDelegate_PersistsAuthSession(t *testing.T) {
+	t.Parallel()
+
+	app := createTestApp(t)
+	user := app.state.CreateUserForTest("ssh-persist-user")
+	src := putTestNodeInStore(t, app, user, "src-node")
+	dst := putTestNodeInStore(t, app, user, "dst-node")
+
+	ns := &noiseServer{headscale: app, machineKey: dst.MachineKey}
+
+	rec := httptest.NewRecorder()
+	ns.SSHActionHandler(rec, newSSHActionRequest(t, src.ID, dst.ID))
+	require.Equal(t, http.StatusOK, rec.Code, "initial poll body=%s", rec.Body.String())
+
+	var action tailcfg.SSHAction
+	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &action))
+	require.NotEmpty(t, action.HoldAndDelegate, "expected HoldAndDelegate, got %+v", action)
+
+	u, err := url.Parse(action.HoldAndDelegate)
+	require.NoError(t, err)
+
+	authIDStr := u.Query().Get("auth_id")
+	require.NotEmpty(t, authIDStr, "HoldAndDelegate URL missing auth_id: %s", action.HoldAndDelegate)
+
+	authID, err := types.AuthIDFromString(authIDStr)
+	require.NoError(t, err)
+
+	_, ok := app.state.GetAuthCacheEntry(authID)
+	require.True(t, ok, "auth session %s must persist after HoldAndDelegate", authID)
+}
+
+// TestSSHActionHandler_RejectsMissingSessionWithoutCheck verifies that without
+// an SSH check covering the pair, a follow-up poll for an unknown auth_id is a
+// genuinely bogus request and is rejected. The re-delegation behaviour for a
+// missing session (issue #3305, exercised end to end with a real client in the
+// servertest package) applies only when the pair is still subject to a check.
+func TestSSHActionHandler_RejectsMissingSessionWithoutCheck(t *testing.T) {
+	t.Parallel()
+
+	app := createTestApp(t)
+	user := app.state.CreateUserForTest("ssh-nocheck-user")
+	src := putTestNodeInStore(t, app, user, "src-node")
+	dst := putTestNodeInStore(t, app, user, "dst-node")
+
+	// No SSH-check policy is set, so the pair is not subject to a check.
+	_, checkFound := app.state.SSHCheckParams(src.ID, dst.ID)
+	require.False(t, checkFound, "test setup: pair must not be subject to a check")
+
+	ns := &noiseServer{headscale: app, machineKey: dst.MachineKey}
+
+	missing := types.MustAuthID()
+
+	rec := httptest.NewRecorder()
+	ns.SSHActionHandler(rec, newSSHActionFollowUpRequest(t, src.ID, dst.ID, missing))
+	require.Equal(t, http.StatusBadRequest, rec.Code,
+		"a bogus auth_id with no active check must be rejected, body=%s", rec.Body.String())
+}
+
+// newSSHActionFollowUpRequest is like newSSHActionRequest but carries the
+// auth_id query parameter that marks a follow-up poll.
+func newSSHActionFollowUpRequest(t *testing.T, src, dst types.NodeID, authID types.AuthID) *http.Request {
+	t.Helper()
+
+	req := newSSHActionRequest(t, src, dst)
+
+	q := req.URL.Query()
+	q.Set("auth_id", authID.String())
+	req.URL.RawQuery = q.Encode()
+
+	return req
+}

hscontrol/servertest/sshcheck_test.go

@@ -0,0 +1,127 @@
+package servertest_test
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/juanfont/headscale/hscontrol/servertest"
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/stretchr/testify/require"
+	"tailscale.com/tailcfg"
+)
+
+// TestSSHCheckReDelegatesWhenSessionMissing exercises the fix for
+// https://github.com/juanfont/headscale/issues/3305 with a real control
+// client. The dst node runs the SSH-check poll over its actual Noise
+// connection: it first obtains a genuine HoldAndDelegate auth_id, that auth
+// session is then dropped from the cache (as it would be on expiry, eviction,
+// or a control-plane restart), and the follow-up poll for the now-missing
+// session must re-delegate a fresh HoldAndDelegate rather than dead-ending the
+// client with an error it keeps retrying until the SSH connection times out.
+func TestSSHCheckReDelegatesWhenSessionMissing(t *testing.T) {
+	t.Parallel()
+
+	h := servertest.NewHarness(t, 2)
+
+	srcID := types.NodeID(h.Client(0).Netmap().SelfNode.ID()) //nolint:gosec
+	dstID := types.NodeID(h.Client(1).Netmap().SelfNode.ID()) //nolint:gosec
+
+	// Subject the same-user (src, dst) pair to an SSH check.
+	h.ChangePolicy(t, []byte(`{
+		"ssh": [{
+			"action": "check",
+			"src": ["harness-default@"],
+			"dst": ["autogroup:self"],
+			"users": ["autogroup:nonroot"]
+		}]
+	}`))
+
+	// Sanity: the policy must actually subject this pair to a check, otherwise
+	// the test would pass for the wrong reason.
+	_, checkFound := h.Server.State().SSHCheckParams(srcID, dstID)
+	require.True(t, checkFound, "test setup: (src, dst) must be subject to an SSH check")
+
+	// The dst node's first poll yields a real HoldAndDelegate carrying a real,
+	// cached auth_id — nothing is fabricated.
+	initial := pollSSHAction(t, h.Server.URL, h.Client(1), srcID, dstID, "")
+	require.NotEmpty(t, initial.HoldAndDelegate, "initial poll must hold and delegate, got %+v", initial)
+
+	authID := authIDFromHoldURL(t, initial.HoldAndDelegate)
+	_, ok := h.Server.State().GetAuthCacheEntry(authID)
+	require.True(t, ok, "the auth session must be cached after the initial poll")
+
+	// Drop the session, reproducing a natural loss (expiry/eviction/restart).
+	h.Server.State().DeleteAuthCacheEntryForTest(authID)
+	_, ok = h.Server.State().GetAuthCacheEntry(authID)
+	require.False(t, ok, "the auth session must be gone before the follow-up poll")
+
+	// The follow-up poll carries the real auth_id whose session is now missing.
+	// With an active check the server must re-delegate a fresh session.
+	followUp := pollSSHAction(t, h.Server.URL, h.Client(1), srcID, dstID, authID.String())
+	require.NotEmpty(t, followUp.HoldAndDelegate,
+		"a missing session under an active check must re-delegate, got %+v", followUp)
+
+	require.NotEqual(t, authID, authIDFromHoldURL(t, followUp.HoldAndDelegate),
+		"re-delegation must mint a fresh auth_id")
+}
+
+// pollSSHAction issues an /machine/ssh/action poll from the given node over its
+// real Noise connection, as tailscaled does. An empty authID is the initial
+// poll; a non-empty one is a follow-up.
+func pollSSHAction(
+	t *testing.T,
+	serverURL string,
+	node *servertest.TestClient,
+	srcID, dstID types.NodeID,
+	authID string,
+) tailcfg.SSHAction {
+	t.Helper()
+
+	actionURL := fmt.Sprintf("%s/machine/ssh/action/%d/to/%d", serverURL, srcID, dstID)
+	if authID != "" {
+		actionURL += "?auth_id=" + authID
+	}
+
+	// Noise requests are addressed with the https scheme; the control client
+	// routes them over the established Noise connection (mirroring how
+	// controlclient issues its own register/map calls).
+	actionURL = strings.Replace(actionURL, "http://", "https://", 1)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, actionURL, nil)
+	require.NoError(t, err)
+
+	resp, err := node.Direct().DoNoiseRequest(req)
+	require.NoError(t, err)
+
+	defer resp.Body.Close()
+
+	require.Equal(t, http.StatusOK, resp.StatusCode, "ssh action poll must return 200")
+
+	var action tailcfg.SSHAction
+	require.NoError(t, json.NewDecoder(resp.Body).Decode(&action))
+
+	return action
+}
+
+// authIDFromHoldURL extracts the auth_id query parameter from a HoldAndDelegate
+// URL.
+func authIDFromHoldURL(t *testing.T, holdURL string) types.AuthID {
+	t.Helper()
+
+	u, err := url.Parse(holdURL)
+	require.NoError(t, err)
+
+	authID, err := types.AuthIDFromString(u.Query().Get("auth_id"))
+	require.NoError(t, err, "HoldAndDelegate URL missing a valid auth_id: %s", holdURL)
+
+	return authID
+}

hscontrol/state/state.go

@@ -1450,6 +1450,14 @@ func (s *State) SetAuthCacheEntry(id types.AuthID, entry *types.AuthRequest) {
 	s.authCache.Add(id, entry)
 }
 
+// DeleteAuthCacheEntryForTest drops a pending auth request from the cache,
+// exposed for testing so a test can reproduce a session that was lost
+// (expired, evicted, or dropped on a control-plane restart) without faking an
+// auth_id.
+func (s *State) DeleteAuthCacheEntryForTest(id types.AuthID) {
+	s.authCache.Remove(id)
+}
+
 // SetLastSSHAuth records a successful SSH check authentication
 // for the given (src, dst) node pair.
 func (s *State) SetLastSSHAuth(src, dst types.NodeID) {

integration/control.go

@@ -23,6 +23,7 @@ type ControlServer interface {
 	GetHealthEndpoint() string
 	GetEndpoint() string
 	WaitForRunning() error
+	Restart() error
 	CreateUser(user string) (*v1.User, error)
 	CreateAuthKey(user uint64, reusable bool, ephemeral bool) (*v1.PreAuthKey, error)
 	CreateAuthKeyWithTags(user uint64, reusable bool, ephemeral bool, tags []string) (*v1.PreAuthKey, error)

integration/hsic/hsic.go

@@ -1533,6 +1533,20 @@ func (h *HeadscaleInContainer) Reload() error {
 	return nil
 }
 
+// Restart restarts the headscale container. The on-disk database and keys
+// persist across the restart, but all in-memory state is dropped — including
+// the bounded cache of pending authentication sessions. This reproduces a
+// control-plane restart, one of the real-world cases where a pending SSH-check
+// auth session is lost.
+func (h *HeadscaleInContainer) Restart() error {
+	err := h.pool.Client.RestartContainer(h.container.Container.ID, 30)
+	if err != nil {
+		return fmt.Errorf("restarting headscale container %s: %w", h.hostname, err)
+	}
+
+	return h.WaitForRunning()
+}
+
 // ApproveRoutes approves routes for a node.
 func (t *HeadscaleInContainer) ApproveRoutes(id uint64, routes []netip.Prefix) (*v1.Node, error) {
 	command := []string{

integration/ssh_test.go

@@ -644,6 +644,20 @@ func doSSHCheck(
 ) chan sshCheckResult {
 	t.Helper()
 
+	return doSSHCheckWithTimeout(t, client, peer, 60*time.Second)
+}
+
+// doSSHCheckWithTimeout is like doSSHCheck but lets the caller extend how long
+// the blocking SSH command may run, for flows that hold the check open longer
+// (e.g. while the control plane restarts).
+func doSSHCheckWithTimeout(
+	t *testing.T,
+	client TailscaleClient,
+	peer TailscaleClient,
+	timeout time.Duration,
+) chan sshCheckResult {
+	t.Helper()
+
 	peerFQDN, _ := peer.FQDN()
 
 	command := []string{
@@ -663,7 +677,7 @@ func doSSHCheck(
 	go func() {
 		stdout, stderr, err := client.Execute(
 			command,
-			dockertestutil.ExecuteCommandTimeout(60*time.Second),
+			dockertestutil.ExecuteCommandTimeout(timeout),
 		)
 		ch <- sshCheckResult{stdout, stderr, err}
 	}()
@@ -1248,6 +1262,90 @@ func TestSSHCheckModeAutoApprove(t *testing.T) {
 	}
 }
 
+// TestSSHCheckModeSessionLossReDelegates reproduces the failure in
+// https://github.com/juanfont/headscale/issues/3305 with a real client: an SSH
+// connection in check mode is pending a verdict when the control plane
+// restarts, which drops the in-memory auth cache so the session the client is
+// still polling for is gone. The client must recover — the server re-delegates
+// a fresh check rather than dead-ending the now-defunct auth_id — and once that
+// fresh check is approved the SSH connection completes.
+func TestSSHCheckModeSessionLossReDelegates(t *testing.T) {
+	IntegrationSkip(t)
+
+	scenario := sshScenario(t, sshCheckPolicy(), "ssh-sessionloss", 1)
+	defer scenario.ShutdownAssertNoPanics(t)
+
+	allClients, err := scenario.ListTailscaleClients()
+	requireNoErrListClients(t, err)
+
+	user1Clients, err := scenario.ListTailscaleClients("user1")
+	requireNoErrListClients(t, err)
+
+	headscale, err := scenario.Headscale()
+	require.NoError(t, err)
+
+	err = scenario.WaitForTailscaleSync()
+	requireNoErrSync(t, err)
+
+	_, err = scenario.ListTailscaleClientsFQDNs()
+	requireNoErrListFQDN(t, err)
+
+	for _, client := range user1Clients {
+		for _, peer := range allClients {
+			if client.Hostname() == peer.Hostname() {
+				continue
+			}
+
+			// Start SSH — blocks waiting for the check verdict while the
+			// pending auth session sits in the control plane's cache. Allow a
+			// generous window: the flow spans a full control-plane restart.
+			sshResult := doSSHCheckWithTimeout(t, client, peer, 120*time.Second)
+
+			firstAuthID := findSSHCheckAuthID(t, headscale)
+
+			// Restart the control plane: the in-memory auth cache is dropped
+			// (the on-disk database and keys persist), so the auth_id the
+			// client is still polling for no longer exists.
+			err := headscale.Restart()
+			require.NoError(t, err, "restarting headscale should succeed")
+
+			err = scenario.WaitForTailscaleSync()
+			requireNoErrSync(t, err)
+
+			// The client keeps polling the now-missing auth_id; with the fix the
+			// server re-delegates a fresh session instead of returning an error
+			// the client cannot recover from. A new auth_id only appears if the
+			// re-delegation happened.
+			secondAuthID := findNewSSHCheckAuthID(t, headscale, firstAuthID)
+			require.NotEqual(t, firstAuthID, secondAuthID,
+				"a lost session under an active check must re-delegate with a new auth_id")
+
+			// Approve the re-delegated session; the SSH connection must now
+			// complete instead of hanging until it times out.
+			_, err = headscale.Execute(
+				[]string{
+					"headscale", "auth", "approve",
+					"--auth-id", secondAuthID,
+				},
+			)
+			require.NoError(t, err)
+
+			select {
+			case result := <-sshResult:
+				require.NoError(t, result.err,
+					"SSH should succeed after re-delegation recovers the lost session")
+				require.Contains(
+					t,
+					peer.ContainerID(),
+					strings.ReplaceAll(result.stdout, "\n", ""),
+				)
+			case <-time.After(90 * time.Second):
+				t.Fatal("SSH did not complete after session-loss re-delegation")
+			}
+		}
+	}
+}
+
 // TestSSHCheckModeNegativeCLI verifies that `headscale auth reject`
 // properly denies an SSH check.
 func TestSSHCheckModeNegativeCLI(t *testing.T) {