From ec8ab870a45c3977eed501459044934eb8af2545 Mon Sep 17 00:00:00 2001
From: Gesa Stupperich <gesa@tailscale.com>
Date: Wed, 3 Jun 2026 12:44:19 +0100
Subject: [PATCH] tstest/integration/testcontrol: expire individual node keys

This adds testcontrol support for expiring individual node keys,
in order to enable test scenarios involving to key-expiry and
 extension.

Updates #19326

Signed-off-by: Gesa Stupperich <gesa@tailscale.com>
---
 tstest/integration/testcontrol/testcontrol.go | 35 +++++++++++++------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/tstest/integration/testcontrol/testcontrol.go b/tstest/integration/testcontrol/testcontrol.go
index 9cabd3ecd..4e0775003 100644
--- a/tstest/integration/testcontrol/testcontrol.go
+++ b/tstest/integration/testcontrol/testcontrol.go
@@ -910,7 +910,8 @@ func (s *Server) serveRegister(w http.ResponseWriter, r *http.Request, mkey key.
 	}
 
 	// If this is a followup request, wait until interactive followup URL visit complete.
-	if req.Followup != "" {
+	isFollowup := req.Followup != ""
+	if isFollowup {
 		followupURL, err := url.Parse(req.Followup)
 		if err != nil {
 			panic(err)
@@ -926,18 +927,22 @@ func (s *Server) serveRegister(w http.ResponseWriter, r *http.Request, mkey key.
 		// some follow-ups? For now all are successes.
 	}
 
-	// The in-memory list of nodes, users, and logins is keyed by
-	// the node key.  If the node key changes, update all the data stores
-	// to use the new node key.
+	// On a key rotation (OldNodeKey set and known to s.nodes), stage
+	// the new key as a candidate entry but keep the old key's entry
+	// alive so an in-flight map poll can still receive updates while
+	// the user completes the auth URL.
 	s.mu.Lock()
 	if _, oldNodeKeyOk := s.nodes[req.OldNodeKey]; oldNodeKeyOk {
 		if _, newNodeKeyOk := s.nodes[req.NodeKey]; !newNodeKeyOk {
-			s.nodes[req.OldNodeKey].Key = req.NodeKey
-			s.nodes[req.NodeKey] = s.nodes[req.OldNodeKey]
-
+			cloned := s.nodes[req.OldNodeKey].Clone()
+			cloned.Key = req.NodeKey
+			s.nodes[req.NodeKey] = cloned
 			s.users[req.NodeKey] = s.users[req.OldNodeKey]
 			s.logins[req.NodeKey] = s.logins[req.OldNodeKey]
-
+		}
+		if isFollowup {
+			// The user has completed the auth URL, the new key
+			// is now authoritative. Retire the old key's entry.
 			delete(s.nodes, req.OldNodeKey)
 			delete(s.users, req.OldNodeKey)
 			delete(s.logins, req.OldNodeKey)
@@ -997,11 +1002,19 @@ func (s *Server) serveRegister(w http.ResponseWriter, r *http.Request, mkey key.
 		}
 		s.nodes[nk] = node
 	}
+	// Consider a node key expired if allExpired is set or if the nodeKey has
+	// an expiry time in the past. This allows tests to set per-node KeyExpiry
+	// via UpdateNode to simulate an admin-triggered or time-based expiry.
+	nodeKeyExpired := s.allExpired
+	if !nodeKeyExpired && req.OldNodeKey.IsZero() {
+		if n, ok := s.nodes[nk]; ok && !n.KeyExpiry.IsZero() && n.KeyExpiry.Before(time.Now()) {
+			nodeKeyExpired = true
+		}
+	}
 	requireAuth := s.RequireAuth
-	if requireAuth && s.nodeKeyAuthed.Contains(nk) {
+	if requireAuth && s.nodeKeyAuthed.Contains(nk) && !nodeKeyExpired {
 		requireAuth = false
 	}
-	allExpired := s.allExpired
 	s.mu.Unlock()
 
 	authURL := ""
@@ -1014,7 +1027,7 @@ func (s *Server) serveRegister(w http.ResponseWriter, r *http.Request, mkey key.
 	res, err := s.encode(false, tailcfg.RegisterResponse{
 		User:              *user,
 		Login:             *login,
-		NodeKeyExpired:    allExpired,
+		NodeKeyExpired:    nodeKeyExpired,
 		MachineAuthorized: machineAuthorized,
 		AuthURL:           authURL,
 	})