disco, wgengine/magicsock: add custom disco message supportbradfitz/dctp_disco

Add an experimental mechanism for registering custom disco message types
on a magicsock.Conn. Message types 0x80 and above are reserved for
external use; types below 0x80 are reserved for the Tailscale protocol.

disco package:
- Add MinCustomMessageType (0x80) constant
- Add ParseHookFunc type and ParseWithHook function

magicsock package:
- Add CustomDiscoMessage struct defining a custom message type
  (parser, handler, and whether to accept unknown peers)
- Conn.AddCustomDiscoMessage registers a type (panics on reserved
  range or duplicate registration)
- Conn.SendCustomDiscoOverDERP sends a custom message via DERP
  (rejects unregistered types)

The feature is gated behind the "customdisco" feature tag
(ts_omit_customdisco build tag) for dead code elimination.

As of 2026-03-10 this is not yet a guaranteed stable API.

Updates tailscale/corp#24454

Change-Id: I33b385f94ef63de9d359b9820203b2a7162dc609
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

8 files changed, 350 insertions, 2 deletions

diff --git a/disco/disco.go b/disco/disco.go
index 8f667b262..dd118878c 100644
--- a/disco/disco.go
+++ b/disco/disco.go
@@ -28,6 +28,7 @@ import (
 	"time"
 
 	"go4.org/mem"
+	"tailscale.com/feature/buildfeatures"
 	"tailscale.com/types/key"
 )
 
@@ -39,6 +40,9 @@ const keyLen = 32
 // NonceLen is the length of the nonces used by nacl box.
 const NonceLen = 24
 
