From a7e7c7b548dd7903326360b4d9740a0412779433 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 9 Apr 2020 14:21:36 -0700
Subject: [PATCH] wgengine/magicsock: close derp connections on rebind

Fixes #276

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 wgengine/magicsock/magicsock.go | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 2264b69fa..8fa2670ee 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -479,12 +479,21 @@ func (c *Conn) setNearestDERP(derpNum int) (wantDERP bool) {
 	for i, ad := range c.activeDerp {
 		go ad.c.NotePreferred(i == c.myDerp)
 	}
-	if derpNum != 0 {
-		go c.derpWriteChanOfAddr(&net.UDPAddr{IP: derpMagicIP, Port: derpNum}, key.Public{})
-	}
+	c.goDerpConnect(derpNum)
 	return true
 }
 
+// goDerpConnect starts a goroutine to start connecting to the given
+// DERP node.
+//
+// c.mu may be held, but does not need to be.
+func (c *Conn) goDerpConnect(node int) {
+	if node == 0 {
+		return
+	}
+	go c.derpWriteChanOfAddr(&net.UDPAddr{IP: derpMagicIP, Port: node}, key.Public{})
+}
+
 // determineEndpoints returns the machine's endpoint addresses. It
 // does a STUN lookup (via netcheck) to determine its public address.
 //
@@ -1226,11 +1235,9 @@ func (c *Conn) SetPrivateKey(privateKey wgcfg.PrivateKey) error {
 	c.closeAllDerpLocked("new-private-key")
 
 	// Key changed. Close existing DERP connections and reconnect to home.
-	oldHome := c.myDerp
-	c.myDerp = 0 // zero it, so setNearestDERP will do work
-	if oldHome != 0 {
-		c.logf("magicsock: private key set, rebooting connection to home derp-%d", oldHome)
-		go c.setNearestDERP(oldHome)
+	if c.myDerp != 0 {
+		c.logf("magicsock: private key changed, reconnecting to home derp-%d", c.myDerp)
+		c.goDerpConnect(c.myDerp)
 	}
 
 	return nil
@@ -1478,6 +1485,11 @@ func (c *Conn) Rebind() {
 		return
 	}
 	c.pconn4.Reset(packetConn.(*net.UDPConn))
+
+	c.mu.Lock()
+	c.closeAllDerpLocked("rebind")
+	c.mu.Unlock()
+	c.goDerpConnect(c.myDerp)
 	c.resetAddrSetStates()
 }