From a2463e8948ae8880897f81770905bb97f315f34e Mon Sep 17 00:00:00 2001
From: David Anderson <danderson@tailscale.com>
Date: Fri, 15 Jan 2021 14:55:44 -0800
Subject: [PATCH] wgengine/magicsock: add an option to disable legacy peer
 handling.

Used in tests to ensure we're not relying on behavior we're going
to remove eventually.

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

diff --git a/wgengine/magicsock/legacy.go b/wgengine/magicsock/legacy.go
index 9bb7bb355..d1973d166 100644
--- a/wgengine/magicsock/legacy.go
+++ b/wgengine/magicsock/legacy.go
@@ -30,9 +30,16 @@
 	"tailscale.com/types/wgkey"
 )
 
-var errNoDestinations = errors.New("magicsock: no destinations")
+var (
+	errNoDestinations = errors.New("magicsock: no destinations")
+	errDisabled       = errors.New("magicsock: legacy networking disabled")
+)
 
 func (c *Conn) createLegacyEndpointLocked(pk key.Public, addrs string) (conn.Endpoint, error) {
+	if c.disableLegacy {
+		return nil, errDisabled
+	}
+
 	a := &addrSet{
 		Logf:      c.logf,
 		publicKey: pk,
@@ -78,6 +85,10 @@ func (c *Conn) createLegacyEndpointLocked(pk key.Public, addrs string) (conn.End
 }
 
 func (c *Conn) findLegacyEndpointLocked(ipp netaddr.IPPort, addr *net.UDPAddr, packet []byte) conn.Endpoint {
+	if c.disableLegacy {
+		return nil
+	}
+
 	// Pre-disco: look up their addrSet.
 	if as, ok := c.addrsByUDP[ipp]; ok {
 		as.updateDst(addr)
@@ -139,6 +150,10 @@ func (c *Conn) resetAddrSetStatesLocked() {
 }
 
 func (c *Conn) sendAddrSet(b []byte, as *addrSet) error {
+	if c.disableLegacy {
+		return errDisabled
+	}
+
 	var addrBuf [8]netaddr.IPPort
 	dsts, roamAddr := as.appendDests(addrBuf[:0], b)
 
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 153fa2392..045f0a788 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -119,6 +119,7 @@ type Conn struct {
 	packetListener   nettype.PacketListener
 	noteRecvActivity func(tailcfg.DiscoKey) // or nil, see Options.NoteRecvActivity
 	simulatedNetwork bool
+	disableLegacy    bool
 
 	// ================================================================
 	// No locking required to access these fields, either because
@@ -382,6 +383,11 @@ type Options struct {
 	// triggering macOS and Windows firwall dialog boxes during
 	// "go test").
 	SimulatedNetwork bool
+
+	// DisableLegacyNetworking disables legacy peer handling. When
+	// enabled, only active discovery-aware nodes will be able to
+	// communicate with Conn.
+	DisableLegacyNetworking bool
 }
 
 func (o *Options) logf() logger.Logf {
@@ -1600,7 +1606,9 @@ func (c *Conn) ReceiveIPv4(b []byte) (n int, ep conn.Endpoint, addr *net.UDPAddr
 				c.logf("magicsock: DERP packet received from idle peer %v; created=%v", dm.src.ShortString(), ep != nil)
 			}
 		}
-		asEp = c.addrsByKey[dm.src]
+		if !c.disableLegacy {
+			asEp = c.addrsByKey[dm.src]
+		}
 		c.mu.Unlock()
 
 		if discoEp != nil {
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index 3836b64d0..2261285e0 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -1485,6 +1485,7 @@ func BenchmarkReceiveFrom(b *testing.B) {
 		EndpointsFunc: func(eps []string) {
 			b.Logf("endpoints: %q", eps)
 		},
+		DisableLegacyNetworking: true,
 	})
 	if err != nil {
 		b.Fatal(err)