net/tstun: add TSMPDiscoAdvertisement to TSMPPing (#17995)

Adds a new types of TSMP messages for advertising disco keys keys
to/from a peer, and implements the advertising triggered by a TSMP ping.

Needed as part of the effort to cache the netmap and still let clients
connect without control being reachable.

Updates #12639

Signed-off-by: Claus Lensbøl <claus@tailscale.com>
Co-authored-by: James Tucker <james@tailscale.com>

net/packet/tsmp.go

@@ -15,7 +15,9 @@ import (
 	"fmt"
 	"net/netip"
 
+	"go4.org/mem"
 	"tailscale.com/types/ipproto"
+	"tailscale.com/types/key"
 )
 
 const minTSMPSize = 7 // the rejected body is 7 bytes
@@ -72,6 +74,9 @@ const (
 
 	// TSMPTypePong is the type byte for a TailscalePongResponse.
 	TSMPTypePong TSMPType = 'o'
+
+	// TSPMTypeDiscoAdvertisement is the type byte for sending disco keys
+	TSMPTypeDiscoAdvertisement TSMPType = 'a'
 )
 
 type TailscaleRejectReason byte
@@ -259,3 +264,53 @@ func (h TSMPPongReply) Marshal(buf []byte) error {
 	binary.BigEndian.PutUint16(buf[9:11], h.PeerAPIPort)
 	return nil
 }
+
+// TSMPDiscoKeyAdvertisement is a TSMP message that's used for distributing Disco Keys.
+//
+// On the wire, after the IP header, it's currently 33 bytes:
+//   - 'a' (TSMPTypeDiscoAdvertisement)
+//   - 32 disco key bytes
+type TSMPDiscoKeyAdvertisement struct {
+	Src, Dst netip.Addr
+	Key      key.DiscoPublic
+}
+
+func (ka *TSMPDiscoKeyAdvertisement) Marshal() ([]byte, error) {
+	var iph Header
+	if ka.Src.Is4() {
+		iph = IP4Header{
+			IPProto: ipproto.TSMP,
+			Src:     ka.Src,
+			Dst:     ka.Dst,
+		}
+	} else {
+		iph = IP6Header{
+			IPProto: ipproto.TSMP,
+			Src:     ka.Src,
+			Dst:     ka.Dst,
+		}
+	}
+	payload := make([]byte, 0, 33)
+	payload = append(payload, byte(TSMPTypeDiscoAdvertisement))
+	payload = ka.Key.AppendTo(payload)
+	if len(payload) != 33 {
+		// Mostly to safeguard against ourselves changing this in the future.
+		return []byte{}, fmt.Errorf("expected payload length 33, got %d", len(payload))
+	}
+
+	return Generate(iph, payload), nil
+}
+
+func (pp *Parsed) AsTSMPDiscoAdvertisement() (tka TSMPDiscoKeyAdvertisement, ok bool) {
+	if pp.IPProto != ipproto.TSMP {
+		return
+	}
+	p := pp.Payload()
+	if len(p) < 33 || p[0] != byte(TSMPTypeDiscoAdvertisement) {
+		return
+	}
+	tka.Src = pp.Src.Addr()
+	tka.Key = key.DiscoPublicFromRaw32(mem.B(p[1:33]))
+
+	return tka, true
+}

net/packet/tsmp_test.go

@@ -4,8 +4,14 @@
 package packet
 
 import (
+	"bytes"
+	"encoding/hex"
 	"net/netip"
+	"slices"
 	"testing"
+
+	"go4.org/mem"
+	"tailscale.com/types/key"
 )
 
 func TestTailscaleRejectedHeader(t *testing.T) {
@@ -71,3 +77,62 @@ func TestTailscaleRejectedHeader(t *testing.T) {
 		}
 	}
 }
+
+func TestTSMPDiscoKeyAdvertisementMarshal(t *testing.T) {
+	var (
+		// IPv4: Ver(4)Len(5), TOS, Len(53), ID, Flags, TTL(64), Proto(99), Cksum
+		headerV4, _ = hex.DecodeString("45000035000000004063705d")
+		// IPv6: Ver(6)TCFlow, Len(33), NextHdr(99), HopLim(64)
+		headerV6, _ = hex.DecodeString("6000000000216340")
+
+		packetType = []byte{'a'}
+		testKey    = bytes.Repeat([]byte{'a'}, 32)
+
+		// IPs
+		srcV4 = netip.MustParseAddr("1.2.3.4")
+		dstV4 = netip.MustParseAddr("4.3.2.1")
+		srcV6 = netip.MustParseAddr("2001:db8::1")
+		dstV6 = netip.MustParseAddr("2001:db8::2")
+	)
+
+	join := func(parts ...[]byte) []byte {
+		return bytes.Join(parts, nil)
+	}
+
+	tests := []struct {
+		name string
+		tka  TSMPDiscoKeyAdvertisement
+		want []byte
+	}{
+		{
+			name: "v4Header",
+			tka: TSMPDiscoKeyAdvertisement{
+				Src: srcV4,
+				Dst: dstV4,
+				Key: key.DiscoPublicFromRaw32(mem.B(testKey)),
+			},
+			want: join(headerV4, srcV4.AsSlice(), dstV4.AsSlice(), packetType, testKey),
+		},
+		{
+			name: "v6Header",
+			tka: TSMPDiscoKeyAdvertisement{
+				Src: srcV6,
+				Dst: dstV6,
+				Key: key.DiscoPublicFromRaw32(mem.B(testKey)),
+			},
+			want: join(headerV6, srcV6.AsSlice(), dstV6.AsSlice(), packetType, testKey),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := tt.tka.Marshal()
+			if err != nil {
+				t.Errorf("error mashalling TSMPDiscoAdvertisement: %s", err)
+			}
+			if !slices.Equal(got, tt.want) {
+				t.Errorf("error mashalling TSMPDiscoAdvertisement, expected: \n%x, \ngot:\n%x", tt.want, got)
+			}
+		})
+	}
+}