derp: add PeerPresentFlags bitmask to Watch messages

Updates tailscale/corp#17816

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

derp/derp.go

@@ -83,9 +83,16 @@
 	// a bug).
 	framePeerGone = frameType(0x08) // 32B pub key of peer that's gone + 1 byte reason
 
-	// framePeerPresent is like framePeerGone, but for other
-	// members of the DERP region when they're meshed up together.
-	framePeerPresent = frameType(0x09) // 32B pub key of peer that's connected + optional 18B ip:port (16 byte IP + 2 byte BE uint16 port)
+	// framePeerPresent is like framePeerGone, but for other members of the DERP
+	// region when they're meshed up together.
+	//
+	// The message is at least 32 bytes (the public key of the peer that's
+	// connected). If there are at least 18 bytes remaining after that, it's the
+	// 16 byte IP + 2 byte BE uint16 port of the client. If there's another byte
+	// remaining after that, it's a PeerPresentFlags byte.
+	// While current servers send 41 bytes, old servers will send fewer, and newer
+	// servers might send more.
+	framePeerPresent = frameType(0x09)
 
 	// frameWatchConns is how one DERP node in a regional mesh
 	// subscribes to the others in the region.
@@ -128,6 +135,19 @@
 	PeerGoneReasonNotHere      = PeerGoneReasonType(0x01) // server doesn't know about this peer, unexpected
 )
 
