coder · deansheather · Jul 26, 2023 · Apr 27, 2023 · May 1, 2023 · May 2, 2023
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()
@@ -1564,6 +1564,109 @@ func TestAgent_Dial(t *testing.T) {
 	}
 }
 
+// TestAgent_UpdatedDERP checks that agents can handle their DERP map being
+// updated, and that clients can also handle it.
+func TestAgent_UpdatedDERP(t *testing.T) {
+	t.Parallel()
+
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+
+	derpMap := *tailnettest.RunDERPAndSTUN(t)
+	metadata := agentsdk.Manifest{
+		DERPMap: &derpMap,
+	}
+	coordinator := tailnet.NewCoordinator(logger, func() *tailcfg.DERPMap {
+		return &derpMap
+	})
+	defer func() {
+		_ = coordinator.Close()
+	}()
+	agentID := uuid.New()
+	statsCh := make(chan *agentsdk.Stats, 50)
+	fs := afero.NewMemMapFs()
+	c := &client{
+		t:           t,
+		agentID:     agentID,
+		manifest:    metadata,
+		statsChan:   statsCh,
+		coordinator: coordinator,
+	}
+	closer := agent.New(agent.Options{
+		Client:                 c,
+		Filesystem:             fs,
+		Logger:                 logger.Named("agent"),
+		ReconnectingPTYTimeout: time.Minute,
+	})
+	defer func() {
+		_ = closer.Close()
+	}()
+
+	// Setup a client connection.
+	newClientConn := func() *codersdk.WorkspaceAgentConn {
+		conn, err := tailnet.NewConn(&tailnet.Options{
+			Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
+			DERPMap:   metadata.DERPMap,
+			Logger:    logger.Named("client"),
+		})
+		require.NoError(t, err)
+		clientConn, serverConn := net.Pipe()
+		serveClientDone := make(chan struct{})
+		t.Cleanup(func() {
+			_ = clientConn.Close()
+			_ = serverConn.Close()
+			_ = conn.Close()
+			<-serveClientDone
+		})
+		go func() {
+			defer close(serveClientDone)
+			err := coordinator.ServeClient(serverConn, uuid.New(), agentID)
+			assert.NoError(t, err)
+		}()
+		sendNode, _ := tailnet.ServeCoordinator(clientConn, func(update tailnet.CoordinatorNodeUpdate) error {
+			if tailnet.CompareDERPMaps(conn.DERPMap(), update.DERPMap) {
+				conn.SetDERPMap(update.DERPMap)
+			}
+			return conn.UpdateNodes(update.Nodes, false)
+		})
+		conn.SetNodeCallback(sendNode)
+
+		sdkConn := &codersdk.WorkspaceAgentConn{
+			Conn: conn,
+		}
+		t.Cleanup(func() {
+			_ = sdkConn.Close()
+		})
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+		if !sdkConn.AwaitReachable(ctx) {
+			t.Fatal("agent not reachable")
+		}
+
+		return sdkConn
+	}
+	conn1 := newClientConn()
+
+	// Change the DERP map.
+	derpMap = *tailnettest.RunDERPAndSTUN(t)
+	// Change the region ID.
+	derpMap.Regions[2] = derpMap.Regions[1]
+	delete(derpMap.Regions, 1)
+	derpMap.Regions[2].RegionID = 2
+	for _, node := range derpMap.Regions[2].Nodes {
+		node.RegionID = 2
+	}
+
+	// Connect from a second client and make sure it uses the new DERP map.
+	conn2 := newClientConn()
+	require.Equal(t, []int{2}, conn2.DERPMap().RegionIDs())
+
+	// Make sure the first client is still reachable (it received a a DERP map
+	// update when the agent's nodes changed).
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+	require.True(t, conn1.AwaitReachable(ctx))
+}
+
 func TestAgent_Speedtest(t *testing.T) {
 	t.Parallel()
 	t.Skip("This test is relatively flakey because of Tailscale's speedtest code...")
@@ -1585,7 +1688,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 +1726,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 +1840,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 +1890,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 +2202,7 @@ func verifyCollectedMetrics(t *testing.T, expected []agentsdk.AgentMetric, actua
 	}
 	return true
 }
+
+func emptyDerpMapFn() *tailcfg.DERPMap {
+	return &tailcfg.DERPMap{}
+}
diff --git a/cli/server.go b/cli/server.go
@@ -766,7 +766,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 
 			if cfg.Prometheus.Enable {
 				// Agent metrics require reference to the tailnet coordinator, so must be initiated after Coder API.
-				closeAgentsFunc, err := prometheusmetrics.Agents(ctx, logger, options.PrometheusRegistry, coderAPI.Database, &coderAPI.TailnetCoordinator, options.DERPMap, coderAPI.Options.AgentInactiveDisconnectTimeout, 0)
+				closeAgentsFunc, err := prometheusmetrics.Agents(ctx, logger, options.PrometheusRegistry, coderAPI.Database, &coderAPI.TailnetCoordinator, coderAPI.DERPMap, coderAPI.Options.AgentInactiveDisconnectTimeout, 0)
 				if err != nil {
 					return xerrors.Errorf("register agents prometheus metric: %w", err)
 				}

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go