state: invalidate SSH check approvals on owner changes and expiry

Honor approval only while its source keeps the same owner and is live.

hscontrol/servertest/sshcheck_test.go

@@ -245,3 +245,42 @@ func TestSSHCheckFollowsReturnedHoldURL(t *testing.T) {
 		expandHoldURL(initial.HoldAndDelegate, srcID, dstID, "root"))
 	assert.True(t, followUp.Accept, "authenticated root login must be accepted, got %+v", followUp)
 }
+
+// TestSSHCheckApprovalDroppedAfterSourceTagConversion: an approval does not
+// outlive its source's owner. A check rule cannot name a tag source, so the
+// converted source is not let in at all.
+func TestSSHCheckApprovalDroppedAfterSourceTagConversion(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
+	h.ChangePolicy(t, []byte(`{
+"tagOwners":{"tag:ci":["harness-default@"],"tag:server":["harness-default@"]},
+"ssh":[{"action":"check","src":["harness-default@"],
+"dst":["tag:server"],"users":["autogroup:nonroot"],"checkPeriod":"1h"}]}`))
+
+	_, c, err := h.Server.State().SetNodeTags(dstID, []string{"tag:server"})
+	require.NoError(t, err)
+	h.Server.App.Change(c)
+
+	initial := pollSSHAction(t, h.Server.URL, h.Client(1), srcID, dstID, "")
+	require.NotEmpty(t, initial.HoldAndDelegate)
+	authID := authIDFromHoldURL(t, initial.HoldAndDelegate)
+	auth, ok := h.Server.State().GetAuthCacheEntry(authID)
+	require.True(t, ok)
+	auth.FinishAuth(types.AuthVerdict{})
+
+	approved := pollSSHAction(t, h.Server.URL, h.Client(1), srcID, dstID, authID.String())
+	require.True(t, approved.Accept)
+	cached := pollSSHAction(t, h.Server.URL, h.Client(1), srcID, dstID, "")
+	require.True(t, cached.Accept)
+
+	_, c, err = h.Server.State().SetNodeTags(srcID, []string{"tag:ci"})
+	require.NoError(t, err)
+	h.Server.App.Change(c)
+	_, checkFound := h.Server.State().SSHCheckParams(srcID, dstID, sshCheckLocalUser)
+	require.False(t, checkFound)
+	after := pollSSHAction(t, h.Server.URL, h.Client(1), srcID, dstID, "")
+	assert.False(t, after.Accept)
+	assert.Empty(t, after.HoldAndDelegate)
+}

hscontrol/state/ssh_check_test.go

