From c848c5a1a1f50767ae62441a1ac7aed5342db403 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Sat, 17 Oct 2026 20:43:42 +0000 Subject: [PATCH] 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 1716eece742b9eb32365e9648830566c4885127b, adapted: adds main's startPollLoop and resetNetmapState) --- hscontrol/servertest/client.go | 165 ++++++++++++++++++++----- hscontrol/servertest/lifecycle_test.go | 39 ++++++ 2 files changed, 176 insertions(+), 28 deletions(-) diff --git a/hscontrol/servertest/client.go b/hscontrol/servertest/client.go index 3f25b39ab..8a92a3411 100644 --- a/hscontrol/servertest/client.go +++ b/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{ - Persist: persist.Persist{}, - GetMachinePrivateKey: func() (key.MachinePrivate, error) { return machineKey, nil }, - ServerURL: server.URL, - AuthKey: authKey, - Hostinfo: &tailcfg.Hostinfo{ - BackendLogID: "servertest-" + name, - Hostname: cc.hostname, - }, - DiscoPublicKey: key.NewDisco().Public(), - Logf: tb.Logf, - HealthTracker: tracker, - Dialer: dialer, - Bus: bus, + deltaUpdates: cc.deltaUpdates, + } + + direct, err := tc.newDirect(persist.Persist{}, authKey, &tailcfg.Hostinfo{ + BackendLogID: "servertest-" + name, + Hostname: cc.hostname, }) if err != nil { tb.Fatalf("servertest: NewDirect(%s): %v", name, err) } - tc := &TestClient{ - 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, - } + tc.direct = direct 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 { diff --git a/hscontrol/servertest/lifecycle_test.go b/hscontrol/servertest/lifecycle_test.go index b83bd1983..de109b378 100644 --- a/hscontrol/servertest/lifecycle_test.go +++ b/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()