From adc961352c92edf608eb752a84434a29093a5c5d Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Tue, 10 Mar 2026 16:04:02 +0000 Subject: [PATCH] disco, wgengine/magicsock: add custom disco message support 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 --- disco/disco.go | 36 +++++ .../feature_customdisco_disabled.go | 13 ++ .../feature_customdisco_enabled.go | 13 ++ feature/featuretags/featuretags.go | 3 +- wgengine/magicsock/custom-disco.go | 125 ++++++++++++++++++ wgengine/magicsock/custom-disco_omit.go | 27 ++++ wgengine/magicsock/custom-disco_test.go | 123 +++++++++++++++++ wgengine/magicsock/magicsock.go | 12 +- 8 files changed, 350 insertions(+), 2 deletions(-) create mode 100644 feature/buildfeatures/feature_customdisco_disabled.go create mode 100644 feature/buildfeatures/feature_customdisco_enabled.go create mode 100644 wgengine/magicsock/custom-disco.go create mode 100644 wgengine/magicsock/custom-disco_omit.go create mode 100644 wgengine/magicsock/custom-disco_test.go 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 }