servertest: add login, profile, restart and register log helpers

TryNewClient returns a refused login; Login covers profile switches and
interactive logins approved through ApproveAuth; RegisterAttempts reads
the client's register log.

(cherry picked from commit 5c2cf6ed6f, adapted: ApproveAuth via State, no v1 REST in servertest)

hscontrol/servertest/client.go

@@ -3,12 +3,14 @@ package servertest
 import (
 	"context"
 	"fmt"
+	"path"
 	"slices"
 	"sync"
 	"testing"
 	"time"
 
 	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/juanfont/headscale/hscontrol/types/testcapture"
 	"tailscale.com/control/controlclient"
 	"tailscale.com/health"
 	"tailscale.com/net/netmon"
@@ -36,12 +38,15 @@ type TestClient struct {
 	authKey string
 	user    *types.User
 
-	// machineKey and discoKey survive [TestClient.Relogin], as they survive
-	// a tailscaled login.
+	// machineKey survives every login and discoKey every login but
+	// [TestClient.Restart], as with tailscaled.
 	machineKey key.MachinePrivate
 	discoKey   key.DiscoPublic
 	logf       logger.Logf
 
+	// approve completes an interactive login; see [WithInteractiveLogin].
+	approve func(authID string) error
+
 	// Connection lifecycle.
 	pollCtx    context.Context //nolint:containedctx // test-only; context stored for cancel control
 	pollCancel context.CancelFunc
@@ -56,6 +61,9 @@ type TestClient struct {
 	deltaUpdates bool
 	deltaRemoved []tailcfg.NodeID
 
+	// registers: see [TestClient.RegisterAttempts].
+	registers []testcapture.RegisterAttempt
+
 	// updates is a buffered channel that receives a signal
 	// each time a new [netmap.NetworkMap] arrives.
 	updates chan *netmap.NetworkMap
@@ -75,6 +83,8 @@ type clientConfig struct {
 	user         *types.User
 	authKey      string
 	deltaUpdates bool
+	hostinfo     *tailcfg.Hostinfo
+	approve      func(authID string) error
 }
 
 // WithDeltaUpdates makes the client accept incremental map updates the way
@@ -105,6 +115,20 @@ func WithAuthKey(key string) ClientOption {
 	return func(c *clientConfig) { c.authKey = key }
 }
 
+// WithHostinfo registers with hi instead of a Hostinfo holding only the
+// client's name, so a client can request tags or announce routes at
+// registration.
+func WithHostinfo(hi *tailcfg.Hostinfo) ClientOption {
+	return func(c *clientConfig) { c.hostinfo = hi }
+}
+
+// WithInteractiveLogin registers without an auth key, as `tailscale up`
+// does: the server answers with an auth URL and approve, given its auth ID,
+// completes the login, as `headscale auth register` would.
+func WithInteractiveLogin(approve func(authID string) error) ClientOption {
+	return func(c *clientConfig) { c.approve = approve }
+}
+
 // WithUser sets the user for the client. If not set, the harness
 // creates a default user.
 func WithUser(user *types.User) ClientOption {
@@ -116,6 +140,20 @@ func WithUser(user *types.User) ClientOption {
 func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOption) *TestClient {
 	tb.Helper()
 
+	tc, err := TryNewClient(tb, server, name, opts...)
+	if err != nil {
+		tb.Fatalf("servertest: NewClient(%s): %v", name, err)
+	}
+
+	return tc
+}
+
+// TryNewClient is [NewClient] for a registration that may be refused: it
+// returns the login error, and a client that holds no node and does not
+// poll.
+func TryNewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOption) (*TestClient, error) {
+	tb.Helper()
+
 	cc := &clientConfig{
 		hostname: name,
 	}
@@ -136,6 +174,8 @@ func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOpt
 	var authKey string
 
 	switch {
+	case cc.approve != nil:
+		// Interactive: no key.
 	case cc.authKey != "":
 		authKey = cc.authKey
 	case cc.ephemeral:
@@ -163,7 +203,7 @@ func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOpt
 		user:       user,
 		machineKey: key.NewMachine(),
 		discoKey:   key.NewDisco().Public(),
-		logf:       tb.Logf,
+		approve:    cc.approve,
 		updates:    make(chan *netmap.NetworkMap, 64),
 		bus:        bus,
 		dialer:     dialer,
@@ -172,10 +212,24 @@ func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOpt
 		deltaUpdates: cc.deltaUpdates,
 	}
 
-	direct, err := tc.newDirect(persist.Persist{}, authKey, &tailcfg.Hostinfo{
-		BackendLogID: "servertest-" + name,
-		Hostname:     cc.hostname,
-	})
+	tc.logf = func(format string, args ...any) {
+		line := fmt.Sprintf(format, args...)
+		tb.Logf("%s", line)
+
+		tc.mu.Lock()
+		tc.registers, _ = testcapture.AppendRegisterLine(tc.registers, line)
+		tc.mu.Unlock()
+	}
+
+	hi := cc.hostinfo
+	if hi == nil {
+		hi = &tailcfg.Hostinfo{
+			BackendLogID: "servertest-" + name,
+			Hostname:     cc.hostname,
+		}
+	}
+
+	direct, err := tc.newDirect(persist.Persist{}, authKey, hi)
 	if err != nil {
 		tb.Fatalf("servertest: NewDirect(%s): %v", name, err)
 	}
@@ -186,13 +240,18 @@ func NewClient(tb testing.TB, server *TestServer, name string, opts ...ClientOpt
 		tc.cleanup()
 	})
 
-	// Register with the server.
-	tc.register(tb)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	err = tc.tryLogin(ctx, controlclient.LoginDefault, tc.approve)
+	if err != nil {
+		return tc, err
+	}
 
 	// Start long-polling in the background.
 	tc.startPoll(tb)
 
-	return tc
+	return tc, nil
 }
 
 // newDirect builds a [controlclient.Direct] for the client's machine.
@@ -222,14 +281,64 @@ func (c *TestClient) register(tb testing.TB) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	url, err := c.direct.TryLogin(ctx, controlclient.LoginDefault)
+	err := c.tryLogin(ctx, controlclient.LoginDefault, c.approve)
 	if err != nil {
 		tb.Fatalf("servertest: TryLogin(%s): %v", c.Name, err)
 	}
+}
 
-	if url != "" {
-		tb.Fatalf("servertest: TryLogin(%s): unexpected auth URL: %s (expected auto-auth with preauth key)", c.Name, url)
+// tryLogin logs the current Direct in. An auth URL is completed through
+// approve, and is an error without it. A failed login is recorded in
+// [TestClient.RegisterAttempts], as tailscaled logs it.
+func (c *TestClient) tryLogin(ctx context.Context, flags controlclient.LoginFlags, approve func(string) error) error {
+	url, err := c.direct.TryLogin(ctx, flags)
+
+	switch {
+	case err != nil:
+	case url != "" && approve != nil:
+		err = c.waitLogin(ctx, url, approve)
+	case url != "":
+		err = fmt.Errorf("unexpected auth URL %q (expected auto-auth with preauth key)", url) //nolint:err113
 	}
+
+	if err != nil {
+		c.mu.Lock()
+		c.registers = append(c.registers, testcapture.RegisterAttempt{Err: err.Error()})
+		c.mu.Unlock()
+	}
+
+	return err
+}
+
+// waitLogin waits on the auth URL while approve completes the login it names.
+func (c *TestClient) waitLogin(ctx context.Context, url string, approve func(string) error) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	done := make(chan error, 1)
+
+	go func() {
+		_, err := c.direct.WaitLoginURL(ctx, url)
+		done <- err
+	}()
+
+	// The URL ends in the auth ID: /register/<id>.
+	err := approve(path.Base(url))
+	if err != nil {
+		return err
+	}
+
+	return <-done
+}
+
+// RegisterAttempts returns the client's /machine/register exchanges as
+// tailscaled logs them, with each failed login as its own entry; see
+// [testcapture.AppendRegisterLine]. Node keys are short forms, not labels.
+func (c *TestClient) RegisterAttempts() []testcapture.RegisterAttempt {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return slices.Clone(c.registers)
 }
 
 // startPoll begins the long-poll [tailcfg.MapRequest] loop.
@@ -481,50 +590,60 @@ func (c *TestClient) ReloginAndPoll(ctx context.Context) error {
 }
 
 // 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.
+// `tailscale up --auth-key` (force: `--force-reauth`): see [TestClient.Login],
+// with the current login. The first request presents the current node key
+// unless force rotates it. hi replaces the Hostinfo, so routes can be
+// restated.
 //
 // Safe to call from non-test goroutines.
 func (c *TestClient) Relogin(ctx context.Context, authKey string, force bool, hi *tailcfg.Hostinfo) error {
+	flags := controlclient.LoginDefault
+	if force {
+		flags = controlclient.LoginInteractive
+	}
+
+	return c.Login(ctx, c.Persist(), authKey, flags, hi, nil)
+}
+
+// Login logs the machine in from p, as LocalBackend.Start does: a new
+// [controlclient.Direct] keeps the machine key and presents p's node key,
+// unless flags ask for a new one. p is the current login ([TestClient.Persist])
+// for a relogin or restart, another profile's for a profile switch, and
+// empty for a new profile. With approve set and no authKey the login is
+// interactive, see [WithInteractiveLogin]. The poll restarts once the login
+// succeeds; after a failed one the client does not poll.
+//
+// Safe to call from non-test goroutines.
+func (c *TestClient) Login(
+	ctx context.Context,
+	p persist.Persist,
+	authKey string,
+	flags controlclient.LoginFlags,
+	hi *tailcfg.Hostinfo,
+	approve func(authID string) error,
+) 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())
+			return fmt.Errorf("servertest: Login(%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)
+		return fmt.Errorf("servertest: Login(%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)
+	err = c.tryLogin(ctx, flags, approve)
 	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
+		return fmt.Errorf("servertest: Login(%s): %w", c.Name, err)
 	}
 
 	c.resetNetmapState()
@@ -533,6 +652,25 @@ func (c *TestClient) Relogin(ctx context.Context, authKey string, force bool, hi
 	return nil
 }
 
+// Restart logs the machine in again with a new disco key and no auth key,
+// as a restarted tailscaled does: its login is kept, its disco key is not.
+//
+// Safe to call from non-test goroutines.
+func (c *TestClient) Restart(ctx context.Context, hi *tailcfg.Hostinfo) error {
+	c.discoKey = key.NewDisco().Public()
+
+	return c.Login(ctx, c.Persist(), "", controlclient.LoginDefault, hi, nil)
+}
+
+// Persist returns the client's current login: its node keys and user.
+func (c *TestClient) Persist() persist.Persist {
+	if v := c.direct.GetPersist(); v.Valid() {
+		return *v.AsStruct()
+	}
+
+	return persist.Persist{}
+}
+
 // 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

hscontrol/servertest/lifecycle_test.go

@@ -17,6 +17,7 @@ import (
 	"tailscale.com/control/controlclient"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/netmap"
+	"tailscale.com/types/persist"
 )
 
 // TestConnectionLifecycle exercises the core node lifecycle:
@@ -400,6 +401,68 @@ func TestReloginKeepsNodeKeyUnlessForced(t *testing.T) {
 		func(nm *netmap.NetworkMap) bool { return nm.SelfNode.ID() == id.NodeID() })
 }
 
