From 99cb5c1863297170ebe6c63c45ab01eb402a8a28 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Thu, 29 Oct 2026 13:08:25 +0000 Subject: [PATCH] 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 5c2cf6ed6fba69b5d7c05d59a1f9e0a718e7750a, adapted: ApproveAuth via State, no v1 REST in servertest) --- hscontrol/servertest/client.go | 210 ++++++++++++++++++++----- hscontrol/servertest/lifecycle_test.go | 63 ++++++++ hscontrol/servertest/server.go | 31 ++++ 3 files changed, 268 insertions(+), 36 deletions(-) diff --git a/hscontrol/servertest/client.go b/hscontrol/servertest/client.go index 8a92a3411..4ee5de9ef 100644 --- a/hscontrol/servertest/client.go +++ b/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/. + 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 diff --git a/hscontrol/servertest/lifecycle_test.go b/hscontrol/servertest/lifecycle_test.go index de109b378..2f33d4698 100644 --- a/hscontrol/servertest/lifecycle_test.go +++ b/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() diff --git a/hscontrol/servertest/server.go b/hscontrol/servertest/server.go index 2e830b0a6..05a660df7 100644 --- a/hscontrol/servertest/server.go +++ b/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()