From c500a49917473fe4ee058650e9a56ff85970cc3b Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Fri, 31 Mar 2023 14:14:06 +0000
Subject: [PATCH] fix: delete old addresses from the active derp mesh mapping

---
 enterprise/derpmesh/derpmesh.go      |  1 +
 enterprise/derpmesh/derpmesh_test.go | 54 ++++++++++++++++++++++++++--
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/enterprise/derpmesh/derpmesh.go b/enterprise/derpmesh/derpmesh.go
index 3982542167073..f98dbb30d0ef4 100644
--- a/enterprise/derpmesh/derpmesh.go
+++ b/enterprise/derpmesh/derpmesh.go
@@ -137,6 +137,7 @@ func (m *Mesh) removeAddress(address string) bool {
 	cancelFunc, isActive := m.active[address]
 	if isActive {
 		cancelFunc()
+		delete(m.active, address)
 	}
 	return isActive
 }
diff --git a/enterprise/derpmesh/derpmesh_test.go b/enterprise/derpmesh/derpmesh_test.go
index 7c7331d21c355..a4a7373c9aab0 100644
--- a/enterprise/derpmesh/derpmesh_test.go
+++ b/enterprise/derpmesh/derpmesh_test.go
@@ -50,9 +50,9 @@ func TestDERPMesh(t *testing.T) {
 		defer firstServer.Close()
 		secondServer, secondServerURL := startDERP(t, tlsConfig)
 		firstMesh := derpmesh.New(slogtest.Make(t, nil).Named("first").Leveled(slog.LevelDebug), firstServer, tlsConfig)
-		firstMesh.SetAddresses([]string{secondServerURL}, true)
+		firstMesh.SetAddresses([]string{secondServerURL}, false)
 		secondMesh := derpmesh.New(slogtest.Make(t, nil).Named("second").Leveled(slog.LevelDebug), secondServer, tlsConfig)
-		secondMesh.SetAddresses([]string{firstServerURL}, true)
+		secondMesh.SetAddresses([]string{firstServerURL}, false)
 		defer firstMesh.Close()
 		defer secondMesh.Close()
 
@@ -183,6 +183,56 @@ func TestDERPMesh(t *testing.T) {
 			}
 		}()
 
+		got := recvData(t, secondClient)
+		require.Equal(t, sent, got)
+		cancelFunc()
+		<-closed
+	})
+	t.Run("CycleAddresses", func(t *testing.T) {
+		t.Parallel()
+		firstServer, firstServerURL := startDERP(t, tlsConfig)
+		defer firstServer.Close()
+		secondServer, secondServerURL := startDERP(t, tlsConfig)
+		firstMesh := derpmesh.New(slogtest.Make(t, nil).Named("first").Leveled(slog.LevelDebug), firstServer, tlsConfig)
+		firstMesh.SetAddresses([]string{secondServerURL}, false)
+		secondMesh := derpmesh.New(slogtest.Make(t, nil).Named("second").Leveled(slog.LevelDebug), secondServer, tlsConfig)
+		// Ensures that the client properly re-adds the address after it's removed.
+		secondMesh.SetAddresses([]string{firstServerURL}, true)
+		secondMesh.SetAddresses([]string{}, true)
+		secondMesh.SetAddresses([]string{firstServerURL}, true)
+		defer firstMesh.Close()
+		defer secondMesh.Close()
+
+		first := key.NewNode()
+		second := key.NewNode()
+		firstClient, err := derphttp.NewClient(first, secondServerURL, tailnet.Logger(slogtest.Make(t, nil)))
+		require.NoError(t, err)
+		firstClient.TLSConfig = tlsConfig
+		secondClient, err := derphttp.NewClient(second, firstServerURL, tailnet.Logger(slogtest.Make(t, nil)))
+		require.NoError(t, err)
+		secondClient.TLSConfig = tlsConfig
+		err = secondClient.Connect(context.Background())
+		require.NoError(t, err)
+
+		closed := make(chan struct{})
+		ctx, cancelFunc := context.WithCancel(context.Background())
+		defer cancelFunc()
+		sent := []byte("hello world")
+		go func() {
+			defer close(closed)
+			ticker := time.NewTicker(50 * time.Millisecond)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-ticker.C:
+				}
+				err = firstClient.Send(second.Public(), sent)
+				require.NoError(t, err)
+			}
+		}()
+
 		got := recvData(t, secondClient)
 		require.Equal(t, sent, got)
 		cancelFunc()