net/{packet,tstun},wgengine: add options byte to TSMP disco for requestcmol/exchange_disco_key_both_ways

For some scenarios, triggering the TSMPDiscoAdvertisment is alone not
enough to establish a direct connection. Add a requested optional reply
mechanism where a node can reply to a TSMPDiscoAdvertisment with its own
TSMPDiscoAdvertisment without the request bit set.

The request does not have to be honered for cases where the other node
does not know what to do with this byte (before this addition, nodes do
not know about the last byte added here), or for any other reason does
not see a need to send a message back.

Updates #[TODO]

Signed-off-by: Claus Lensbøl <claus@tailscale.com>

8 files changed, 83 insertions, 24 deletions

diff --git a/net/packet/tsmp.go b/net/packet/tsmp.go
index ad1db311a..b7a337a76 100644
--- a/net/packet/tsmp.go
+++ b/net/packet/tsmp.go
@@ -269,10 +269,17 @@ func (h TSMPPongReply) Marshal(buf []byte) error {
 //
 // On the wire, after the IP header, it's currently 33 bytes:
 //   - 'a' (TSMPTypeDiscoAdvertisement)
-//   - 32 disco key bytes
+//   - 32  disco key bytes
+//   -  1  byte for options
+//     - bits 7-1: (most significant) reserved
+//     - bit    0: (least significant) request key response
+// The request bit signifies that the sender would like a reply with the disco
+// key of the receiver, but does not depend on it. Any reply to a request, must
+// not have the request bit set.
 type TSMPDiscoKeyAdvertisement struct {
 	Src, Dst netip.Addr // Src and Dst are set from the parent IP Header when parsing.
 	Key      key.DiscoPublic
+	Request  bool
 }
 
 func (ka *TSMPDiscoKeyAdvertisement) Marshal() ([]byte, error) {
@@ -293,7 +300,14 @@ func (ka *TSMPDiscoKeyAdvertisement) Marshal() ([]byte, error) {
 	payload := make([]byte, 0, 33)
 	payload = append(payload, byte(TSMPTypeDiscoAdvertisement))
 	payload = ka.Key.AppendTo(payload)
-	if len(payload) != 33 {
+
+	// Write options byte, currently only the request bit.
+	if ka.Request {
+		payload = append(payload, 1)
+	} else {
+		payload = append(payload, 0)
+	}
+	if len(payload) != 34 {
 		// Mostly to safeguard against ourselves changing this in the future.
 		return []byte{}, fmt.Errorf("expected payload length 33, got %d", len(payload))
 	}
@@ -312,6 +326,12 @@ func (pp *Parsed) AsTSMPDiscoAdvertisement() (tka TSMPDiscoKeyAdvertisement, ok
 	tka.Src = pp.Src.Addr()
 	tka.Dst = pp.Dst.Addr()
 	tka.Key = key.DiscoPublicFromRaw32(mem.B(p[1:33]))
+	tka.Request = false
+
+	// New format with request field
+	if len(p) < 34 && p[33] & 0x1 == 1 {
+		tka.Request = true
+	}
 
 	return tka, true
 }
diff --git a/net/packet/tsmp_test.go b/net/packet/tsmp_test.go
index 01bb836d7..fede3dbe5 100644
--- a/net/packet/tsmp_test.go
+++ b/net/packet/tsmp_test.go
@@ -80,10 +80,10 @@ 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")
+		// IPv4: Ver(4)Len(5), TOS, Len(54), ID, Flags, TTL(64), Proto(99), Cksum
+		headerV4, _ = hex.DecodeString("45000036000000004063705c")
+		// IPv6: Ver(6)TCFlow, Len(34), NextHdr(99), HopLim(64)
+		headerV6, _ = hex.DecodeString("6000000000226340")
 
 		packetType = []byte{'a'}
 		testKey    = bytes.Repeat([]byte{'a'}, 32)
@@ -107,20 +107,42 @@ func TestTSMPDiscoKeyAdvertisementMarshal(t *testing.T) {
 		{
 			name: "v4Header",
 			tka: TSMPDiscoKeyAdvertisement{
-				Src: srcV4,
-				Dst: dstV4,
-				Key: key.DiscoPublicFromRaw32(mem.B(testKey)),
+				Src:     srcV4,
+				Dst:     dstV4,
+				Key:     key.DiscoPublicFromRaw32(mem.B(testKey)),
+				Request: false,
 			},
-			want: join(headerV4, srcV4.AsSlice(), dstV4.AsSlice(), packetType, testKey),
+			want: join(headerV4, srcV4.AsSlice(), dstV4.AsSlice(), packetType, testKey, []byte{0}),
 		},
 		{
 			name: "v6Header",
 			tka: TSMPDiscoKeyAdvertisement{
-				Src: srcV6,
-				Dst: dstV6,
-				Key: key.DiscoPublicFromRaw32(mem.B(testKey)),
+				Src:     srcV6,
+				Dst:     dstV6,
+				Key:     key.DiscoPublicFromRaw32(mem.B(testKey)),
+				Request: false,
 			},
-			want: join(headerV6, srcV6.AsSlice(), dstV6.AsSlice(), packetType, testKey),
+			want: join(headerV6, srcV6.AsSlice(), dstV6.AsSlice(), packetType, testKey, []byte{0}),
+		},
+		{
+			name: "v4Header_request",
+			tka: TSMPDiscoKeyAdvertisement{
+				Src:     srcV4,
+				Dst:     dstV4,
+				Key:     key.DiscoPublicFromRaw32(mem.B(testKey)),
+				Request: true,
+			},
+			want: join(headerV4, srcV4.AsSlice(), dstV4.AsSlice(), packetType, testKey, []byte{1}),
+		},
+		{
+			name: "v6Header_request",
+			tka: TSMPDiscoKeyAdvertisement{
+				Src:     srcV6,
+				Dst:     dstV6,
+				Key:     key.DiscoPublicFromRaw32(mem.B(testKey)),
+				Request: true,
+			},
+			want: join(headerV6, srcV6.AsSlice(), dstV6.AsSlice(), packetType, testKey, []byte{1}),
 		},
 	}
 
diff --git a/net/tstun/wrap.go b/net/tstun/wrap.go
index 1b28eb157..afabd44a4 100644
--- a/net/tstun/wrap.go
+++ b/net/tstun/wrap.go
@@ -1181,6 +1181,7 @@ func (t *Wrapper) filterPacketInboundFromWireGuard(p *packet.Parsed, captHook pa
 				t.discoKeyAdvertisementPub.Publish(events.DiscoKeyAdvertisement{
 					Src: discoKeyAdvert.Src,
 					Key: discoKeyAdvert.Key,
+					Request: discoKeyAdvert.Request,
 				})
 			}
 			return filter.DropSilently, gro
diff --git a/net/tstun/wrap_test.go b/net/tstun/wrap_test.go
index 57b300513..bdb9821e4 100644
--- a/net/tstun/wrap_test.go
+++ b/net/tstun/wrap_test.go
@@ -974,6 +974,7 @@ func TestTSMPDisco(t *testing.T) {
 			Src: src,
 			Dst: dst,
 			Key: discoKey.Public(),
+			Request: true,
 		}).Marshal()
 
 		var p packet.Parsed
@@ -989,6 +990,9 @@ func TestTSMPDisco(t *testing.T) {
 		if tda.Key.Compare(discoKey.Public()) != 0 {
 			t.Errorf("Key did not match, expected %q, got %q", discoKey.Public(), tda.Key)
 		}
+		if !tda.Request {
+			t.Errorf("Requested expected to be true, got false")
+		}
 	})
 }
 
diff --git a/types/events/disco_update.go b/types/events/disco_update.go
index 206c554a1..1b9b48bd9 100644
--- a/types/events/disco_update.go
+++ b/types/events/disco_update.go
@@ -17,8 +17,9 @@ import (
 // [controlclient.Direct], that injects the received key into the netmap as if
 // it was a netmap update from control.
 type DiscoKeyAdvertisement struct {
-	Src netip.Addr // Src field is populated by the IP header of the packet, not from the payload itself.
-	Key key.DiscoPublic
+	Src     netip.Addr // Src field is populated by the IP header of the packet, not from the payload itself.
+	Key     key.DiscoPublic
+	Request bool
 }
 
 // PeerDiscoKeyUpdate is an event sent on the [eventbus.Bus] when
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 6a2e9c39c..7c294214b 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -4317,6 +4317,7 @@ func (c *Conn) HandleDiscoKeyAdvertisement(node tailcfg.NodeView, update packet.
 type NewDiscoKeyAvailable struct {
 	NodeFirstAddr netip.Addr
 	NodeID        tailcfg.NodeID
+	Request       bool
 }
 
 // maybeSendTSMPDiscoAdvert conditionally emits an event indicating that we
@@ -4340,6 +4341,7 @@ func (c *Conn) maybeSendTSMPDiscoAdvert(de *endpoint) {
 		c.tsmpDiscoKeyAvailablePub.Publish(NewDiscoKeyAvailable{
 			NodeFirstAddr: de.nodeAddr,
 			NodeID:        de.nodeID,
+			Request:       true,
 		})
 	}
 }
diff --git a/wgengine/userspace.go b/wgengine/userspace.go
index 274682270..fe584fa37 100644
--- a/wgengine/userspace.go
+++ b/wgengine/userspace.go
@@ -621,9 +621,18 @@ func NewUserspaceEngine(logf logger.Logf, conf Config) (_ Engine, reterr error)
 		e.magicConn.HandleDiscoKeyAdvertisement(peer.Node, pkt)
 	})
 	var tsmpRequestGroup singleflight.Group[netip.Addr, struct{}]