+// TestLoginProfilesAndRestart pins the shapes of the login helpers: an
+// interactive login approved for a user, a refused key login recorded with
+// its answers, a second profile on the same machine, a switch back, and a
+// restart that keeps the login but not the disco key.
+func TestLoginProfilesAndRestart(t *testing.T) {
+	t.Parallel()
+
+	srv := servertest.NewServer(t)
+	owner := srv.CreateUser(t, "owner")
+	other := srv.CreateUser(t, "other")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	approve := func(user *types.User) func(string) error {
+		return func(authID string) error { return srv.ApproveAuth(ctx, authID, user.Name) }
+	}
+
+	client := servertest.NewClient(t, srv, "machine", servertest.WithInteractiveLogin(approve(owner)))
+
+	first := findNodeID(t, srv, client.Name)
+	node, ok := srv.State().GetNodeByID(first)
+	require.True(t, ok)
+	assert.Equal(t, owner.Name, node.User().Name(), "the approving user owns the node")
+
+	once := srv.CreateEphemeralPreAuthKey(t, types.UserID(owner.ID))
+	_, err := servertest.TryNewClient(t, srv, "spent", servertest.WithAuthKey(once))
+	require.NoError(t, err)
+
+	refused, err := servertest.TryNewClient(t, srv, "late", servertest.WithAuthKey(once))
+	require.Error(t, err, "a spent single-use key is refused")
+
+	attempts := refused.RegisterAttempts()
+	require.NotEmpty(t, attempts)
+	assert.Equal(t, err.Error(), attempts[len(attempts)-1].Err, "the refusal is the last attempt")
+	assert.Empty(t, srv.State().GetNodesByMachineKeyAllUsers(refused.MachineKey()), "a refused login leaves no node")
+
+	hi := &tailcfg.Hostinfo{BackendLogID: "servertest-machine", Hostname: "machine"}
+	ownerLogin := client.Persist()
+
+	require.NoError(t, client.Login(ctx, persist.Persist{}, "", controlclient.LoginInteractive, hi, approve(other)))
+	assert.Len(t, srv.State().GetNodesByMachineKeyAllUsers(client.MachineKey()), 2,
+		"another user's profile adds a node on the machine")
+
+	require.NoError(t, client.Login(ctx, ownerLogin, "", controlclient.LoginDefault, hi, nil))
+	client.WaitForCondition(t, "back on the first profile", 10*time.Second,
+		func(nm *netmap.NetworkMap) bool { return nm.SelfNode.ID() == first.NodeID() })
+
+	before, ok := srv.State().GetNodeByID(first)
+	require.True(t, ok)
+
+	require.NoError(t, client.Restart(ctx, hi))
+
+	require.EventuallyWithT(t, func(c *assert.CollectT) {
+		after, ok := srv.State().GetNodeByID(first)
+		if assert.True(c, ok) {
+			assert.Equal(c, before.NodeKey(), after.NodeKey(), "a restart keeps the node key")
+			assert.NotEqual(c, before.DiscoKey(), after.DiscoKey(), "a restart brings a new disco key")
+		}
+	}, 10*time.Second, 50*time.Millisecond)
+}
+
 func TestNodeExpiryRouteFailover(t *testing.T) {
 	t.Parallel()
 

hscontrol/servertest/server.go

@@ -5,6 +5,8 @@
 package servertest
 
 import (
+	"context"
+	"fmt"
 	"net"
 	"net/http"
 	"net/netip"
@@ -14,6 +16,7 @@ import (
 	hscontrol "github.com/juanfont/headscale/hscontrol"
 	"github.com/juanfont/headscale/hscontrol/state"
 	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/juanfont/headscale/hscontrol/util"
 	"tailscale.com/net/memnet"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/dnstype"
@@ -272,6 +275,34 @@ func (s *TestServer) CreateUser(tb testing.TB, name string) *types.User {
 	return u
 }
 
+// ApproveAuth registers the pending login authID to user, as
+// `headscale auth register` does.
+func (s *TestServer) ApproveAuth(_ context.Context, authID, user string) error {
+	id, err := types.AuthIDFromString(authID)
+	if err != nil {
+		return fmt.Errorf("servertest: ApproveAuth: %w", err)
+	}
+
+	u, err := s.st.GetUserByName(user)
+	if err != nil {
+		return fmt.Errorf("servertest: ApproveAuth: %w", err)
+	}
+
+	node, nodeChange, err := s.st.HandleNodeFromAuthPath(id, types.UserID(u.ID), nil, util.RegisterMethodCLI)
+	if err != nil {
+		return fmt.Errorf("servertest: ApproveAuth: %w", err)
+	}
+
+	routeChange, err := s.st.AutoApproveRoutes(node)
+	if err != nil {
+		return fmt.Errorf("servertest: ApproveAuth: %w", err)
+	}
+
+	s.App.Change(nodeChange, routeChange)
+
+	return nil
+}
+
 // CreatePreAuthKey creates a reusable pre-auth key for the given user.
 func (s *TestServer) CreatePreAuthKey(tb testing.TB, userID types.UserID) string {
 	tb.Helper()