From a514df71ede165061715ac4ddda7e7dedbe67d5d Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Thu, 18 Jan 2024 09:34:30 +0400
Subject: [PATCH] chore: add setDERPMap to configMaps (#11590)

Add setDERPMap
---
 tailnet/configmaps.go               | 18 ++++++
 tailnet/configmaps_internal_test.go | 89 ++++++++++++++++++++++++++++-
 tailnet/proto/compare.go            | 12 ++++
 3 files changed, 116 insertions(+), 3 deletions(-)

diff --git a/tailnet/configmaps.go b/tailnet/configmaps.go
index 2a2266913b..9f4faa76eb 100644
--- a/tailnet/configmaps.go
+++ b/tailnet/configmaps.go
@@ -253,6 +253,24 @@ func (c *configMaps) setBlockEndpoints(blockEndpoints bool) {
 	c.Broadcast()
 }
 
+// setDERPMap sets the DERP map, triggering a configuration of the engine if it has changed.
+// c.L MUST NOT be held.
+func (c *configMaps) setDERPMap(derpMap *proto.DERPMap) {
+	c.L.Lock()
+	defer c.L.Unlock()
+	eq, err := c.derpMap.Equal(derpMap)
+	if err != nil {
+		c.logger.Critical(context.Background(), "failed to compare DERP maps", slog.Error(err))
+		return
+	}
+	if eq {
+		return
+	}
+	c.derpMap = derpMap
+	c.derpMapDirty = true
+	c.Broadcast()
+}
+
 // derMapLocked returns the current DERPMap.  c.L must be held
 func (c *configMaps) derpMapLocked() *tailcfg.DERPMap {
 	m := DERPMapFromProto(c.derpMap)
diff --git a/tailnet/configmaps_internal_test.go b/tailnet/configmaps_internal_test.go
index 334bc43017..ddd3f7789e 100644
--- a/tailnet/configmaps_internal_test.go
+++ b/tailnet/configmaps_internal_test.go
@@ -571,6 +571,88 @@ func TestConfigMaps_setBlockEndpoints_same(t *testing.T) {
 	_ = testutil.RequireRecvCtx(ctx, t, done)
 }
 
+func TestConfigMaps_setDERPMap_different(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	fEng := newFakeEngineConfigurable()
+	nodePrivateKey := key.NewNode()
+	nodeID := tailcfg.NodeID(5)
+	discoKey := key.NewDisco()
+	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public(), nil)
+	defer uut.close()
+
+	derpMap := &proto.DERPMap{
+		HomeParams: &proto.DERPMap_HomeParams{RegionScore: map[int64]float64{1: 0.025}},
+		Regions: map[int64]*proto.DERPMap_Region{
+			1: {
+				RegionCode: "AUH",
+				Nodes: []*proto.DERPMap_Region_Node{
+					{Name: "AUH0"},
+				},
+			},
+		},
+	}
+	uut.setDERPMap(derpMap)
+
+	dm := testutil.RequireRecvCtx(ctx, t, fEng.setDERPMap)
+	require.Len(t, dm.HomeParams.RegionScore, 1)
+	require.Equal(t, dm.HomeParams.RegionScore[1], 0.025)
+	require.Len(t, dm.Regions, 1)
+	r1 := dm.Regions[1]
+	require.Equal(t, "AUH", r1.RegionCode)
+	require.Len(t, r1.Nodes, 1)
+	require.Equal(t, "AUH0", r1.Nodes[0].Name)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		uut.close()
+	}()
+	_ = testutil.RequireRecvCtx(ctx, t, done)
+}
+
+func TestConfigMaps_setDERPMap_same(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	fEng := newFakeEngineConfigurable()
+	nodePrivateKey := key.NewNode()
+	nodeID := tailcfg.NodeID(5)
+	discoKey := key.NewDisco()
+	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public(), nil)
+	defer uut.close()
+
+	// Given: DERP Map already set
+	derpMap := &proto.DERPMap{
+		HomeParams: &proto.DERPMap_HomeParams{RegionScore: map[int64]float64{1: 0.025}},
+		Regions: map[int64]*proto.DERPMap_Region{
+			1: {
+				RegionCode: "AUH",
+				Nodes: []*proto.DERPMap_Region_Node{
+					{Name: "AUH0"},
+				},
+			},
+		},
+	}
+	uut.L.Lock()
+	uut.derpMap = derpMap
+	uut.L.Unlock()
+
+	// Then: we don't configure
+	requireNeverConfigures(ctx, t, &uut.phased)
+
+	// When we set the same DERP map
+	uut.setDERPMap(derpMap)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		uut.close()
+	}()
+	_ = testutil.RequireRecvCtx(ctx, t, done)
+}
+
 func expectStatusWithHandshake(
 	ctx context.Context, t testing.TB, fEng *fakeEngineConfigurable, k key.NodePublic, lastHandshake time.Time,
 ) <-chan struct{} {
@@ -696,6 +778,7 @@ type fakeEngineConfigurable struct {
 	setNetworkMap chan *netmap.NetworkMap
 	reconfig      chan reconfigCall
 	filter        chan *filter.Filter
+	setDERPMap    chan *tailcfg.DERPMap
 
 	// To fake these fields the test should read from status, do stuff to the
 	// StatusBuilder, then write to statusDone
@@ -713,6 +796,7 @@ func newFakeEngineConfigurable() *fakeEngineConfigurable {
 		setNetworkMap: make(chan *netmap.NetworkMap),
 		reconfig:      make(chan reconfigCall),
 		filter:        make(chan *filter.Filter),
+		setDERPMap:    make(chan *tailcfg.DERPMap),
 		status:        make(chan *ipnstate.StatusBuilder),
 		statusDone:    make(chan struct{}),
 	}
@@ -727,9 +811,8 @@ func (f fakeEngineConfigurable) Reconfig(wg *wgcfg.Config, r *router.Config, _ *
 	return nil
 }
 
-func (fakeEngineConfigurable) SetDERPMap(*tailcfg.DERPMap) {
-	// TODO implement me
-	panic("implement me")
+func (f fakeEngineConfigurable) SetDERPMap(d *tailcfg.DERPMap) {
+	f.setDERPMap <- d
 }
 
 func (f fakeEngineConfigurable) SetFilter(flt *filter.Filter) {
diff --git a/tailnet/proto/compare.go b/tailnet/proto/compare.go
index 012ac293a0..7a2b158aa1 100644
--- a/tailnet/proto/compare.go
+++ b/tailnet/proto/compare.go
@@ -18,3 +18,15 @@ func (s *Node) Equal(o *Node) (bool, error) {
 	}
 	return bytes.Equal(sBytes, oBytes), nil
 }
+
+func (s *DERPMap) Equal(o *DERPMap) (bool, error) {
+	sBytes, err := gProto.Marshal(s)
+	if err != nil {
+		return false, err
+	}
+	oBytes, err := gProto.Marshal(o)
+	if err != nil {
+		return false, err
+	}
+	return bytes.Equal(sBytes, oBytes), nil
+}