coder · deansheather · Jul 26, 2023 · Apr 27, 2023 · May 1, 2023 · May 2, 2023
diff --git a/agent/agent.go b/agent/agent.go
@@ -27,6 +27,7 @@ import (
 	"github.com/spf13/afero"
 	"go.uber.org/atomic"
 	"golang.org/x/exp/slices"
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/xerrors"
 	"tailscale.com/net/speedtest"
 	"tailscale.com/tailcfg"
@@ -72,6 +73,7 @@ type Options struct {
 type Client interface {
 	Manifest(ctx context.Context) (agentsdk.Manifest, error)
 	Listen(ctx context.Context) (net.Conn, error)
+	DERPMapUpdates(ctx context.Context) (<-chan agentsdk.DERPMapUpdate, io.Closer, error)
 	ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
 	PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
 	PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
@@ -699,12 +701,26 @@ func (a *agent) run(ctx context.Context) error {
 		network.SetBlockEndpoints(manifest.DisableDirectConnections)
 	}
 
-	a.logger.Debug(ctx, "running tailnet connection coordinator")
-	err = a.runCoordinator(ctx, network)
-	if err != nil {
-		return xerrors.Errorf("run coordinator: %w", err)
-	}
-	return nil
+	eg, egCtx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		a.logger.Debug(egCtx, "running tailnet connection coordinator")
+		err := a.runCoordinator(egCtx, network)
+		if err != nil {
+			return xerrors.Errorf("run coordinator: %w", err)
+		}
+		return nil
+	})
+
+	eg.Go(func() error {
+		a.logger.Debug(egCtx, "running derp map subscriber")
+		err := a.runDERPMapSubscriber(egCtx, network)
+		if err != nil {
+			return xerrors.Errorf("run derp map subscriber: %w", err)
+		}
+		return nil
+	})
+
+	return eg.Wait()
 }
 
 func (a *agent) wireguardAddresses(agentID uuid.UUID) []netip.Prefix {
@@ -927,6 +943,34 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
 	}
 }
 
