servertest: add relogin and auth key helpers

Relogin builds a new Direct on the same machine key and persisted
node key, as tailscaled does on up --auth-key; force rotates the key.

(cherry picked from commit 1716eece74, adapted: adds main's startPollLoop and resetNetmapState)

hscontrol/servertest/client.go

@@ -15,6 +15,7 @@ import (
 	"tailscale.com/net/tsdial"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
+	"tailscale.com/types/logger"
 	"tailscale.com/types/netmap"
 	"tailscale.com/types/persist"
 	"tailscale.com/types/views"
@@ -35,6 +36,12 @@ type TestClient struct {
 	authKey string
 	user    *types.User
 
+	// machineKey and discoKey survive [TestClient.Relogin], as they survive
+	// a tailscaled login.
+	machineKey key.MachinePrivate
+	discoKey   key.DiscoPublic
+	logf       logger.Logf
+
 	// Connection lifecycle.
 	pollCtx    context.Context //nolint:containedctx // test-only; context stored for cancel control
 	pollCancel context.CancelFunc
@@ -66,6 +73,7 @@ type clientConfig struct {
 	hostname     string
 	tags         []string
 	user         *types.User
+	authKey      string
 	deltaUpdates bool
 }
 
@@ -92,6 +100,11 @@ func WithTags(tags ...string) ClientOption {
 	return func(c *clientConfig) { c.tags = tags }
 }
 
+// WithAuthKey registers with key instead of a key minted for the client.
+func WithAuthKey(key string) ClientOption {
+	return func(c *clientConfig) { c.authKey = key }
+}
+
 // WithUser sets the user for the client. If not set, the harness
 // creates a default user.
 func WithUser(user *types.User) ClientOption {
@@ -123,6 +136,8 @@ func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOpt
 	var authKey string
 
 	switch {
+	case cc.authKey != "":
+		authKey = cc.authKey
 	case cc.ephemeral:
 		authKey = server.CreateEphemeralPreAuthKey(tb, uid)
 	case len(cc.tags) > 0:
@@ -141,40 +156,31 @@ func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOpt
 	// so that no real TCP sockets are used.
 	dialer.SetSystemDialerForTest(server.MemNet().Dial)
 
-	machineKey := key.NewMachine()
+	tc := &TestClient{
+		Name:       name,
+		server:     server,
+		authKey:    authKey,
+		user:       user,
+		machineKey: key.NewMachine(),
+		discoKey:   key.NewDisco().Public(),
+		logf:       tb.Logf,
+		updates:    make(chan *netmap.NetworkMap, 64),
+		bus:        bus,
+		dialer:     dialer,
+		tracker:    tracker,
 
-	direct, err := controlclient.NewDirect(controlclient.Options{
+		deltaUpdates: cc.deltaUpdates,
-		Persist:              persist.Persist{},
+	}
-		GetMachinePrivateKey: func() (key.MachinePrivate, error) { return machineKey, nil },
+
-		ServerURL:            server.URL,
+	direct, err := tc.newDirect(persist.Persist{}, authKey, &tailcfg.Hostinfo{
-		AuthKey:              authKey,
+		BackendLogID: "servertest-" + name,
-		Hostinfo: &tailcfg.Hostinfo{
+		Hostname:     cc.hostname,
-			BackendLogID: "servertest-" + name,
-			Hostname:     cc.hostname,
-		},
-		DiscoPublicKey: key.NewDisco().Public(),
-		Logf:           tb.Logf,
-		HealthTracker:  tracker,
-		Dialer:         dialer,
-		Bus:            bus,
 	})
 	if err != nil {
 		tb.Fatalf("servertest: NewDirect(%s): %v", name, err)
 	}
 
-	tc := &TestClient{
+	tc.direct = direct
-		Name:    name,
-		server:  server,
-		direct:  direct,
-		authKey: authKey,
-		user:    user,
-		updates: make(chan *netmap.NetworkMap, 64),
-		bus:     bus,
-		dialer:  dialer,
-		tracker: tracker,
-
-		deltaUpdates: cc.deltaUpdates,
-	}
 
 	tb.Cleanup(func() {
 		tc.cleanup()
@@ -189,6 +195,26 @@ func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOpt
 	return tc
 }
 
+// newDirect builds a [controlclient.Direct] for the client's machine.
+func (c *TestClient) newDirect(
+	p persist.Persist,
+	authKey string,
+	hi *tailcfg.Hostinfo,
+) (*controlclient.Direct, error) {
+	return controlclient.NewDirect(controlclient.Options{
+		Persist:              p,
+		GetMachinePrivateKey: func() (key.MachinePrivate, error) { return c.machineKey, nil },
+		ServerURL:            c.server.URL,
+		AuthKey:              authKey,
+		Hostinfo:             hi,
+		DiscoPublicKey:       c.discoKey,
+		Logf:                 c.logf,
+		HealthTracker:        c.tracker,
+		Dialer:               c.dialer,
+		Bus:                  c.bus,
+	})
+}
+
 // register performs the initial [controlclient.Direct.TryLogin] to register the client.
 func (c *TestClient) register(tb testing.TB) {
 	tb.Helper()
@@ -210,6 +236,13 @@ func (c *TestClient) register(tb testing.TB) {
 func (c *TestClient) startPoll(tb testing.TB) {
 	tb.Helper()
 
+	c.startPollLoop()
+}
+
+// startPollLoop creates a fresh poll context and launches the background
+// [controlclient.Direct.PollNetMap] goroutine, which blocks until the
+// context is cancelled or the server closes the connection.
+func (c *TestClient) startPollLoop() {
 	c.pollCtx, c.pollCancel = context.WithCancel(context.Background())
 	c.pollDone = make(chan struct{})
 
@@ -225,6 +258,23 @@ func (c *TestClient) startPoll(tb testing.TB) {
 	}()
 }
 
+// resetNetmapState clears the cached netmap and drains any pending
+// updates from a previous session so that convergence waits observe
+// only the new session's maps.
+func (c *TestClient) resetNetmapState() {
+	c.mu.Lock()
+	c.netmap = nil
+	c.mu.Unlock()
+
+	for {
+		select {
+		case <-c.updates:
+		default:
+			return
+		}
+	}
+}
+
 // UpdateFullNetmap implements [controlclient.NetmapUpdater].
 // Called by [controlclient.Direct] when a new [netmap.NetworkMap] is received.
 func (c *TestClient) UpdateFullNetmap(nm *netmap.NetworkMap) {
@@ -430,6 +480,59 @@ func (c *TestClient) ReloginAndPoll(ctx context.Context) error {
 	return nil
 }
 
+// Relogin logs the client in again with authKey, as tailscaled does for
+// `tailscale up --auth-key` (force: `--force-reauth`). Like
+// LocalBackend.Start it builds a new [controlclient.Direct] that keeps the
+// machine key and the persisted node key, so the first request presents the
+// current node key unless force rotates it. hi replaces the Hostinfo, so
+// routes can be restated. The poll restarts once the login succeeds.
+//
+// Safe to call from non-test goroutines.
+func (c *TestClient) Relogin(ctx context.Context, authKey string, force bool, hi *tailcfg.Hostinfo) error {
+	if c.pollCancel != nil {
+		c.pollCancel()
+
+		select {
+		case <-c.pollDone:
+		case <-ctx.Done():
+			return fmt.Errorf("servertest: Relogin(%s): poll did not exit: %w", c.Name, ctx.Err())
+		}
+	}
+
+	var p persist.Persist
+	if v := c.direct.GetPersist(); v.Valid() {
+		p = *v.AsStruct()
+	}
+
+	direct, err := c.newDirect(p, authKey, hi)
+	if err != nil {
+		return fmt.Errorf("servertest: Relogin(%s): NewDirect: %w", c.Name, err)
+	}
+
+	c.direct.Close()
+	c.direct = direct
+	c.authKey = authKey
+
+	flags := controlclient.LoginDefault
+	if force {
+		flags = controlclient.LoginInteractive
+	}
+
+	url, err := direct.TryLogin(ctx, flags)
+	if err != nil {
+		return fmt.Errorf("servertest: Relogin(%s): %w", c.Name, err)
+	}
+
+	if url != "" {
+		return fmt.Errorf("servertest: Relogin(%s): unexpected auth URL %q", c.Name, url) //nolint:err113
+	}
+
+	c.resetNetmapState()
+	c.startPollLoop()
+
+	return nil
+}
+
 // RestartPoll tears down the current long-poll session and immediately
 // starts a new one without re-registering, the way tailscaled restarts
 // its map poll on state-machine transitions (pause/unpause around
@@ -656,6 +759,12 @@ func (c *TestClient) WaitForCondition(tb testing.TB, desc string, timeout time.D
 	}
 }
 
+// MachineKey is the client's machine key, which identifies its nodes across
+// logins.
+func (c *TestClient) MachineKey() key.MachinePublic {
+	return c.machineKey.Public()
+}
+
 // Direct returns the underlying [controlclient.Direct] for
 // advanced operations like [controlclient.Direct.SetHostinfo] or SendUpdate.
 func (c *TestClient) Direct() *controlclient.Direct {

hscontrol/servertest/lifecycle_test.go

@@ -361,6 +361,45 @@ func TestReloginPeerVisibleHostinfoReachesPeers(t *testing.T) {
 	}, 5*time.Second, 10*time.Millisecond, "observer must see the relogin's peer-visible Hostinfo")
 }
 
+// TestReloginKeepsNodeKeyUnlessForced pins the shape of
+// [servertest.TestClient.Relogin]: a login with a given key presents the
+// current node key, and only a forced one rotates it. The node stays the
+// same either way.
+func TestReloginKeepsNodeKeyUnlessForced(t *testing.T) {
+	t.Parallel()
+
+	srv := servertest.NewServer(t)
+	user := srv.CreateUser(t, "relogin")
+	key := srv.CreatePreAuthKey(t, types.UserID(user.ID))
+	client := servertest.NewClient(t, srv, "relogin",
+		servertest.WithUser(user), servertest.WithAuthKey(key))
+
+	id := findNodeID(t, srv, client.Name)
+	nodeKey := func() string {
+		node, ok := srv.State().GetNodeByID(id)
+		require.True(t, ok)
+
+		return node.NodeKey().String()
+	}
+	first := nodeKey()
+
+	hi := &tailcfg.Hostinfo{BackendLogID: "servertest-relogin", Hostname: "relogin"}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	require.NoError(t, client.Relogin(ctx, key, false, hi))
+	assert.Equal(t, first, nodeKey(), "a login without force keeps the node key")
+
+	require.NoError(t, client.Relogin(ctx, key, true, hi))
+	assert.NotEqual(t, first, nodeKey(), "a forced login rotates the node key")
+
+	assert.Len(t, srv.State().GetNodesByMachineKeyAllUsers(client.MachineKey()), 1,
+		"relogin keeps one node per machine")
+	client.WaitForCondition(t, "netmap after relogin", 10*time.Second,
+		func(nm *netmap.NetworkMap) bool { return nm.SelfNode.ID() == id.NodeID() })
+}
+
 func TestNodeExpiryRouteFailover(t *testing.T) {
 	t.Parallel()