From c81814e4f850a14d0225903b0e2d936bd4950ed9 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Fri, 12 Mar 2021 09:45:37 -0800
Subject: [PATCH] derp{,/derphttp},magicsock: tell DERP server when ping acks
 can be expected

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 derp/derp_client.go              | 50 ++++++++++++++++++++------------
 derp/derphttp/derphttp_client.go | 16 +++++++++-
 wgengine/magicsock/magicsock.go  |  1 +
 3 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/derp/derp_client.go b/derp/derp_client.go
index c152611a1..11979d93f 100644
--- a/derp/derp_client.go
+++ b/derp/derp_client.go
@@ -21,13 +21,14 @@
 
 // Client is a DERP client.
 type Client struct {
-	serverKey  key.Public // of the DERP server; not a machine or node key
-	privateKey key.Private
-	publicKey  key.Public // of privateKey
-	logf       logger.Logf
-	nc         Conn
-	br         *bufio.Reader
-	meshKey    string
+	serverKey   key.Public // of the DERP server; not a machine or node key
+	privateKey  key.Private
+	publicKey   key.Public // of privateKey
+	logf        logger.Logf
+	nc          Conn
+	br          *bufio.Reader
+	meshKey     string
+	canAckPings bool
 
 	wmu sync.Mutex // hold while writing to bw
 	bw  *bufio.Writer
@@ -48,8 +49,9 @@ func (f clientOptFunc) update(o *clientOpt) { f(o) }
 
 // clientOpt are the options passed to newClient.
 type clientOpt struct {
-	MeshKey   string
-	ServerPub key.Public
+	MeshKey     string
+	ServerPub   key.Public
+	CanAckPings bool
 }
 
 // MeshKey returns a ClientOpt to pass to the DERP server during connect to get
@@ -64,6 +66,12 @@ func ServerPublicKey(key key.Public) ClientOpt {
 	return clientOptFunc(func(o *clientOpt) { o.ServerPub = key })
 }
 
+// CanAckPings returns a ClientOpt to set whether it advertises to the
+// server that it's capable of acknowledging ping requests.
+func CanAckPings(v bool) ClientOpt {
+	return clientOptFunc(func(o *clientOpt) { o.CanAckPings = v })
+}
+
 func NewClient(privateKey key.Private, nc Conn, brw *bufio.ReadWriter, logf logger.Logf, opts ...ClientOpt) (*Client, error) {
 	var opt clientOpt
 	for _, o := range opts {
@@ -77,13 +85,14 @@ func NewClient(privateKey key.Private, nc Conn, brw *bufio.ReadWriter, logf logg
 
 func newClient(privateKey key.Private, nc Conn, brw *bufio.ReadWriter, logf logger.Logf, opt clientOpt) (*Client, error) {
 	c := &Client{
-		privateKey: privateKey,
-		publicKey:  privateKey.Public(),
-		logf:       logf,
-		nc:         nc,
-		br:         brw.Reader,
-		bw:         brw.Writer,
-		meshKey:    opt.MeshKey,
+		privateKey:  privateKey,
+		publicKey:   privateKey.Public(),
+		logf:        logf,
+		nc:          nc,
+		br:          brw.Reader,
+		bw:          brw.Writer,
+		meshKey:     opt.MeshKey,
+		canAckPings: opt.CanAckPings,
 	}
 	if opt.ServerPub.IsZero() {
 		if err := c.recvServerKey(); err != nil {
@@ -147,6 +156,10 @@ type clientInfo struct {
 	// connection list & forward packets. It's empty for regular
 	// users.
 	MeshKey string `json:"meshKey,omitempty"`
+
+	// CanAckPings is whether the client declares it's able to ack
+	// pings.
+	CanAckPings bool
 }
 
 func (c *Client) sendClientKey() error {
@@ -155,8 +168,9 @@ func (c *Client) sendClientKey() error {
 		return err
 	}
 	msg, err := json.Marshal(clientInfo{
-		Version: ProtocolVersion,
-		MeshKey: c.meshKey,
+		Version:     ProtocolVersion,
+		MeshKey:     c.meshKey,
+		CanAckPings: c.canAckPings,
 	})
 	if err != nil {
 		return err
diff --git a/derp/derphttp/derphttp_client.go b/derp/derphttp/derphttp_client.go
index 28fdef059..e5a4a6be0 100644
--- a/derp/derphttp/derphttp_client.go
+++ b/derp/derphttp/derphttp_client.go
@@ -63,6 +63,7 @@ type Client struct {
 
 	mu           sync.Mutex
 	preferred    bool
+	canAckPings  bool
 	closed       bool
 	netConn      io.Closer
 	client       *derp.Client
@@ -333,7 +334,11 @@ func (c *Client) connect(ctx context.Context, caller string) (client *derp.Clien
 			return nil, 0, fmt.Errorf("GET failed: %v: %s", err, b)
 		}
 	}
-	derpClient, err = derp.NewClient(c.privateKey, httpConn, brw, c.logf, derp.MeshKey(c.MeshKey), derp.ServerPublicKey(serverPub))
+	derpClient, err = derp.NewClient(c.privateKey, httpConn, brw, c.logf,
+		derp.MeshKey(c.MeshKey),
+		derp.ServerPublicKey(serverPub),
+		derp.CanAckPings(c.canAckPings),
+	)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -665,6 +670,15 @@ func (c *Client) SendPong(data [8]byte) error {
 	return dc.SendPong(data)
 }
 
+// SetCanAckPings sets whether this client will reply to ping requests from the server.
+//
+// This only affects future connections.
+func (c *Client) SetCanAckPings(v bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.canAckPings = v
+}
+
 // NotePreferred notes whether this Client is the caller's preferred
 // (home) DERP node. It's only used for stats.
 func (c *Client) NotePreferred(v bool) {
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 24b9ee7d1..945418702 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -1308,6 +1308,7 @@ func (c *Conn) derpWriteChanOfAddr(addr netaddr.IPPort, peer key.Public) chan<-
 		return c.derpMap.Regions[regionID]
 	})
 
+	dc.SetCanAckPings(true)
 	dc.NotePreferred(c.myDerp == regionID)
 	dc.DNSCache = dnscache.Get()