From 268d331cb5aa9ae1f4debdf304664de552533467 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Sat, 18 Apr 2020 08:48:01 -0700
Subject: [PATCH] wgengine/magicsock: prune key.Public-keyed on peer removals

Fixes #215
---
 wgengine/magicsock/magicsock.go | 31 +++++++++++++++++++++++++++----
 wgengine/userspace.go           |  4 ++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index fc358df58..f06ea918d 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -81,6 +81,7 @@ type Conn struct {
 	endpointsUpdateActive bool
 	wantEndpointsUpdate   string // non-empty for why reason
 	lastEndpoints         []string
+	peerSet               map[key.Public]struct{}
 
 	// addrsByUDP is a map of every remote ip:port to a priority
 	// list of endpoint addresses for a peer.
@@ -95,9 +96,9 @@ type Conn struct {
 	//	10.0.0.3:3 -> [10.0.0.3:3]
 	addrsByUDP map[netaddr.IPPort]*AddrSet
 
-	// addsByKey maps from public keys (as seen by incoming DERP
+	// addrsByKey maps from public keys (as seen by incoming DERP
 	// packets) to its AddrSet (the same values as in addrsByUDP).
-	addrsByKey map[key.Public]*AddrSet // TODO(#215): remove entries when peers go away
+	addrsByKey map[key.Public]*AddrSet
 
 	netInfoFunc func(*tailcfg.NetInfo) // nil until set
 	netInfoLast *tailcfg.NetInfo
@@ -116,11 +117,11 @@ type Conn struct {
 	// home connection, or what was once our home), then we
 	// remember that route here to optimistically use instead of
 	// creating a new DERP connection back to their home.
-	derpRoute map[key.Public]derpRoute // TODO(#215): remove entries when peers go away
+	derpRoute map[key.Public]derpRoute
 
 	// peerLastDerp tracks which DERP node we last used to speak with a
 	// peer. It's only used to quiet logging, so we only log on change.
-	peerLastDerp map[key.Public]int // TODO(#215): remove entries when peers go away
+	peerLastDerp map[key.Public]int
 }
 
 // derpRoute is a route entry for a public key, saying that a certain
@@ -1240,6 +1241,28 @@ func (c *Conn) SetPrivateKey(privateKey wgcfg.PrivateKey) error {
 	return nil
 }
 
+// UpdatePeers is called when the set of WireGuard peers changes. It
+// then removes any state for old peers.
+//
+// The caller passes ownership of newPeers map to UpdatePeers.
+func (c *Conn) UpdatePeers(newPeers map[key.Public]struct{}) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	oldPeers := c.peerSet
+	c.peerSet = newPeers
+
+	// Clean up any key.Public-keyed maps for peers that no longer
+	// exist.
+	for peer := range oldPeers {
+		if _, ok := newPeers[peer]; !ok {
+			delete(c.addrsByKey, peer)
+			delete(c.derpRoute, peer)
+			delete(c.peerLastDerp, peer)
+		}
+	}
+}
+
 // SetDERPEnabled controls whether DERP is used.
 // New connections have it enabled by default.
 func (c *Conn) SetDERPEnabled(wantDerp bool) {
diff --git a/wgengine/userspace.go b/wgengine/userspace.go
index f4e080cc0..66ef65b00 100644
--- a/wgengine/userspace.go
+++ b/wgengine/userspace.go
@@ -327,10 +327,12 @@ func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, dnsDomains []string) error
 	e.wgLock.Lock()
 	defer e.wgLock.Unlock()
 
+	peerSet := make(map[key.Public]struct{}, len(cfg.Peers))
 	e.mu.Lock()
 	e.peerSequence = e.peerSequence[:0]
 	for _, p := range cfg.Peers {
 		e.peerSequence = append(e.peerSequence, p.PublicKey)
+		peerSet[key.Public(p.PublicKey)] = struct{}{}
 	}
 	e.mu.Unlock()
 
@@ -362,6 +364,8 @@ func (e *userspaceEngine) Reconfig(cfg *wgcfg.Config, dnsDomains []string) error
 		return err
 	}
 
+	e.magicConn.UpdatePeers(peerSet)
+
 	// TODO(apenwarr): only handling the first local address.
 	//   Currently we never use more than one anyway.
 	var cidr wgcfg.CIDR