+// PeerPresentFlags is an optional byte of bit flags sent after a framePeerPresent message.
+//
+// For a modern server, the value should always be non-zero. If the value is zero,
+// that means the server doesn't support this field.
+type PeerPresentFlags byte
+
+// PeerPresentFlags bits.
+const (
+	PeerPresentIsRegular  = 1 << 0
+	PeerPresentIsMeshPeer = 1 << 1
+	PeerPresentIsProber   = 1 << 2
+)
+
 var bin = binary.BigEndian
 
 func writeUint32(bw *bufio.Writer, v uint32) error {

derp/derp_client.go

@@ -368,6 +368,8 @@ type PeerPresentMessage struct {
 	Key key.NodePublic
 	// IPPort is the remote IP and port of the client.
 	IPPort netip.AddrPort
+	// Flags is a bitmask of info about the client.
+	Flags PeerPresentFlags
 }
 
 func (PeerPresentMessage) msg() {}
@@ -547,18 +549,33 @@ func (c *Client) recvTimeout(timeout time.Duration) (m ReceivedMessage, err erro
 			return pg, nil
 
 		case framePeerPresent:
-			if n < keyLen {
+			remain := b
+			chunk, remain, ok := cutLeadingN(remain, keyLen)
+			if !ok {
 				c.logf("[unexpected] dropping short peerPresent frame from DERP server")
 				continue
 			}
 			var msg PeerPresentMessage
-			msg.Key = key.NodePublicFromRaw32(mem.B(b[:keyLen]))
-			if n >= keyLen+16+2 {
-				msg.IPPort = netip.AddrPortFrom(
-					netip.AddrFrom16([16]byte(b[keyLen:keyLen+16])).Unmap(),
-					binary.BigEndian.Uint16(b[keyLen+16:keyLen+16+2]),
-				)
+			msg.Key = key.NodePublicFromRaw32(mem.B(chunk))
+
+			const ipLen = 16
+			const portLen = 2
+			chunk, remain, ok = cutLeadingN(remain, ipLen+portLen)
+			if !ok {
+				// Older server which didn't send the IP.
+				return msg, nil
 			}
+			msg.IPPort = netip.AddrPortFrom(
+				netip.AddrFrom16([16]byte(chunk[:ipLen])).Unmap(),
+				binary.BigEndian.Uint16(chunk[ipLen:]),
+			)
+
+			chunk, _, ok = cutLeadingN(remain, 1)
+			if !ok {
+				// Older server which doesn't send PeerPresentFlags.
+				return msg, nil
+			}
+			msg.Flags = PeerPresentFlags(chunk[0])
 			return msg, nil
 
 		case frameRecvPacket:
@@ -636,3 +653,10 @@ func (c *Client) LocalAddr() (netip.AddrPort, error) {
 	}
 	return netip.ParseAddrPort(a.String())
 }
+
+func cutLeadingN(b []byte, n int) (chunk, remain []byte, ok bool) {
+	if len(b) >= n {
+		return b[:n], b[n:], true
+	}
+	return nil, b, false
+}

derp/derp_server.go

@@ -566,7 +566,7 @@ func (s *Server) registerClient(c *sclient) {
 	}
 	s.keyOfAddr[c.remoteIPPort] = c.key
 	s.curClients.Add(1)
-	s.broadcastPeerStateChangeLocked(c.key, c.remoteIPPort, true)
+	s.broadcastPeerStateChangeLocked(c.key, c.remoteIPPort, c.presentFlags(), true)
 }
 
 // broadcastPeerStateChangeLocked enqueues a message to all watchers
@@ -574,12 +574,13 @@ func (s *Server) registerClient(c *sclient) {
 // presence changed.
 //
 // s.mu must be held.
-func (s *Server) broadcastPeerStateChangeLocked(peer key.NodePublic, ipPort netip.AddrPort, present bool) {
+func (s *Server) broadcastPeerStateChangeLocked(peer key.NodePublic, ipPort netip.AddrPort, flags PeerPresentFlags, present bool) {
 	for w := range s.watchers {
 		w.peerStateChange = append(w.peerStateChange, peerConnState{
 			peer:    peer,
 			present: present,
 			ipPort:  ipPort,
+			flags:   flags,
 		})
 		go w.requestMeshUpdate()
 	}
@@ -601,7 +602,7 @@ func (s *Server) unregisterClient(c *sclient) {
 			delete(s.clientsMesh, c.key)
 			s.notePeerGoneFromRegionLocked(c.key)
 		}
-		s.broadcastPeerStateChangeLocked(c.key, netip.AddrPort{}, false)
+		s.broadcastPeerStateChangeLocked(c.key, netip.AddrPort{}, 0, false)
 	case *dupClientSet:
 		c.debugLogf("removed duplicate client")
 		if set.removeClient(c) {
@@ -700,6 +701,7 @@ func (s *Server) addWatcher(c *sclient) {
 			peer:    peer,
 			present: true,
 			ipPort:  ac.remoteIPPort,
+			flags:   ac.presentFlags(),
 		})
 	}
 
@@ -1435,11 +1437,26 @@ type sclient struct {
 	peerGoneLim *rate.Limiter
 }
 
+func (c *sclient) presentFlags() PeerPresentFlags {
+	var f PeerPresentFlags
+	if c.info.IsProber {
+		f |= PeerPresentIsProber
+	}
+	if c.canMesh {
+		f |= PeerPresentIsMeshPeer
+	}
+	if f == 0 {
+		return PeerPresentIsRegular
+	}
+	return f
+}
+
 // peerConnState represents whether a peer is connected to the server
 // or not.
 type peerConnState struct {
 	ipPort  netip.AddrPort // if present, the peer's IP:port
 	peer    key.NodePublic
+	flags   PeerPresentFlags
 	present bool
 }
 
@@ -1634,9 +1651,9 @@ func (c *sclient) sendPeerGone(peer key.NodePublic, reason PeerGoneReasonType) e
 }
 
 // sendPeerPresent sends a peerPresent frame, without flushing.
-func (c *sclient) sendPeerPresent(peer key.NodePublic, ipPort netip.AddrPort) error {
+func (c *sclient) sendPeerPresent(peer key.NodePublic, ipPort netip.AddrPort, flags PeerPresentFlags) error {
 	c.setWriteDeadline()
-	const frameLen = keyLen + 16 + 2
+	const frameLen = keyLen + 16 + 2 + 1 // 16 byte IP + 2 byte port + 1 byte flags
 	if err := writeFrameHeader(c.bw.bw(), framePeerPresent, frameLen); err != nil {
 		return err
 	}
@@ -1645,6 +1662,7 @@ func (c *sclient) sendPeerPresent(peer key.NodePublic, ipPort netip.AddrPort) er
 	a16 := ipPort.Addr().As16()
 	copy(payload[keyLen:], a16[:])
 	binary.BigEndian.PutUint16(payload[keyLen+16:], ipPort.Port())
+	payload[keyLen+18] = byte(flags)
 	_, err := c.bw.Write(payload)
 	return err
 }
@@ -1675,7 +1693,7 @@ func (c *sclient) sendMeshUpdates() error {
 		}
 		var err error
 		if pcs.present {
-			err = c.sendPeerPresent(pcs.peer, pcs.ipPort)
+			err = c.sendPeerPresent(pcs.peer, pcs.ipPort, pcs.flags)
 		} else {
 			err = c.sendPeerGone(pcs.peer, PeerGoneReasonDisconnected)
 		}

derp/derp_test.go

@@ -623,7 +623,13 @@ func (tc *testClient) wantPresent(t *testing.T, peers ...key.NodePublic) {
 					}
 				}))
 			}
-			t.Logf("got present with IP %v", m.IPPort)
+			t.Logf("got present with IP %v, flags=%v", m.IPPort, m.Flags)
+			switch m.Flags {
+			case PeerPresentIsMeshPeer, PeerPresentIsRegular:
+				// Okay
+			default:
+				t.Errorf("unexpected PeerPresentIsMeshPeer flags %v", m.Flags)
+			}
 			delete(want, got)
 			if len(want) == 0 {
 				return