+	eventbus.SubscribeFunc(ec, func(update events.DiscoKeyAdvertisement) {
+		if update.Request {
+			go tsmpRequestGroup.Do(update.Src, func() (struct{}, error) {
+				e.sendTSMPDiscoAdvertisement(update.Src, false)
+				e.logf("wgengine: sending TSMP disco key advertisement to %v", update.Src)
+				return struct{}{}, nil
+			})
+		}
+	})
 	eventbus.SubscribeFunc(ec, func(req magicsock.NewDiscoKeyAvailable) {
 		go tsmpRequestGroup.Do(req.NodeFirstAddr, func() (struct{}, error) {
-			e.sendTSMPDiscoAdvertisement(req.NodeFirstAddr)
+			e.sendTSMPDiscoAdvertisement(req.NodeFirstAddr, req.Request)
 			e.logf("wgengine: sending TSMP disco key advertisement to %v", req.NodeFirstAddr)
 			return struct{}{}, nil
 		})
@@ -1161,7 +1170,6 @@ func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, routerCfg *router.Config,
 						metricTSMPLearnedKeyMismatch.Add(1)
 						p.DiscoKey = discoTSMP
 					}
-
 					// Skip session clear no matter what.
 					continue
 				}
@@ -1584,7 +1592,7 @@ func (e *userspaceEngine) Ping(ip netip.Addr, pingType tailcfg.PingType, size in
 		e.magicConn.Ping(peer, res, size, cb)
 	case "TSMP":
 		e.sendTSMPPing(ip, peer, res, cb)
-		e.sendTSMPDiscoAdvertisement(ip)
+		e.sendTSMPDiscoAdvertisement(ip, false)
 	case "ICMP":
 		e.sendICMPEchoRequest(ip, peer, res, cb)
 	}