+// MessageType is the type byte at the start of a disco message payload.
+// Values 0x01 through 0x7f are reserved for the Tailscale disco protocol.
+// Values 0x80 and above are available for external/custom disco protocols.
 type MessageType byte
 
 const (
@@ -53,10 +57,21 @@ const (
 	TypeAllocateUDPRelayEndpointResponse = MessageType(0x09)
 )
 
+// MinCustomMessageType is the minimum MessageType value available for
+// external/custom disco protocols. Types below this value are reserved
+// for the Tailscale disco protocol.
+const MinCustomMessageType = MessageType(0x80)
+
 const v0 = byte(0)
 
 var errShort = errors.New("short message")
 
+// ParseHookFunc is the signature of a hook function that can parse custom
+// disco message types. It is called by ParseWithHook for message types
+// >= MinCustomMessageType (0x80) that are not handled by the standard parser.
+// If the hook returns (nil, nil), the default "unknown message type" error is used.
+type ParseHookFunc func(msgType MessageType, ver uint8, p []byte) (Message, error)
+
 // LooksLikeDiscoWrapper reports whether p looks like it's a packet
 // containing an encrypted disco message.
 func LooksLikeDiscoWrapper(p []byte) bool {
@@ -108,6 +123,27 @@ func Parse(p []byte) (Message, error) {
 	}
 }
 
+// ParseWithHook is like Parse but calls hook for message types
+// >= MinCustomMessageType (0x80) that are not handled by the standard parser.
+// If hook is nil, it behaves identically to Parse.
+func ParseWithHook(p []byte, hook ParseHookFunc) (Message, error) {
+	if !buildfeatures.HasCustomDisco {
+		return Parse(p)
+	}
+	if len(p) < 2 {
+		return nil, errShort
+	}
+	t := MessageType(p[0])
+	if t < MinCustomMessageType || hook == nil {
+		return Parse(p)
+	}
+	ver, rest := p[1], p[2:]
+	if m, err := hook(t, ver, rest); m != nil || err != nil {
+		return m, err
+	}
+	return nil, fmt.Errorf("unknown message type 0x%02x", byte(t))
+}
+
 // Message a discovery message.
 type Message interface {
 	// AppendMarshal appends the message's marshaled representation.
diff --git a/feature/buildfeatures/feature_customdisco_disabled.go b/feature/buildfeatures/feature_customdisco_disabled.go
new file mode 100644
index 000000000..d65decc67
--- /dev/null
+++ b/feature/buildfeatures/feature_customdisco_disabled.go
@@ -0,0 +1,13 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Code generated by gen.go; DO NOT EDIT.
+
+//go:build ts_omit_customdisco
+
+package buildfeatures
+
+// HasCustomDisco is whether the binary was built with support for modular feature "Custom disco message support".
+// Specifically, it's whether the binary was NOT built with the "ts_omit_customdisco" build tag.
+// It's a const so it can be used for dead code elimination.
+const HasCustomDisco = false
diff --git a/feature/buildfeatures/feature_customdisco_enabled.go b/feature/buildfeatures/feature_customdisco_enabled.go
new file mode 100644
index 000000000..44fca2bc0
--- /dev/null
+++ b/feature/buildfeatures/feature_customdisco_enabled.go
@@ -0,0 +1,13 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Code generated by gen.go; DO NOT EDIT.
+
+//go:build !ts_omit_customdisco
+
+package buildfeatures
+
+// HasCustomDisco is whether the binary was built with support for modular feature "Custom disco message support".
+// Specifically, it's whether the binary was NOT built with the "ts_omit_customdisco" build tag.
+// It's a const so it can be used for dead code elimination.
+const HasCustomDisco = true
diff --git a/feature/featuretags/featuretags.go b/feature/featuretags/featuretags.go
index 4220c02b7..ae0e2cd94 100644
--- a/feature/featuretags/featuretags.go
+++ b/feature/featuretags/featuretags.go
@@ -144,7 +144,8 @@ var Features = map[FeatureTag]FeatureMeta{
 		Sym: "CompletionScripts", Desc: "embed CLI shell completion scripts",
 		Deps: []FeatureTag{"completion"},
 	},
-	"cloud": {Sym: "Cloud", Desc: "detect cloud environment to learn instances IPs and DNS servers"},
+	"cloud":       {Sym: "Cloud", Desc: "detect cloud environment to learn instances IPs and DNS servers"},
+	"customdisco": {Sym: "CustomDisco", Desc: "Custom disco message support"},
 	"dbus": {
 		Sym:                  "DBus",
 		Desc:                 "Linux DBus support",
diff --git a/wgengine/magicsock/custom-disco.go b/wgengine/magicsock/custom-disco.go
new file mode 100644
index 000000000..c00fb9c00
--- /dev/null
+++ b/wgengine/magicsock/custom-disco.go
@@ -0,0 +1,125 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+//go:build !ts_omit_customdisco
+
+package magicsock
+
+import (
+	"errors"
+	"fmt"
+	"net/netip"
+
+	"tailscale.com/disco"
+	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
+)
+
+// customDiscoRegistry is the type of Conn.customDisco when custom disco
+// support is compiled in.
+type customDiscoRegistry map[disco.MessageType]*CustomDiscoMessage
+
+// CustomDiscoMessage defines a custom disco message type for use with
+// AddCustomDiscoMessage. This is an experimental interface for extending the
+// disco protocol; as of 2026-03-10 it is not yet a guaranteed stable API.
+type CustomDiscoMessage struct {
+	// MessageType is the disco message type byte. It must be >=
+	// disco.MinCustomMessageType (0x80); AddCustomDiscoMessage panics
+	// otherwise. Values below 0x80 are reserved for the Tailscale
+	// disco protocol.
+	MessageType disco.MessageType
+
+	// Parse parses the raw message payload (after the type and version
+	// header bytes) into a disco.Message. If it returns (nil, nil) the
+	// message is treated as an unknown type.
+	Parse disco.ParseHookFunc
+
+	// AcceptUnknownPeers, if true, causes disco messages to be
+	// accepted even from peers not present in the netmap.
+	AcceptUnknownPeers bool
+
+	// HandleMessage, if non-nil, is called after a received disco message is
+	// parsed.
+	HandleMessage func(dm disco.Message, sender key.DiscoPublic, derpNodeSrc key.NodePublic)
+}
+
+// AddCustomDiscoMessage registers a custom disco message type on the Conn.
+// See CustomDiscoMessage for field documentation.
+//
+// It panics if m.MessageType < disco.MinCustomMessageType (0x80) or if a
+// handler for the same MessageType has already been registered.
+func (c *Conn) AddCustomDiscoMessage(m *CustomDiscoMessage) {
+	if m.MessageType < disco.MinCustomMessageType {
+		panic(fmt.Sprintf("disco message type 0x%02x is in the reserved range (must be >= 0x%02x)", byte(m.MessageType), byte(disco.MinCustomMessageType)))
+	}
+	if _, dup := c.customDisco[m.MessageType]; dup {
+		panic(fmt.Sprintf("duplicate registration for disco message type 0x%02x", byte(m.MessageType)))
+	}
+	if c.customDisco == nil {
+		c.customDisco = make(customDiscoRegistry)
+	}
+	c.customDisco[m.MessageType] = m
+}
+
+// customDiscoAcceptsUnknownPeers reports whether any registered custom disco
+// message type has AcceptUnknownPeers set.
+func (c *Conn) customDiscoAcceptsUnknownPeers() bool {
+	for _, cd := range c.customDisco {
+		if cd.AcceptUnknownPeers {
+			return true
+		}
+	}
+	return false
+}
+
+// customDiscoParseHook dispatches to the Parse hook of the registered custom
+// disco message type matching msgType.
+func (c *Conn) customDiscoParseHook(msgType disco.MessageType, ver uint8, p []byte) (disco.Message, error) {
+	if cd, ok := c.customDisco[msgType]; ok && cd.Parse != nil {
+		return cd.Parse(msgType, ver, p)
+	}
+	return nil, nil
+}
+
+// handleMessage dispatches a parsed disco message to the registered handler
+// for the given message type.
+func (r customDiscoRegistry) handleMessage(msgType disco.MessageType, dm disco.Message, sender key.DiscoPublic, derpNodeSrc key.NodePublic) {
+	if cd, ok := r[msgType]; ok && cd.HandleMessage != nil {
+		cd.HandleMessage(dm, sender, derpNodeSrc)
+	}
+}
+
+// SendCustomDiscoOverDERP sends a disco message to a peer identified by
+// its disco and node public keys via the specified DERP region.
+//
+// It returns an error if the message's type byte (the first byte of its
+// marshaled form) has not been registered via AddCustomDiscoMessage.
+func (c *Conn) SendCustomDiscoOverDERP(dstDisco key.DiscoPublic, dstNode key.NodePublic, derpRegion int, m disco.Message) (sent bool, err error) {
+	payload := m.AppendMarshal(nil)
+	if len(payload) == 0 {
+		return false, errors.New("empty disco message")
+	}
+	msgType := disco.MessageType(payload[0])
+	if _, ok := c.customDisco[msgType]; !ok {
+		return false, fmt.Errorf("unregistered custom disco message type 0x%02x", byte(msgType))
+	}
+
+	dstAddr := netip.AddrPortFrom(tailcfg.DerpMagicIPAddr, uint16(derpRegion))
+
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return false, errConnClosed
+	}
+	pkt := make([]byte, 0, 512)
+	pkt = append(pkt, disco.Magic...)
+	pkt = c.discoAtomic.Public().AppendTo(pkt)
+	di := c.discoInfoForKnownPeerLocked(dstDisco)
+	c.mu.Unlock()
+
+	box := di.sharedKey.Seal(payload)
+	pkt = append(pkt, box...)
+	const isDisco = true
+	const isGeneveEncap = false
+	return c.sendAddr(dstAddr, dstNode, pkt, isDisco, isGeneveEncap)
+}
diff --git a/wgengine/magicsock/custom-disco_omit.go b/wgengine/magicsock/custom-disco_omit.go
new file mode 100644
index 000000000..1fb5f7a0e
--- /dev/null
+++ b/wgengine/magicsock/custom-disco_omit.go
@@ -0,0 +1,27 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+//go:build ts_omit_customdisco
+
+package magicsock
+
+import (
+	"tailscale.com/disco"
+	"tailscale.com/types/key"
+)
+
+type customDiscoRegistry struct{}
+
+// CustomDiscoMessage is a stub when custom disco support is omitted.
+type CustomDiscoMessage struct{}
+
+func (c *Conn) AddCustomDiscoMessage(*CustomDiscoMessage)        {}
+func (c *Conn) customDiscoAcceptsUnknownPeers() bool             { return false }
+func (c *Conn) customDiscoParseHook(disco.MessageType, uint8, []byte) (disco.Message, error) {
+	return nil, nil
+}
+func (customDiscoRegistry) handleMessage(disco.MessageType, disco.Message, key.DiscoPublic, key.NodePublic) {
+}
+func (c *Conn) SendCustomDiscoOverDERP(key.DiscoPublic, key.NodePublic, int, disco.Message) (bool, error) {
+	return false, nil
+}
diff --git a/wgengine/magicsock/custom-disco_test.go b/wgengine/magicsock/custom-disco_test.go
new file mode 100644
index 000000000..f1ee64fdb
--- /dev/null
+++ b/wgengine/magicsock/custom-disco_test.go
@@ -0,0 +1,123 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+//go:build !ts_omit_customdisco
+
+package magicsock
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"tailscale.com/disco"
+	"tailscale.com/net/netaddr"
+	"tailscale.com/types/key"
+	"tailscale.com/types/logger"
+)
+
+const testCustomDiscoType = disco.MessageType(0x80)
+
+// testCustomDiscoMsg is a minimal custom disco message for testing.
+type testCustomDiscoMsg struct {
+	Data [4]byte
+}
+
+func (m *testCustomDiscoMsg) AppendMarshal(b []byte) []byte {
+	b = append(b, byte(testCustomDiscoType), 0) // type, version
+	b = append(b, m.Data[:]...)
+	return b
+}
+
+func TestCustomDiscoMessage(t *testing.T) {
+	ln, ip := localhostListener{}, netaddr.IPv4(127, 0, 0, 1)
+	d := &devices{
+		m1:     ln,
+		m1IP:   ip,
+		m2:     ln,
+		m2IP:   ip,
+		stun:   ln,
+		stunIP: ip,
+	}
+
+	logf, closeLogf := logger.LogfCloser(t.Logf)
+	defer closeLogf()
+
+	derpMap, cleanup := runDERPAndStun(t, logf, d.stun, d.stunIP)
+	defer cleanup()
+
+	m1 := newMagicStack(t, logger.WithPrefix(logf, "m1: "), d.m1, derpMap)
+	defer m1.Close()
+	m2 := newMagicStack(t, logger.WithPrefix(logf, "m2: "), d.m2, derpMap)
+	defer m2.Close()
+
+	cleanupMesh := meshStacks(logf, nil, m1, m2)
+	defer cleanupMesh()
+
+	// Channel to receive the custom disco message on m2.
+	gotMsg := make(chan *testCustomDiscoMsg, 1)
+
+	parseHook := func(msgType disco.MessageType, ver uint8, p []byte) (disco.Message, error) {
+		if msgType != testCustomDiscoType {
+			return nil, nil
+		}
+		if len(p) < 4 {
+			return nil, errors.New("short message")
+		}
+		m := &testCustomDiscoMsg{}
+		copy(m.Data[:], p[:4])
+		return m, nil
+	}
+
+	handleMsg := func(dm disco.Message, sender key.DiscoPublic, derpNodeSrc key.NodePublic) {
+		if cm, ok := dm.(*testCustomDiscoMsg); ok {
+			gotMsg <- cm
+		}
+	}
+
+	// Register on both sides so m1 can send and m2 can receive.
+	msgDef := &CustomDiscoMessage{
+		MessageType:   testCustomDiscoType,
+		Parse:         parseHook,
+		HandleMessage: handleMsg,
+	}
+	m1.conn.AddCustomDiscoMessage(msgDef)
+	m2.conn.AddCustomDiscoMessage(msgDef)
+
+	// Wait for the mesh to be fully established.
+	deadline := time.Now().Add(5 * time.Second)
+	for time.Now().Before(deadline) {
+		st1 := m1.Status()
+		st2 := m2.Status()
+		if p := st1.Peer[m2.Public()]; p != nil && p.InMagicSock {
+			if p := st2.Peer[m1.Public()]; p != nil && p.InMagicSock {
+				break
+			}
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	// Send a custom disco message from m1 to m2 over DERP region 1.
+	want := [4]byte{'t', 'e', 's', 't'}
+	sent, err := m1.conn.SendCustomDiscoOverDERP(
+		m2.conn.DiscoPublicKey(),
+		m2.privateKey.Public(),
+		1, // DERP region
+		&testCustomDiscoMsg{Data: want},
+	)
+	if err != nil {
+		t.Fatalf("SendCustomDiscoOverDERP: %v", err)
+	}
+	if !sent {
+		t.Fatal("SendCustomDiscoOverDERP reported not sent")
+	}
+
+	select {
+	case got := <-gotMsg:
+		if got.Data != want {
+			t.Errorf("got data %q, want %q", got.Data, want)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for custom disco message")
+	}
+}
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 78ffd0cd0..e3eee2a4c 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -168,6 +168,10 @@ type Conn struct {
 	health                 *health.Tracker      // or nil
 	controlKnobs           *controlknobs.Knobs  // or nil
 
+	// customDisco contains registered custom disco message types,
+	// keyed by MessageType. Set via AddCustomDiscoMessage.
+	customDisco customDiscoRegistry
+
 	// ================================================================
 	// No locking required to access these fields, either because
 	// they're static after construction, or are wholly owned by a
@@ -2151,6 +2155,8 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 		}
 	case c.peerMap.knownPeerDiscoKey(sender):
 		di = c.discoInfoForKnownPeerLocked(sender)
+	case c.customDiscoAcceptsUnknownPeers():
+		di = c.discoInfoForKnownPeerLocked(sender)
 	default:
 		metricRecvDiscoBadPeer.Add(1)
 		if debugDisco() {
@@ -2199,7 +2205,7 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 		cb(packet.PathDisco, time.Now(), disco.ToPCAPFrame(src.ap, derpNodeSrc, payload), packet.CaptureMeta{})
 	}
 
-	dm, err := disco.Parse(payload)
+	dm, err := disco.ParseWithHook(payload, c.customDiscoParseHook)
 	if debugDisco() {
 		c.logf("magicsock: disco: disco.Parse = %T, %v", dm, err)
 	}
@@ -2421,6 +2427,10 @@ func (c *Conn) handleDiscoMessage(msg []byte, src epAddr, shouldBeRelayHandshake
 			RxFromNodeKey:  nodeKey,
 			Message:        req,
 		})
+	default:
+		if buildfeatures.HasCustomDisco {
+			c.customDisco.handleMessage(disco.MessageType(payload[0]), dm, sender, derpNodeSrc)
+		}
 	}
 	return
 }