From d456bfdc6d48f2fac1d704c619bf570ac6b41fa0 Mon Sep 17 00:00:00 2001
From: David Anderson <danderson@tailscale.com>
Date: Fri, 15 Jan 2021 12:50:33 -0800
Subject: [PATCH] wgengine/magicsock: fix BenchmarkReceiveFrom.

Previously, this benchmark relied on behavior of the legacy
receive codepath, which I changed in 22507adf. With this
change, the benchmark instead relies on the new active discovery
path.

Signed-off-by: David Anderson <danderson@tailscale.com>
---
 wgengine/magicsock/magicsock.go      | 11 ++++++++++-
 wgengine/magicsock/magicsock_test.go | 15 +++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 40089cec2..153fa2392 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -681,6 +681,16 @@ func (c *Conn) callNetInfoCallback(ni *tailcfg.NetInfo) {
 	}
 }
 
+// addValidDiscoPathForTest makes addr a validated disco address for
+// discoKey. It's used in tests to enable receiving of packets from
+// addr without having to spin up the entire active discovery
+// machinery.
+func (c *Conn) addValidDiscoPathForTest(discoKey tailcfg.DiscoKey, addr netaddr.IPPort) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.discoOfAddr[addr] = discoKey
+}
+
 func (c *Conn) SetNetInfoCallback(fn func(*tailcfg.NetInfo)) {
 	if fn == nil {
 		panic("nil NetInfoCallback")
@@ -2147,7 +2157,6 @@ func (c *Conn) SetNetworkMap(nm *controlclient.NetworkMap) {
 			delete(c.sharedDiscoKey, dk)
 		}
 	}
-
 }
 
 func (c *Conn) wantDerpLocked() bool { return c.derpMap != nil }
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index 05e4715ba..3836b64d0 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -1497,6 +1497,21 @@ func BenchmarkReceiveFrom(b *testing.B) {
 	}
 	defer sendConn.Close()
 
+	// Give conn just enough state that it'll recognize sendConn as a
+	// valid peer and not fall through to the legacy magicsock
+	// codepath.
+	discoKey := tailcfg.DiscoKey{31: 1}
+	conn.SetNetworkMap(&controlclient.NetworkMap{
+		Peers: []*tailcfg.Node{
+			{
+				DiscoKey:  discoKey,
+				Endpoints: []string{sendConn.LocalAddr().String()},
+			},
+		},
+	})
+	conn.CreateEndpoint([32]byte{1: 1}, "0000000000000000000000000000000000000000000000000000000000000001.disco.tailscale:12345")
+	conn.addValidDiscoPathForTest(discoKey, netaddr.MustParseIPPort(sendConn.LocalAddr().String()))
+
 	var dstAddr net.Addr = conn.pconn4.LocalAddr()
 	sendBuf := make([]byte, 1<<10)
 	for i := range sendBuf {