From f873da5b1687cd559974ee03e39856cef16337ff Mon Sep 17 00:00:00 2001
From: David Anderson <danderson@tailscale.com>
Date: Fri, 18 Dec 2020 00:31:48 -0800
Subject: [PATCH] wgengine/magicsock: move more legacy endpoint handling.

Signed-off-by: David Anderson <danderson@tailscale.com>
---
 wgengine/magicsock/legacy.go    | 21 +++++++++++++++++++++
 wgengine/magicsock/magicsock.go | 22 +++++-----------------
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/wgengine/magicsock/legacy.go b/wgengine/magicsock/legacy.go
index aafa4a82a..af87cb147 100644
--- a/wgengine/magicsock/legacy.go
+++ b/wgengine/magicsock/legacy.go
@@ -13,6 +13,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/tailscale/wireguard-go/conn"
 	"github.com/tailscale/wireguard-go/device"
 	"github.com/tailscale/wireguard-go/wgcfg"
 	"inet.af/netaddr"
@@ -23,6 +24,26 @@ import (
 
 var errNoDestinations = errors.New("magicsock: no destinations")
 
+func (c *Conn) findLegacyEndpointLocked(ipp netaddr.IPPort, addr *net.UDPAddr) conn.Endpoint {
+	// Pre-disco: look up their addrSet.
+	if as, ok := c.addrsByUDP[ipp]; ok {
+		return as
+	}
+
+	// Pre-disco: the peer that sent this packet has roamed beyond
+	// the knowledge provided by the control server.  If the
+	// packet is valid wireguard will call UpdateDst on the
+	// original endpoint using this addr.
+	return (*singleEndpoint)(addr)
+}
+
+func (c *Conn) resetAddrSetStatesLocked() {
+	for _, as := range c.addrsByKey {
+		as.curAddr = -1
+		as.stopSpray = as.timeNow().Add(sprayPeriod)
+	}
+}
+
 func (c *Conn) sendSingleEndpoint(b []byte, se *singleEndpoint) error {
 	addr := (*net.UDPAddr)(se)
 	if addr.IP.Equal(derpMagicIP) {
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 95b0e20a0..8e3215eff 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -1395,16 +1395,7 @@ func (c *Conn) findEndpoint(ipp netaddr.IPPort, addr *net.UDPAddr) conn.Endpoint
 		}
 	}
 
-	// Pre-disco: look up their addrSet.
-	if as, ok := c.addrsByUDP[ipp]; ok {
-		return as
-	}
-
-	// Pre-disco: the peer that sent this packet has roamed beyond
-	// the knowledge provided by the control server.  If the
-	// packet is valid wireguard will call UpdateDst on the
-	// original endpoint using this addr.
-	return (*singleEndpoint)(addr)
+	return c.findLegacyEndpointLocked(ipp, addr)
 }
 
 type udpReadResult struct {
@@ -2476,23 +2467,20 @@ func (c *Conn) Rebind() {
 	if haveKey {
 		c.goDerpConnect(c.myDerp)
 	}
-	c.resetAddrSetStates()
+	c.resetEndpointStates()
 }
 
-// resetAddrSetStates resets the preferred address for all peers and
+// resetEndpointStates resets the preferred address for all peers and
 // re-enables spraying.
 // This is called when connectivity changes enough that we no longer
 // trust the old routes.
-func (c *Conn) resetAddrSetStates() {
+func (c *Conn) resetEndpointStates() {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	for _, as := range c.addrsByKey {
-		as.curAddr = -1
-		as.stopSpray = as.timeNow().Add(sprayPeriod)
-	}
 	for _, de := range c.endpointOfDisco {
 		de.noteConnectivityChange()
 	}
+	c.resetAddrSetStatesLocked()
 }
 
 // packIPPort packs an IPPort into the form wanted by WireGuard.