From a83ca9e734c2e23b349c9eb8d06551956825907b Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 29 Jun 2020 14:26:25 -0700
Subject: [PATCH] wgengine/magicsock: cache precomputed nacl/box shared keys

Updates #483
---
 wgengine/magicsock/magicsock.go      | 20 +++++++++++++++++++-
 wgengine/magicsock/magicsock_test.go | 11 +++++------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index be0183ec3..0df25a155 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -95,6 +95,7 @@ type Conn struct {
 	discoOfNode  map[tailcfg.NodeKey]tailcfg.DiscoKey
 
 	endpointOfDisco map[tailcfg.DiscoKey]*discoEndpoint
+	sharedDiscoKey  map[tailcfg.DiscoKey]*[32]byte // nacl/box precomputed key
 
 	// addrsByUDP is a map of every remote ip:port to a priority
 	// list of endpoint addresses for a peer.
@@ -250,6 +251,7 @@ func newConn() *Conn {
 		derpStarted:     make(chan struct{}),
 		peerLastDerp:    make(map[key.Public]int),
 		endpointOfDisco: make(map[tailcfg.DiscoKey]*discoEndpoint),
+		sharedDiscoKey:  make(map[tailcfg.DiscoKey]*[32]byte),
 	}
 	c.endpointsUpdateWaiter = sync.NewCond(&c.mu)
 	return c
@@ -499,6 +501,11 @@ func (c *Conn) SetNetInfoCallback(fn func(*tailcfg.NetInfo)) {
 func (c *Conn) SetDiscoPrivateKey(k key.Private) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	if !c.discoPrivate.IsZero() && c.discoPrivate != k {
+		// TODO: support changing a key at runtime; need to
+		// clear a bunch of maps at least
+		panic("unsupported")
+	}
 	c.discoPrivate = k
 	c.logf("magicsock: disco key set; public: %x", k.Public())
 }
@@ -1348,7 +1355,7 @@ func (c *Conn) handleDiscoMessage(msg []byte, addr *net.UDPAddr) bool {
 	var nonce [nonceLen]byte
 	copy(nonce[:], msg[len(magic)+len(key.Public{}):])
 	sealedBox := msg[headerLen:]
-	payload, ok := box.Open(nil, sealedBox, &nonce, key.Public(sender).B32(), c.discoPrivate.B32())
+	payload, ok := box.OpenAfterPrecomputation(nil, sealedBox, &nonce, c.sharedDiscoKeyLocked(sender))
 	if !ok {
 		c.logf("magicsock: failed to open disco message box purportedly from %s (disco key %x)", senderNode.Key.ShortString(), sender[:])
 		return false
@@ -1358,6 +1365,16 @@ func (c *Conn) handleDiscoMessage(msg []byte, addr *net.UDPAddr) bool {
 	return true
 }
 
+func (c *Conn) sharedDiscoKeyLocked(k tailcfg.DiscoKey) *[32]byte {
+	if v, ok := c.sharedDiscoKey[k]; ok {
+		return v
+	}
+	shared := new([32]byte)
+	box.Precompute(shared, key.Public(k).B32(), c.discoPrivate.B32())
+	c.sharedDiscoKey[k] = shared
+	return shared
+}
+
 // SetPrivateKey sets the connection's private key.
 //
 // This is only used to be able prove our identity when connecting to
@@ -1491,6 +1508,7 @@ func (c *Conn) SetNetworkMap(nm *controlclient.NetworkMap) {
 		if _, ok := c.nodeOfDisco[dk]; !ok {
 			de.cleanup()
 			delete(c.endpointOfDisco, dk)
+			delete(c.sharedDiscoKey, dk)
 		}
 	}
 
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index 174be429e..3cd0f359d 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -841,12 +841,11 @@ func TestDiscoMessage(t *testing.T) {
 	peer1Priv := key.NewPrivate()
 	peer1Pub := peer1Priv.Public()
 
-	c := &Conn{
-		logf:         t.Logf,
-		discoPrivate: key.NewPrivate(),
-		nodeOfDisco: map[tailcfg.DiscoKey]*tailcfg.Node{
-			tailcfg.DiscoKey(peer1Pub): &tailcfg.Node{Key: tailcfg.NodeKey{1: 1}},
-		},
+	c := newConn()
+	c.logf = t.Logf
+	c.SetDiscoPrivateKey(key.NewPrivate())
+	c.nodeOfDisco = map[tailcfg.DiscoKey]*tailcfg.Node{
+		tailcfg.DiscoKey(peer1Pub): &tailcfg.Node{Key: tailcfg.NodeKey{1: 1}},
 	}
 
 	const payload = "why hello"