From 6284454ae5995924f69e8d55660d36aff700aaf2 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 24 Mar 2020 22:24:59 -0700
Subject: [PATCH] wgengine/magicsock: if UDP blocked, pick DERP where most
 peers are

Updates #207

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 wgengine/magicsock/magicsock.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 8a42638ad..9f8aba14a 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -389,6 +389,26 @@ func (c *Conn) pickDERPFallback() int {
 		return 0
 	}
 
+	// See where our peers are. Pick wherever the most nodes are.
+	var (
+		peersOnDerp = map[int]int{}
+		best        int
+		bestCount   int
+	)
+	for _, as := range c.addrsByKey {
+		if id := as.derpID(); id != 0 {
+			peersOnDerp[id]++
+			if v := peersOnDerp[id]; v > bestCount {
+				bestCount = v
+				best = id
+			}
+		}
+	}
+	if best != 0 {
+		return best
+	}
+
+	// Otherwise just pick something randomly.
 	h := fnv.New64()
 	h.Write([]byte(fmt.Sprintf("%p/%d", c, processStartUnixNano))) // arbitrary
 	return ids[rand.New(rand.NewSource(int64(h.Sum64()))).Intn(len(ids))]
@@ -1487,6 +1507,16 @@ type AddrSet struct {
 	lastSpray time.Time
 }
 
+// derpID returns this AddrSet's home DERP node, or 0 if none is found.
+func (as *AddrSet) derpID() int {
+	for _, ua := range as.addrs {
+		if ua.IP.Equal(derpMagicIP) {
+			return ua.Port
+		}
+	}
+	return 0
+}
+
 func (as *AddrSet) timeNow() time.Time {
 	if as.clock != nil {
 		return as.clock()