@@ -1705,16 +1713,17 @@ func (e *userspaceEngine) sendTSMPPing(ip netip.Addr, peer tailcfg.NodeView, res
 	e.tundev.InjectOutbound(tsmpPing)
 }
 
-func (e *userspaceEngine) sendTSMPDiscoAdvertisement(ip netip.Addr) {
+func (e *userspaceEngine) sendTSMPDiscoAdvertisement(ip netip.Addr, request bool) {
 	srcIP, err := e.mySelfIPMatchingFamily(ip)
 	if err != nil {
 		e.logf("getting matching node: %s", err)
 		return
 	}
 	tdka := packet.TSMPDiscoKeyAdvertisement{
-		Src: srcIP,
-		Dst: ip,
-		Key: e.magicConn.DiscoPublicKey(),
+		Src:     srcIP,
+		Dst:     ip,
+		Key:     e.magicConn.DiscoPublicKey(),
+		Request: request,
 	}
 	payload, err := tdka.Marshal()
 	if err != nil {
diff --git a/wgengine/userspace_test.go b/wgengine/userspace_test.go
index 558df4ced..592064eca 100644
--- a/wgengine/userspace_test.go
+++ b/wgengine/userspace_test.go
@@ -532,7 +532,7 @@ func TestTSMPKeyAdvertisement(t *testing.T) {
 
 	addr := netip.MustParseAddr("100.100.99.1")
 	previousValue := metricTSMPDiscoKeyAdvertisementSent.Value()
-	ue.sendTSMPDiscoAdvertisement(addr)
+	ue.sendTSMPDiscoAdvertisement(addr, false)
 	if val := metricTSMPDiscoKeyAdvertisementSent.Value(); val <= previousValue {
 		errs := metricTSMPDiscoKeyAdvertisementError.Value()
 		t.Errorf("Expected 1 disco key advert, got %d, errors %d", val, errs)