state: forget a node's SSH check approvals when it is expired

hscontrol/state/ssh_check_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
 )
 
@@ -182,3 +183,134 @@ func TestSSHCheckAuthFollowsSourceOwner(t *testing.T) {
 		})
 	}
 }
+
+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

@@ -974,7 +974,7 @@ func (s *State) ListEphemeralNodes() views.Slice[types.NodeView] {
 // SetNodeExpiry updates the expiration time for a node.
 // If expiry is nil, the node's expiry is disabled (node will never expire).
 func (s *State) SetNodeExpiry(nodeID types.NodeID, expiry *time.Time) (types.NodeView, change.Change, error) {
-	var onlineChanged bool
+	var onlineChanged, expired bool
 
 	genBefore := s.polMan.NodesGeneration()
 
@@ -989,12 +989,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