@@ -1,6 +1,7 @@
 package state
 
 import (
+	"fmt"
 	"sync"
 	"testing"
 	"time"
@@ -8,16 +9,27 @@ import (
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
 )
 
-func newTestStateForSSHCheck() *State {
-	return &State{
-		sshCheckAuth: make(map[sshCheckPair]time.Time),
+func newTestStateForSSHCheck(t *testing.T) *State {
+	t.Helper()
+
+	s := newRetagTestState(t)
+
+	user := s.CreateUserForTest("ssh-user")
+	for i := range 10 {
+		node := s.CreateRegisteredNodeForTest(user, fmt.Sprintf("ssh-node-%d", i))
+		node.User = user
+		s.nodeStore.PutNode(*node)
 	}
+
+	return s
 }
 
 func TestSSHCheckAuth(t *testing.T) {
-	s := newTestStateForSSHCheck()
+	s := newTestStateForSSHCheck(t)
 
 	src := types.NodeID(1)
 	dst := types.NodeID(2)
@@ -46,7 +58,7 @@ func TestSSHCheckAuth(t *testing.T) {
 }
 
 func TestSSHCheckAuthClear(t *testing.T) {
-	s := newTestStateForSSHCheck()
+	s := newTestStateForSSHCheck(t)
 
 	s.SetLastSSHAuth(types.NodeID(1), types.NodeID(2))
 	s.SetLastSSHAuth(types.NodeID(1), types.NodeID(3))
@@ -68,13 +80,13 @@ func TestSSHCheckAuthClear(t *testing.T) {
 }
 
 func TestSSHCheckAuthConcurrent(t *testing.T) {
-	s := newTestStateForSSHCheck()
+	s := newTestStateForSSHCheck(t)
 
 	var wg sync.WaitGroup
 
 	for i := range 100 {
 		wg.Go(func() {
-			src := types.NodeID(uint64(i % 10))   //nolint:gosec
+			src := types.NodeID(uint64(i%10 + 1)) //nolint:gosec
 			dst := types.NodeID(uint64(i%5 + 10)) //nolint:gosec
 
 			s.SetLastSSHAuth(src, dst)
@@ -95,3 +107,210 @@ func TestSSHCheckAuthConcurrent(t *testing.T) {
 
 	wg.Wait()
 }
+
+func TestSSHCheckAuthFollowsSourceOwner(t *testing.T) {
+	for _, tt := range []struct {
+		name   string
+		tagged bool
+		kept   bool
+	}{
+		{"convert", false, false},
+		{"retag", true, false},
+		{"user_change", false, false},
+		{"delete", false, false},
+		{"destination_change", false, true},
+		{"rekey", false, true},
+		{"same_tags", true, true},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			s := newTestStateForSSHCheck(t)
+			_, err := s.SetPolicy([]byte(`{"tagOwners":{"tag:ci":["ssh-user@"],"tag:ops":["ssh-user@"]}}`))
+			require.NoError(t, err)
+
+			src, dst := types.NodeID(1), types.NodeID(2)
+			if tt.tagged {
+				_, _, err = s.SetNodeTags(src, []string{"tag:ci", "tag:ops"})
+				require.NoError(t, err)
+			}
+
+			s.SetLastSSHAuth(src, dst)
+			at, ok := s.GetLastSSHAuth(src, dst)
+			require.True(t, ok)
+
+			switch tt.name {
+			case "convert", "destination_change":
+				id := src
+				if tt.name == "destination_change" {
+					id = dst
+				}
+
+				_, _, err = s.SetNodeTags(id, []string{"tag:ci"})
+				require.NoError(t, err)
+			case "retag":
+				_, _, err = s.SetNodeTags(src, []string{"tag:ops"})
+				require.NoError(t, err)
+			case "user_change":
+				other := s.CreateUserForTest("other-user")
+				_, ok = s.nodeStore.UpdateNode(src, func(n *types.Node) {
+					n.UserID = &other.ID
+					n.User = other
+				})
+				require.True(t, ok)
+			case "delete":
+				node, ok := s.GetNodeByID(src)
+				require.True(t, ok)
+
+				_, err = s.DeleteNode(node)
+				require.NoError(t, err)
+			case "rekey":
+				_, ok = s.nodeStore.UpdateNode(src, func(n *types.Node) { n.NodeKey = key.NewNode().Public() })
+				require.True(t, ok)
+			case "same_tags":
+				_, ok = s.nodeStore.UpdateNode(src, func(n *types.Node) {
+					n.Tags = []string{"tag:ops", "tag:ci"}
+				})
+				require.True(t, ok)
+			}
+
+			last, ok := s.GetLastSSHAuth(src, dst)
+			assert.Equal(t, tt.kept, ok)
+
+			if tt.kept {
+				assert.Equal(t, at, last)
+			} else {
+				assert.True(t, last.IsZero())
+			}
+		})
+	}
+}
+
+func TestSSHCheckAuthExpiresWithSource(t *testing.T) {
+	s := newTestStateForSSHCheck(t)
+	src, dst, other := types.NodeID(1), types.NodeID(2), types.NodeID(3)
+	s.SetLastSSHAuth(src, dst)
+	s.SetLastSSHAuth(src, other)
+	s.SetLastSSHAuth(dst, src)
+	s.SetLastSSHAuth(dst, other)
+	node, ok := s.GetNodeByID(src)
+	require.True(t, ok)
+
+	expiry := time.Now()
+	_, _, err := s.SetNodeExpiry(src, &expiry)
+	require.NoError(t, err)
+
+	for _, destination := range []types.NodeID{dst, other} {
+		_, ok = s.GetLastSSHAuth(src, destination)
+		assert.False(t, ok)
+	}
+
+	for _, destination := range []types.NodeID{src, other} {
+		_, ok = s.GetLastSSHAuth(dst, destination)
+		assert.True(t, ok)
+	}
+
+	userID := types.UserID(node.UserID().Get())
+	pak, err := s.CreatePreAuthKey(&userID, true, false, nil, nil)
+	require.NoError(t, err)
+	reauthenticated, _, err := s.HandleNodeFromPreAuthKey(tailcfg.RegisterRequest{
+		Auth:     &tailcfg.RegisterResponseAuth{AuthKey: pak.Key},
+		NodeKey:  key.NewNode().Public(),
+		Hostinfo: &tailcfg.Hostinfo{Hostname: node.Hostname()},
+		Expiry:   time.Now().Add(24 * time.Hour),
+	}, node.MachineKey())
+	require.NoError(t, err)
+	require.Equal(t, src, reauthenticated.ID())
+	require.False(t, reauthenticated.IsExpired())
+
+	for _, destination := range []types.NodeID{dst, other} {
+		_, ok = s.GetLastSSHAuth(src, destination)
+		assert.False(t, ok, "pre-auth key relogin must require a new SSH check")
+	}
+}
+
+func TestSSHCheckAuthKeepsApprovalOnExpiryExtension(t *testing.T) {
+	future := time.Now().Add(24 * time.Hour)
+
+	zero := time.Time{}
+	for name, expiry := range map[string]*time.Time{
+		"future": &future, "disabled": nil, "zero": &zero,
+	} {
+		t.Run(name, func(t *testing.T) {
+			s := newTestStateForSSHCheck(t)
+			src, dst := types.NodeID(1), types.NodeID(2)
+			s.SetLastSSHAuth(src, dst)
+			before, ok := s.GetLastSSHAuth(src, dst)
+			require.True(t, ok)
+
+			_, _, err := s.SetNodeExpiry(src, expiry)
+			require.NoError(t, err)
+
+			after, ok := s.GetLastSSHAuth(src, dst)
+			require.True(t, ok)
+			assert.Equal(t, before, after)
+		})
+	}
+}
+
+// TestSSHCheckAuthExpiresWithinBatch: a relogin that shares the expiry's
+// NodeStore batch renews the node before SetNodeExpiry sees it, so the expiry
+// must forget approvals from what its own write saw.
+func TestSSHCheckAuthExpiresWithinBatch(t *testing.T) {
+	s := newTestStateForSSHCheck(t)
+	src, dst := types.NodeID(1), types.NodeID(2)
+	s.SetLastSSHAuth(src, dst)
+	node, ok := s.GetNodeByID(src)
+	require.True(t, ok)
+
+	userID := types.UserID(node.UserID().Get())
+	pak, err := s.CreatePreAuthKey(&userID, true, false, nil, nil)
+	require.NoError(t, err)
+
+	// Hold the writer so the expiry and the relogin land in one batch, in
+	// this order.
+	entered, release := make(chan struct{}), make(chan struct{})
+
+	releaseOnce := sync.OnceFunc(func() { close(release) })
+	defer releaseOnce()
+
+	go s.nodeStore.UpdateNode(dst, func(*types.Node) {
+		close(entered)
+		<-release
+	})
+
+	<-entered
+
+	expiry := time.Now()
+	expiryDone := make(chan error, 1)
+
+	go func() {
+		_, _, err := s.SetNodeExpiry(src, &expiry)
+		expiryDone <- err
+	}()
+
+	waitParkedOnWriteQueue(t, "SetNodeExpiry")
+
+	regDone := make(chan error, 1)
+
+	go func() {
+		_, _, err := s.HandleNodeFromPreAuthKey(tailcfg.RegisterRequest{
+			Auth:     &tailcfg.RegisterResponseAuth{AuthKey: pak.Key},
+			NodeKey:  key.NewNode().Public(),
+			Hostinfo: &tailcfg.Hostinfo{Hostname: node.Hostname()},
+			Expiry:   time.Now().Add(24 * time.Hour),
+		}, node.MachineKey())
+		regDone <- err
+	}()
+
+	waitParkedOnWriteQueue(t, "reregisterNodeWithPAK")
+
+	releaseOnce()
+	require.NoError(t, <-expiryDone)
+	require.NoError(t, <-regDone)
+
+	renewed, ok := s.GetNodeByID(src)
+	require.True(t, ok)
+	require.False(t, renewed.IsExpired(), "precondition: the relogin renewed the node")
+
+	_, ok = s.GetLastSSHAuth(src, dst)
+	assert.False(t, ok, "expiry must forget the approval")
+}

hscontrol/state/state.go

@@ -162,6 +162,11 @@ type sshCheckPair struct {
 	Dst types.NodeID
 }
 
+type sshCheckAuthRecord struct {
+	at  time.Time
+	src types.NodeView
+}
+
 // State manages Headscale's core state, coordinating between database, policy management,
 // IP allocation, and DERP routing. All methods are thread-safe.
 //
@@ -216,7 +221,7 @@ type State struct {
 	//
 	// Ref: https://github.com/tailscale/tailscale/issues/10480
 	// Ref: https://github.com/tailscale/tailscale/issues/7125
-	sshCheckAuth map[sshCheckPair]time.Time
+	sshCheckAuth map[sshCheckPair]sshCheckAuthRecord
 	sshCheckMu   sync.RWMutex
 
 	// persistMu serialises node-row persistence and deletion so the database
@@ -391,7 +396,7 @@ func NewState(cfg *types.Config) (*State, error) {
 		nodeStore:           nodeStore,
 		pings:               newPingTracker(),
 
-		sshCheckAuth:  make(map[sshCheckPair]time.Time),
+		sshCheckAuth:  make(map[sshCheckPair]sshCheckAuthRecord),
 		registerLocks: xsync.NewMap[key.MachinePublic, *registerLock](),
 	}
 
@@ -1083,7 +1088,7 @@ func (s *State) SetNodeExpiry(nodeID types.NodeID, expiry *time.Time) (types.Nod
 		expiry = &utc
 	}
 
-	var onlineChanged bool
+	var onlineChanged, expired bool
 
 	genBefore := s.polMan.NodesGeneration()
 
@@ -1098,12 +1103,25 @@ func (s *State) SetNodeExpiry(nodeID types.NodeID, expiry *time.Time) (types.Nod
 		// recover the client, but an expired key is not online.
 		node.IsOnline = new(node.ShouldBeOnline())
 		onlineChanged = wasOnline != node.Online()
+		// A relogin later in the batch can renew the node before n is read.
+		expired = node.IsExpired()
 	})
 
 	if !ok {
 		return types.NodeView{}, change.Change{}, fmt.Errorf("%w: %d", ErrNodeNotInNodeStore, nodeID)
 	}
 
+	if expired {
+		// Expiry revokes the human check even if a pre-auth key renews the node.
+		s.sshCheckMu.Lock()
+		for pair := range s.sshCheckAuth {
+			if pair.Src == nodeID {
+				delete(s.sshCheckAuth, pair)
+			}
+		}
+		s.sshCheckMu.Unlock()
+	}
+
 	// The online flip already re-elected primaries in the NodeStore, so
 	// peers need it even when the database write below fails.
 	var recompute change.Change
@@ -2107,18 +2125,33 @@ func (s *State) SetLastSSHAuth(src, dst types.NodeID) {
 	s.sshCheckMu.Lock()
 	defer s.sshCheckMu.Unlock()
 
-	s.sshCheckAuth[sshCheckPair{Src: src, Dst: dst}] = time.Now()
+	node, ok := s.nodeStore.GetNode(src)
+	if !ok {
+		return
+	}
+
+	s.sshCheckAuth[sshCheckPair{Src: src, Dst: dst}] = sshCheckAuthRecord{
+		at: time.Now(), src: node,
+	}
 }
 
 // GetLastSSHAuth returns when src last authenticated for SSH check
-// to dst.
+// to dst, while the source still has the owner that authenticated.
 func (s *State) GetLastSSHAuth(src, dst types.NodeID) (time.Time, bool) {
 	s.sshCheckMu.RLock()
 	defer s.sshCheckMu.RUnlock()
 
-	t, ok := s.sshCheckAuth[sshCheckPair{Src: src, Dst: dst}]
+	record, ok := s.sshCheckAuth[sshCheckPair{Src: src, Dst: dst}]
+	if !ok {
+		return time.Time{}, false
+	}
 
-	return t, ok
+	node, exists := s.nodeStore.GetNode(src)
+	if !exists || !sameOwner(record.src, node) {
+		return time.Time{}, false
+	}
+
+	return record.at, true
 }
 
 // ClearSSHCheckAuth clears all recorded SSH check auth times.
@@ -2127,7 +2160,7 @@ func (s *State) ClearSSHCheckAuth() {
 	s.sshCheckMu.Lock()
 	defer s.sshCheckMu.Unlock()
 
-	s.sshCheckAuth = make(map[sshCheckPair]time.Time)
+	s.sshCheckAuth = make(map[sshCheckPair]sshCheckAuthRecord)
 }
 
 // preserveNetInfo preserves NetInfo from an existing node for faster DERP connectivity.
@@ -2178,6 +2211,20 @@ type authNodeUpdateParams struct {
 	IsConvertFromTag bool
 }
 
+// sameOwner reports whether a and b have one owner: the same tag set for
+// tagged nodes, the same user otherwise.
+func sameOwner(a, b types.NodeView) bool {
+	if !a.Valid() || !b.Valid() || a.IsTagged() != b.IsTagged() {
+		return false
+	}
+
+	if a.IsTagged() {
+		return views.SliceEqualAnyOrder(a.Tags(), b.Tags())
+	}
+
+	return a.UserID().Valid() && b.UserID().Valid() && a.UserID().Get() == b.UserID().Get()
+}
+
 // applyAuthNodeUpdate applies common update logic for re-authenticating or converting
 // an existing node. It updates the node in [NodeStore], processes RequestTags, and
 // persists changes to the database. The bool reports whether peers read the
@@ -2315,6 +2362,9 @@ func (s *State) applyAuthNodeUpdate(params authNodeUpdateParams) (types.NodeView
 			exp := time.Now().Add(s.cfg.Node.Expiry)
 			node.Expiry = &exp
 		}
+
+		// Route approvals stay on an owner change, as on Tailscale SaaS;
+		// autoApprovers only add.
 	})
 
 	if !ok {
@@ -3361,6 +3411,9 @@ func (s *State) reregisterNodeWithPAK(
 		default:
 			node.Expiry = nil
 		}
+
+		// Route approvals stay on an owner change, as on Tailscale SaaS;
+		// autoApprovers only add.
 	})
 
 	// A nil prior means the writer never saw the node (deleted, or the store