update derp map on new connection

coder · deansheather · Jul 26, 2023 · Apr 27, 2023 · May 1, 2023 · May 2, 2023
commit 404c3e4d94ec1d3b49c400871b1ba132ecc7dd03
diff --git a/agent/agent.go b/agent/agent.go
@@ -810,8 +810,14 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
 	}
 	defer coordinator.Close()
 	a.logger.Info(ctx, "connected to coordination endpoint")
-	sendNodes, errChan := tailnet.ServeCoordinator(coordinator, func(nodes []*tailnet.Node) error {
-		return network.UpdateNodes(nodes, false)
+	sendNodes, errChan := tailnet.ServeCoordinator(coordinator, func(update tailnet.CoordinatorNodeUpdate) error {
+		// Check if we need to update our DERP map.
+		if !tailnet.CompareDERPMaps(network.DERPMap(), update.DERPMap) {
+			a.logger.Info(ctx, "updating DERP map on connection request due to changes", slog.F("old", network.DERPMap()), slog.F("new", update.DERPMap))
+			network.SetDERPMap(update.DERPMap)
+		}
+
+		return network.UpdateNodes(update.Nodes, false)
 	})
 	network.SetNodeCallback(sendNodes)
 	select {

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -889,7 +889,7 @@ func TestAgent_StartupScript(t *testing.T) {
 				DERPMap:       &tailcfg.DERPMap{},
 			},
 			statsChan:   make(chan *agentsdk.Stats),
-			coordinator: tailnet.NewCoordinator(logger),
+			coordinator: tailnet.NewCoordinator(logger, emptyDerpMapFn),
 		}
 		closer := agent.New(agent.Options{
 			Client:                 client,
@@ -930,7 +930,7 @@ func TestAgent_StartupScript(t *testing.T) {
 				return codersdk.ReadBodyAsError(res)
 			},
 			statsChan:   make(chan *agentsdk.Stats),
-			coordinator: tailnet.NewCoordinator(logger),
+			coordinator: tailnet.NewCoordinator(logger, emptyDerpMapFn),
 		}
 		closer := agent.New(agent.Options{
 			Client:                 client,
@@ -1327,7 +1327,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 				ShutdownScript: "echo " + expected,
 			},
 			statsChan:   make(chan *agentsdk.Stats),
-			coordinator: tailnet.NewCoordinator(logger),
+			coordinator: tailnet.NewCoordinator(logger, emptyDerpMapFn),
 		}
 
 		fs := afero.NewMemMapFs()
@@ -1585,7 +1585,7 @@ func TestAgent_Reconnect(t *testing.T) {
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 	// After the agent is disconnected from a coordinator, it's supposed
 	// to reconnect!
-	coordinator := tailnet.NewCoordinator(logger)
+	coordinator := tailnet.NewCoordinator(logger, emptyDerpMapFn)
 	defer coordinator.Close()
 
 	agentID := uuid.New()
@@ -1623,7 +1623,7 @@ func TestAgent_Reconnect(t *testing.T) {
 func TestAgent_WriteVSCodeConfigs(t *testing.T) {
 	t.Parallel()
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-	coordinator := tailnet.NewCoordinator(logger)
+	coordinator := tailnet.NewCoordinator(logger, emptyDerpMapFn)
 	defer coordinator.Close()
 
 	client := &client{
@@ -1737,7 +1737,9 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
 	if metadata.DERPMap == nil {
 		metadata.DERPMap = tailnettest.RunDERPAndSTUN(t)
 	}
-	coordinator := tailnet.NewCoordinator(logger)
+	coordinator := tailnet.NewCoordinator(logger, func() *tailcfg.DERPMap {
+		return metadata.DERPMap
+	})
 	t.Cleanup(func() {
 		_ = coordinator.Close()
 	})
@@ -1785,8 +1787,10 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
 		defer close(serveClientDone)
 		coordinator.ServeClient(serverConn, uuid.New(), agentID)
 	}()
-	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
-		return conn.UpdateNodes(node, false)
+	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(update tailnet.CoordinatorNodeUpdate) error {
+		// Don't need to worry about updating the DERP map since it'll never
+		// change in this test (as we aren't dealing with proxies etc.)
+		return conn.UpdateNodes(update.Nodes, false)
 	})
 	conn.SetNodeCallback(sendNode)
 	agentConn := &codersdk.WorkspaceAgentConn{
@@ -2095,3 +2099,7 @@ func verifyCollectedMetrics(t *testing.T, expected []agentsdk.AgentMetric, actua
 	}
 	return true
 }
+
+func emptyDerpMapFn() *tailcfg.DERPMap {
+	return &tailcfg.DERPMap{}
+}
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -224,9 +224,6 @@ func New(options *Options) *API {
 	if options.PrometheusRegistry == nil {
 		options.PrometheusRegistry = prometheus.NewRegistry()
 	}
-	if options.TailnetCoordinator == nil {
-		options.TailnetCoordinator = tailnet.NewCoordinator(options.Logger)
-	}
 	if options.DERPServer == nil {
 		options.DERPServer = derp.NewServer(key.NewNode(), tailnet.Logger(options.Logger.Named("derp")))
 	}
@@ -314,6 +311,9 @@ func New(options *Options) *API {
 		Experiments:           experiments,
 		healthCheckGroup:      &singleflight.Group[string, *healthcheck.Report]{},
 	}
+	if options.TailnetCoordinator == nil {
+		options.TailnetCoordinator = tailnet.NewCoordinator(options.Logger, api.DERPMap)
+	}
 	if options.HealthcheckFunc == nil {
 		options.HealthcheckFunc = func(ctx context.Context, apiKey string) (*healthcheck.Report, error) {
 			return healthcheck.Run(ctx, &healthcheck.ReportOptions{

diff --git a/coderd/prometheusmetrics/prometheusmetrics_test.go b/coderd/prometheusmetrics/prometheusmetrics_test.go
@@ -300,13 +300,13 @@ func TestAgents(t *testing.T) {
 	coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 
 	// given
-	coordinator := tailnet.NewCoordinator(slogtest.Make(t, nil).Leveled(slog.LevelDebug))
-	coordinatorPtr := atomic.Pointer[tailnet.Coordinator]{}
-	coordinatorPtr.Store(&coordinator)
 	derpMap := tailnettest.RunDERPAndSTUN(t)
 	derpMapFn := func() *tailcfg.DERPMap {
 		return derpMap
 	}
+	coordinator := tailnet.NewCoordinator(slogtest.Make(t, nil).Leveled(slog.LevelDebug), derpMapFn)
+	coordinatorPtr := atomic.Pointer[tailnet.Coordinator]{}
+	coordinatorPtr.Store(&coordinator)
 	agentInactiveDisconnectTimeout := 1 * time.Hour // don't need to focus on this value in tests
 	registry := prometheus.NewRegistry()
 

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -687,8 +687,13 @@ func (api *API) dialWorkspaceAgentTailnet(agentID uuid.UUID) (*codersdk.Workspac
 		return left
 	})
 
-	sendNodes, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
-		err = conn.UpdateNodes(node, true)
+	sendNodes, _ := tailnet.ServeCoordinator(clientConn, func(update tailnet.CoordinatorNodeUpdate) error {
+		// Check if we need to update the DERP map used by the connection.
+		if !tailnet.CompareDERPMaps(conn.DERPMap(), update.DERPMap) {
+			conn.SetDERPMap(update.DERPMap)
+		}
+
+		err = conn.UpdateNodes(update.Nodes, true)
 		if err != nil {
 			return xerrors.Errorf("update nodes: %w", err)
 		}

diff --git a/coderd/wsconncache/wsconncache_test.go b/coderd/wsconncache/wsconncache_test.go
@@ -20,6 +20,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 	"go.uber.org/goleak"
+	"tailscale.com/tailcfg"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
@@ -159,7 +160,9 @@ func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Durati
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 	manifest.DERPMap = tailnettest.RunDERPAndSTUN(t)
 
-	coordinator := tailnet.NewCoordinator(logger)
+	coordinator := tailnet.NewCoordinator(logger, func() *tailcfg.DERPMap {
+		return manifest.DERPMap
+	})
 	t.Cleanup(func() {
 		_ = coordinator.Close()
 	})
@@ -190,8 +193,10 @@ func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Durati
 		_ = conn.Close()
 	})
 	go coordinator.ServeClient(serverConn, uuid.New(), agentID)
-	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
-		return conn.UpdateNodes(node, false)
+	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(update tailnet.CoordinatorNodeUpdate) error {
+		// Don't need to worry about updating the DERP map since it'll never
+		// change in this test (as we aren't dealing with proxies etc.)
+		return conn.UpdateNodes(update.Nodes, false)
 	})
 	conn.SetNodeCallback(sendNode)
 	agentConn := &codersdk.WorkspaceAgentConn{

diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -248,8 +248,13 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, opti
 				options.Logger.Debug(ctx, "failed to dial", slog.Error(err))
 				continue
 			}
-			sendNode, errChan := tailnet.ServeCoordinator(websocket.NetConn(ctx, ws, websocket.MessageBinary), func(node []*tailnet.Node) error {
-				return conn.UpdateNodes(node, false)
+			sendNode, errChan := tailnet.ServeCoordinator(websocket.NetConn(ctx, ws, websocket.MessageBinary), func(update tailnet.CoordinatorNodeUpdate) error {
+				// Check if we need to update the DERP map used by the connection.
+				if !tailnet.CompareDERPMaps(conn.DERPMap(), update.DERPMap) {
+					options.Logger.Debug(ctx, "updating DERP map on connection request due to changes", slog.F("old", conn.DERPMap()), slog.F("new", update.DERPMap))
+					conn.SetDERPMap(update.DERPMap)
+				}
+				return conn.UpdateNodes(update.Nodes, false)
 			})
 			conn.SetNodeCallback(sendNode)
 			options.Logger.Debug(ctx, "serving coordinator")

diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -403,9 +403,9 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 	}
 
 	if changed, enabled := featureChanged(codersdk.FeatureHighAvailability); changed {
-		coordinator := agpltailnet.NewCoordinator(api.Logger)
+		coordinator := agpltailnet.NewCoordinator(api.Logger, api.AGPL.DERPMap)
 		if enabled {
-			haCoordinator, err := tailnet.NewCoordinator(api.Logger, api.Pubsub)
+			haCoordinator, err := tailnet.NewCoordinator(api.Logger, api.Pubsub, api.AGPL.DERPMap)
 			if err != nil {
 				api.Logger.Error(ctx, "unable to set up high availability coordinator", slog.Error(err))
 				// If we try to setup the HA coordinator and it fails, nothing

diff --git a/enterprise/tailnet/coordinator.go b/enterprise/tailnet/coordinator.go
@@ -14,6 +14,7 @@ import (
 	"github.com/google/uuid"
 	lru "github.com/hashicorp/golang-lru/v2"
 	"golang.org/x/xerrors"
+	"tailscale.com/tailcfg"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/coderd/database"
@@ -22,7 +23,7 @@ import (
 
 // NewCoordinator creates a new high availability coordinator
 // that uses PostgreSQL pubsub to exchange handshakes.
-func NewCoordinator(logger slog.Logger, pubsub database.Pubsub) (agpl.Coordinator, error) {
+func NewCoordinator(logger slog.Logger, pubsub database.Pubsub, derpMapFn func() *tailcfg.DERPMap) (agpl.Coordinator, error) {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 
 	nameCache, err := lru.New[uuid.UUID, string](512)
@@ -34,8 +35,9 @@ func NewCoordinator(logger slog.Logger, pubsub database.Pubsub) (agpl.Coordinato
 		id:                       uuid.New(),
 		log:                      logger,
 		pubsub:                   pubsub,
-		closeFunc:                cancelFunc,
 		close:                    make(chan struct{}),
+		closeFunc:                cancelFunc,
+		derpMapFn:                derpMapFn,
 		nodes:                    map[uuid.UUID]*agpl.Node{},
 		agentSockets:             map[uuid.UUID]*agpl.TrackedConn{},
 		agentToConnectionSockets: map[uuid.UUID]map[uuid.UUID]*agpl.TrackedConn{},
@@ -57,6 +59,8 @@ type haCoordinator struct {
 	close     chan struct{}
 	closeFunc context.CancelFunc
 
+	derpMapFn func() *tailcfg.DERPMap
+
 	// nodes maps agent and connection IDs their respective node.
 	nodes map[uuid.UUID]*agpl.Node
 	// agentSockets maps agent IDs to their open websocket.
@@ -109,7 +113,10 @@ func (c *haCoordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID
 	// node of the agent. This allows the connection to establish.
 	node, ok := c.nodes[agent]
 	if ok {
-		err := tc.Enqueue([]*agpl.Node{node})
+		err := tc.Enqueue(agpl.CoordinatorNodeUpdate{
+			DERPMap: c.derpMapFn(),
+			Nodes:   []*agpl.Node{node},
+		})
 		c.mutex.Unlock()
 		if err != nil {
 			return xerrors.Errorf("enqueue node: %w", err)
@@ -177,7 +184,10 @@ func (c *haCoordinator) handleNextClientMessage(id, agent uuid.UUID, decoder *js
 		}
 		return nil
 	}
-	err = agentSocket.Enqueue([]*agpl.Node{&node})
+	err = agentSocket.Enqueue(agpl.CoordinatorNodeUpdate{
+		DERPMap: c.derpMapFn(),
+		Nodes:   []*agpl.Node{&node},
+	})
 	c.mutex.Unlock()
 	if err != nil {
 		return xerrors.Errorf("enqueu nodes: %w", err)
@@ -212,7 +222,10 @@ func (c *haCoordinator) ServeAgent(conn net.Conn, id uuid.UUID, name string) err
 	// Publish all nodes on this instance that want to connect to this agent.
 	nodes := c.nodesSubscribedToAgent(id)
 	if len(nodes) > 0 {
-		err := tc.Enqueue(nodes)
+		err := tc.Enqueue(agpl.CoordinatorNodeUpdate{
+			DERPMap: c.derpMapFn(),
+			Nodes:   nodes,
+		})
 		if err != nil {
 			c.mutex.Unlock()
 			return xerrors.Errorf("enqueue nodes: %w", err)
@@ -308,8 +321,12 @@ func (c *haCoordinator) handleAgentUpdate(id uuid.UUID, decoder *json.Decoder) (
 	}
 
 	// Publish the new node to every listening socket.
+	derpMap := c.derpMapFn()
 	for _, connectionSocket := range connectionSockets {
-		_ = connectionSocket.Enqueue([]*agpl.Node{&node})
+		_ = connectionSocket.Enqueue(agpl.CoordinatorNodeUpdate{
+			DERPMap: derpMap,
+			Nodes:   []*agpl.Node{&node},
+		})
 	}
 	c.mutex.Unlock()
 	return &node, nil
@@ -486,7 +503,10 @@ func (c *haCoordinator) handlePubsubMessage(ctx context.Context, message []byte)
 		if err != nil {
 			c.log.Error(ctx, "invalid nodes JSON", slog.F("id", agentID), slog.Error(err), slog.F("node", string(nodeJSON)))
 		}
-		err = agentSocket.Enqueue(nodes)
+		err = agentSocket.Enqueue(agpl.CoordinatorNodeUpdate{
+			DERPMap: c.derpMapFn(),
+			Nodes:   nodes,
+		})
 		if err != nil {
 			c.log.Error(ctx, "send callmemaybe to agent", slog.Error(err))
 			return