+// runDERPMapSubscriber runs a coordinator and returns if a reconnect should occur.
+func (a *agent) runDERPMapSubscriber(ctx context.Context, network *tailnet.Conn) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	updates, closer, err := a.client.DERPMapUpdates(ctx)
+	if err != nil {
+		return err
+	}
+	defer closer.Close()
+
+	a.logger.Info(ctx, "connected to derp map endpoint")
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case update := <-updates:
+			if update.Err != nil {
+				return update.Err
+			}
+			if update.DERPMap != nil && !tailnet.CompareDERPMaps(network.DERPMap(), update.DERPMap) {
+				a.logger.Info(ctx, "updating derp map due to detected changes")
+				network.SetDERPMap(update.DERPMap)
+			}
+		}
+	}
+}
+
 func (a *agent) runStartupScript(ctx context.Context, script string) error {
 	return a.runScript(ctx, "startup", script)
 }

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -1717,6 +1717,120 @@ 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)
+
+	originalDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
+	require.NotNil(t, originalDerpMap)
+
+	coordinator := tailnet.NewCoordinator(logger)
+	defer func() {
+		_ = coordinator.Close()
+	}()
+	agentID := uuid.New()
+	statsCh := make(chan *agentsdk.Stats, 50)
+	fs := afero.NewMemMapFs()
+	client := agenttest.NewClient(t,
+		logger.Named("agent"),
+		agentID,
+		agentsdk.Manifest{
+			DERPMap: originalDerpMap,
+			// Force DERP.
+			DisableDirectConnections: true,
+		},
+		statsCh,
+		coordinator,
+	)
+	closer := agent.New(agent.Options{
+		Client:                 client,
+		Filesystem:             fs,
+		Logger:                 logger.Named("agent"),
+		ReconnectingPTYTimeout: time.Minute,
+	})
+	defer func() {
+		_ = closer.Close()
+	}()
+
+	// Setup a client connection.
+	newClientConn := func(derpMap *tailcfg.DERPMap) *codersdk.WorkspaceAgentConn {
+		conn, err := tailnet.NewConn(&tailnet.Options{
+			Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
+			DERPMap:   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(nodes []*tailnet.Node) error {
+			return conn.UpdateNodes(nodes, false)
+		})
+		conn.SetNodeCallback(sendNode)
+		// Force DERP.
+		conn.SetBlockEndpoints(true)
+
+		sdkConn := codersdk.NewWorkspaceAgentConn(conn, codersdk.WorkspaceAgentConnOptions{
+			AgentID:   agentID,
+			CloseFunc: func() error { return codersdk.ErrSkipClose },
+		})
+		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(originalDerpMap)
+
+	// Change the DERP map.
+	newDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
+	require.NotNil(t, newDerpMap)
+
+	// Change the region ID.
+	newDerpMap.Regions[2] = newDerpMap.Regions[1]
+	delete(newDerpMap.Regions, 1)
+	newDerpMap.Regions[2].RegionID = 2
+	for _, node := range newDerpMap.Regions[2].Nodes {
+		node.RegionID = 2
+	}
+
+	// Push a new DERP map to the agent.
+	err := client.PushDERPMapUpdate(agentsdk.DERPMapUpdate{
+		DERPMap: newDerpMap,
+	})
+	require.NoError(t, err)
+
+	// Connect from a second client and make sure it uses the new DERP map.
+	conn2 := newClientConn(newDerpMap)
+	require.Equal(t, []int{2}, conn2.DERPMap().RegionIDs())
+
+	// If the first client gets a DERP map update, it should be able to
+	// reconnect just fine.
+	conn1.SetDERPMap(newDerpMap)
+	require.Equal(t, []int{2}, conn1.DERPMap().RegionIDs())
+	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...")
@@ -1940,8 +2054,8 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
 		defer close(serveClientDone)
 		coordinator.ServeClient(serverConn, uuid.New(), metadata.AgentID)
 	}()
-	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
-		return conn.UpdateNodes(node, false)
+	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(nodes []*tailnet.Node) error {
+		return conn.UpdateNodes(nodes, false)
 	})
 	conn.SetNodeCallback(sendNode)
 	agentConn := codersdk.NewWorkspaceAgentConn(conn, codersdk.WorkspaceAgentConnOptions{

diff --git a/agent/agenttest/client.go b/agent/agenttest/client.go
@@ -10,11 +10,13 @@ import (
 
 	"github.com/google/uuid"
 	"golang.org/x/exp/maps"
+	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
 	"github.com/coder/coder/tailnet"
+	"github.com/coder/coder/testutil"
 )
 
 func NewClient(t testing.TB,
@@ -28,12 +30,13 @@ func NewClient(t testing.TB,
 		manifest.AgentID = agentID
 	}
 	return &Client{
-		t:           t,
-		logger:      logger.Named("client"),
-		agentID:     agentID,
-		manifest:    manifest,
-		statsChan:   statsChan,
-		coordinator: coordinator,
+		t:              t,
+		logger:         logger.Named("client"),
+		agentID:        agentID,
+		manifest:       manifest,
+		statsChan:      statsChan,
+		coordinator:    coordinator,
+		derpMapUpdates: make(chan agentsdk.DERPMapUpdate),
 	}
 }
 
@@ -53,6 +56,7 @@ type Client struct {
 	lifecycleStates []codersdk.WorkspaceAgentLifecycle
 	startup         agentsdk.PostStartupRequest
 	logs            []agentsdk.StartupLog
+	derpMapUpdates  chan agentsdk.DERPMapUpdate
 }
 
 func (c *Client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
@@ -191,6 +195,26 @@ func (c *Client) GetServiceBanner(ctx context.Context) (codersdk.ServiceBannerCo
 	return codersdk.ServiceBannerConfig{}, nil
 }
 
+func (c *Client) PushDERPMapUpdate(update agentsdk.DERPMapUpdate) error {
+	timer := time.NewTimer(testutil.WaitShort)
+	defer timer.Stop()
+	select {
+	case c.derpMapUpdates <- update:
+	case <-timer.C:
+		return xerrors.New("timeout waiting to push derp map update")
+	}
+
+	return nil
+}
+
+func (c *Client) DERPMapUpdates(_ context.Context) (<-chan agentsdk.DERPMapUpdate, io.Closer, error) {
+	closed := make(chan struct{})
+	return c.derpMapUpdates, closeFunc(func() error {
+		close(closed)
+		return nil
+	}), nil
+}
+
 type closeFunc func() error
 
 func (c closeFunc) Close() error {

diff --git a/cli/netcheck.go b/cli/netcheck.go
@@ -26,7 +26,7 @@ func (r *RootCmd) netcheck() *clibase.Cmd {
 			ctx, cancel := context.WithTimeout(inv.Context(), 30*time.Second)
 			defer cancel()
 
-			connInfo, err := client.WorkspaceAgentConnectionInfo(ctx)
+			connInfo, err := client.WorkspaceAgentConnectionInfoGeneric(ctx)
 			if err != nil {
 				return err
 			}

diff --git a/cli/server.go b/cli/server.go
@@ -477,7 +477,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				AppHostnameRegex:            appHostnameRegex,
 				Logger:                      logger.Named("coderd"),
 				Database:                    dbfake.New(),
-				DERPMap:                     derpMap,
+				BaseDERPMap:                 derpMap,
 				Pubsub:                      pubsub.NewInMemory(),
 				CacheDir:                    cacheDir,
 				GoogleTokenValidator:        googleTokenValidator,
@@ -822,7 +822,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)
 				}