types/key,wgengine/magicsock,control/controlclient,ipn: add debug disco key rotation

Adds the ability to rotate discovery keys on running clients, needed for
testing upcoming disco key distribution changes.

Introduces key.DiscoKey, an atomic container for a disco private key,
public key, and the public key's ShortString, replacing the prior
separate atomic fields.

magicsock.Conn has a new RotateDiscoKey method, and access to this is
provided via localapi and a CLI debug command.

Note that this implementation is primarily for testing as it stands, and
regular use should likely introduce an additional mechanism that allows
the old key to be used for some time, to provide a seamless key rotation
rather than one that invalidates all sessions.

Updates tailscale/corp#34037

Signed-off-by: James Tucker <james@tailscale.com>

4 files changed, 52 insertions, 3 deletions

diff --git a/control/controlclient/auto.go b/control/controlclient/auto.go
index 3cbfe8581..336a8d491 100644
--- a/control/controlclient/auto.go
+++ b/control/controlclient/auto.go
@@ -767,6 +767,13 @@ func (c *Auto) UpdateEndpoints(endpoints []tailcfg.Endpoint) {
 	}
 }
 
+// SetDiscoPublicKey sets the client's Disco public to key and sends the change
+// to the control server.
+func (c *Auto) SetDiscoPublicKey(key key.DiscoPublic) {
+	c.direct.SetDiscoPublicKey(key)
+	c.updateControl()
+}
+
 func (c *Auto) Shutdown() {
 	c.mu.Lock()
 	if c.closed {
diff --git a/control/controlclient/client.go b/control/controlclient/client.go
index d0aa129ae..41b39622b 100644
--- a/control/controlclient/client.go
+++ b/control/controlclient/client.go
@@ -12,6 +12,7 @@ import (
 	"context"
 
 	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
 )
 
 // LoginFlags is a bitmask of options to change the behavior of Client.Login
@@ -80,7 +81,12 @@ type Client interface {
 	// TODO: a server-side change would let us simply upload this
 	// in a separate http request. It has nothing to do with the rest of
 	// the state machine.
+	// Note: the auto client uploads the new endpoints to control immediately.
 	UpdateEndpoints(endpoints []tailcfg.Endpoint)
+	// SetDiscoPublicKey updates the disco public key that will be sent in
+	// future map requests. This should be called after rotating the discovery key.
+	// Note: the auto client uploads the new key to control immediately.
+	SetDiscoPublicKey(key.DiscoPublic)
 	// ClientID returns the ClientID of a client. This ID is meant to
 	// distinguish one client from another.
 	ClientID() int64
diff --git a/control/controlclient/direct.go b/control/controlclient/direct.go
index 62bbb3586..006a801ef 100644
--- a/control/controlclient/direct.go
+++ b/control/controlclient/direct.go
@@ -74,7 +74,6 @@ type Direct struct {
 	logf                  logger.Logf
 	netMon                *netmon.Monitor // non-nil
 	health                *health.Tracker
-	discoPubKey           key.DiscoPublic
 	busClient             *eventbus.Client
 	clientVersionPub      *eventbus.Publisher[tailcfg.ClientVersion]
 	autoUpdatePub         *eventbus.Publisher[AutoUpdate]
@@ -95,6 +94,7 @@ type Direct struct {
 	mu              syncs.Mutex       // mutex guards the following fields
 	serverLegacyKey key.MachinePublic // original ("legacy") nacl crypto_box-based public key; only used for signRegisterRequest on Windows now
 	serverNoiseKey  key.MachinePublic
+	discoPubKey     key.DiscoPublic // protected by mu; can be updated via [SetDiscoPublicKey]
 
 	sfGroup     singleflight.Group[struct{}, *ts2021.Client] // protects noiseClient creation.
 	noiseClient *ts2021.Client                               // also protected by mu
@@ -316,7 +316,6 @@ func NewDirect(opts Options) (*Direct, error) {
 		logf:                  opts.Logf,
 		persist:               opts.Persist.View(),
 		authKey:               opts.AuthKey,
-		discoPubKey:           opts.DiscoPublicKey,
 		debugFlags:            opts.DebugFlags,
 		netMon:                netMon,
 		health:                opts.HealthTracker,
@@ -329,6 +328,7 @@ func NewDirect(opts Options) (*Direct, error) {
 		dnsCache:              dnsCache,
 		dialPlan:              opts.DialPlan,
 	}
+	c.discoPubKey = opts.DiscoPublicKey
 	c.closedCtx, c.closeCtx = context.WithCancel(context.Background())
 
 	c.controlClientID = nextControlClientID.Add(1)
@@ -853,6 +853,14 @@ func (c *Direct) SendUpdate(ctx context.Context) error {
 	return c.sendMapRequest(ctx, false, nil)
 }
 
+// SetDiscoPublicKey updates the disco public key in local state.
+// It does not implicitly trigger [SendUpdate]; callers should arrange for that.
+func (c *Direct) SetDiscoPublicKey(key key.DiscoPublic) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.discoPubKey = key
+}
+
 // ClientID returns the controlClientID of the controlClient.
 func (c *Direct) ClientID() int64 {
 	return c.controlClientID
@@ -902,6 +910,7 @@ func (c *Direct) sendMapRequest(ctx context.Context, isStreaming bool, nu Netmap
 	persist := c.persist
 	serverURL := c.serverURL
 	serverNoiseKey := c.serverNoiseKey
+	discoKey := c.discoPubKey
 	hi := c.hostInfoLocked()
 	backendLogID := hi.BackendLogID
 	connectionHandleForTest := c.connectionHandleForTest
@@ -945,11 +954,12 @@ func (c *Direct) sendMapRequest(ctx context.Context, isStreaming bool, nu Netmap
 	}
 
 	nodeKey := persist.PublicNodeKey()
+
 	request := &tailcfg.MapRequest{
 		Version:                 tailcfg.CurrentCapabilityVersion,
 		KeepAlive:               true,
 		NodeKey:                 nodeKey,
-		DiscoKey:                c.discoPubKey,
+		DiscoKey:                discoKey,
 		Endpoints:               eps,
 		EndpointTypes:           epTypes,
 		Stream:                  isStreaming,
diff --git a/control/controlclient/direct_test.go b/control/controlclient/direct_test.go
index dd93dc7b3..4329fc878 100644
--- a/control/controlclient/direct_test.go
+++ b/control/controlclient/direct_test.go
@@ -20,6 +20,32 @@ import (
 	"tailscale.com/util/eventbus/eventbustest"
 )
 
+func TestSetDiscoPublicKey(t *testing.T) {
+	initialKey := key.NewDisco().Public()
+
+	c := &Direct{
+		discoPubKey: initialKey,
+	}
+
+	c.mu.Lock()
+	if c.discoPubKey != initialKey {
+		t.Fatalf("initial disco key mismatch: got %v, want %v", c.discoPubKey, initialKey)
+	}
+	c.mu.Unlock()
+
+	newKey := key.NewDisco().Public()
+	c.SetDiscoPublicKey(newKey)
+
+	c.mu.Lock()
+	if c.discoPubKey != newKey {
+		t.Fatalf("disco key not updated: got %v, want %v", c.discoPubKey, newKey)
+	}
+	if c.discoPubKey == initialKey {
+		t.Fatal("disco key should have changed")
+	}
+	c.mu.Unlock()
+}
+
 func TestNewDirect(t *testing.T) {
 	hi := hostinfo.New()
 	ni := tailcfg.NetInfo{LinkType: "wired"}