feat: set DNS hostnames in workspace updates controller (#15507)

spikecurtis · web-flow · commit 916df4d411cb · 2024-11-15T11:00:19.000+04:00
re: #14730 Adds support for the workspace updates protocol controller to also program DNS names for each agent. Right now, we only program names like `myagent.myworkspace.me.coder` and `myworkspace.coder.` (if there is exactly one agent in the workspace). We also want to support `myagent.myworkspace.username.coder.`, but for that we need to update WorkspaceUpdates RPC to also send the workspace owner's username, which will be in a separate PR.
diff --git a/tailnet/configmaps.go b/tailnet/configmaps.go
@@ -277,16 +277,6 @@ func (c *configMaps) setAddresses(ips []netip.Prefix) {
 	c.Broadcast()
 }
 
-func (c *configMaps) addHosts(hosts map[dnsname.FQDN][]netip.Addr) {
-	c.L.Lock()
-	defer c.L.Unlock()
-	for name, addrs := range hosts {
-		c.hosts[name] = slices.Clone(addrs)
-	}
-	c.netmapDirty = true
-	c.Broadcast()
-}
-
 func (c *configMaps) setHosts(hosts map[dnsname.FQDN][]netip.Addr) {
 	c.L.Lock()
 	defer c.L.Unlock()
@@ -298,16 +288,6 @@ func (c *configMaps) setHosts(hosts map[dnsname.FQDN][]netip.Addr) {
 	c.Broadcast()
 }
 
-func (c *configMaps) removeHosts(names []dnsname.FQDN) {
-	c.L.Lock()
-	defer c.L.Unlock()
-	for _, name := range names {
-		delete(c.hosts, name)
-	}
-	c.netmapDirty = true
-	c.Broadcast()
-}
-
 // setBlockEndpoints sets whether we should block configuring endpoints we learn
 // from peers.  It triggers a configuration of the engine if the value changes.
 // nolint: revive
diff --git a/tailnet/configmaps_internal_test.go b/tailnet/configmaps_internal_test.go
@@ -10,7 +10,6 @@ import (
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/exp/maps"
 	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/net/dns"
 	"tailscale.com/tailcfg"
@@ -1177,8 +1176,8 @@ func TestConfigMaps_addRemoveHosts(t *testing.T) {
 	addr3 := CoderServicePrefix.AddrFromUUID(uuid.New())
 	addr4 := CoderServicePrefix.AddrFromUUID(uuid.New())
 
-	// WHEN: we add two hosts
-	uut.addHosts(map[dnsname.FQDN][]netip.Addr{
+	// WHEN: we set two hosts
+	uut.setHosts(map[dnsname.FQDN][]netip.Addr{
 		"agent.myws.me.coder.": {
 			addr1,
 		},
@@ -1207,36 +1206,6 @@ func TestConfigMaps_addRemoveHosts(t *testing.T) {
 		OnlyIPv6: true,
 	})
 
-	// WHEN: we add a new host
-	newHost := map[dnsname.FQDN][]netip.Addr{
-		"agent2.myws.me.coder.": {
-			addr4,
-		},
-	}
-	uut.addHosts(newHost)
-
-	// THEN: the engine is reconfigured with both the old and new hosts
-	_ = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
-	req = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
-	require.Equal(t, req.dnsCfg, &dns.Config{
-		Routes: map[dnsname.FQDN][]*dnstype.Resolver{
-			CoderDNSSuffix: nil,
-		},
-		Hosts: map[dnsname.FQDN][]netip.Addr{
-			"agent.myws.me.coder.": {
-				addr1,
-			},
-			"dev.main.me.coder.": {
-				addr2,
-				addr3,
-			},
-			"agent2.myws.me.coder.": {
-				addr4,
-			},
-		},
-		OnlyIPv6: true,
-	})
-
 	// WHEN: We replace the hosts with a new set
 	uut.setHosts(map[dnsname.FQDN][]netip.Addr{
 		"newagent.myws.me.coder.": {
@@ -1265,8 +1234,8 @@ func TestConfigMaps_addRemoveHosts(t *testing.T) {
 		OnlyIPv6: true,
 	})
 
-	// WHEN: we remove all the hosts, and a bad host
-	uut.removeHosts(append(maps.Keys(req.dnsCfg.Hosts), "badhostname"))
+	// WHEN: we remove all the hosts
+	uut.setHosts(map[dnsname.FQDN][]netip.Addr{})
 	_ = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
 	req = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
 
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -451,22 +451,6 @@ func (c *Conn) SetAddresses(ips []netip.Prefix) error {
 	return nil
 }
 
-func (c *Conn) AddDNSHosts(hosts map[dnsname.FQDN][]netip.Addr) error {
-	if c.dnsConfigurator == nil {
-		return xerrors.New("no DNSConfigurator set")
-	}
-	c.configMaps.addHosts(hosts)
-	return nil
-}
-
-func (c *Conn) RemoveDNSHosts(names []dnsname.FQDN) error {
-	if c.dnsConfigurator == nil {
-		return xerrors.New("no DNSConfigurator set")
-	}
-	c.configMaps.removeHosts(names)
-	return nil
-}
-
 // SetDNSHosts replaces the map of DNS hosts for the connection.
 func (c *Conn) SetDNSHosts(hosts map[dnsname.FQDN][]netip.Addr) error {
 	if c.dnsConfigurator == nil {
diff --git a/tailnet/controllers.go b/tailnet/controllers.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"maps"
 	"math"
+	"net/netip"
 	"strings"
 	"sync"
 	"time"
@@ -15,6 +16,7 @@ import (
 	"storj.io/drpc"
 	"storj.io/drpc/drpcerr"
 	"tailscale.com/tailcfg"
+	"tailscale.com/util/dnsname"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/codersdk"
@@ -104,6 +106,12 @@ type WorkspaceUpdatesController interface {
 	New(WorkspaceUpdatesClient) CloserWaiter
 }
 
+// DNSHostsSetter is something that you can set a mapping of DNS names to IPs on. It's the subset
+// of the tailnet.Conn that we use to configure DNS records.
+type DNSHostsSetter interface {
+	SetDNSHosts(hosts map[dnsname.FQDN][]netip.Addr) error
+}
+
 // ControlProtocolClients represents an abstract interface to the tailnet control plane via a set
 // of protocol clients. The Closer should close all the clients (e.g. by closing the underlying
 // connection).
@@ -835,8 +843,9 @@ func (r *basicResumeTokenRefresher) refresh() {
 }
 
 type tunnelAllWorkspaceUpdatesController struct {
-	coordCtrl *TunnelSrcCoordController
-	logger    slog.Logger
+	coordCtrl     *TunnelSrcCoordController
+	dnsHostSetter DNSHostsSetter
+	logger        slog.Logger
 }
 
 type workspace struct {
@@ -845,30 +854,48 @@ type workspace struct {
 	agents map[uuid.UUID]agent
 }
 
+// addAllDNSNames adds names for all of its agents to the given map of names
+func (w workspace) addAllDNSNames(names map[dnsname.FQDN][]netip.Addr) error {
+	for _, a := range w.agents {
+		// TODO: technically, DNS labels cannot start with numbers, but the rules are often not
+		//       strictly enforced.
+		// TODO: support <agent>.<workspace>.<username>.coder
+		fqdn, err := dnsname.ToFQDN(fmt.Sprintf("%s.%s.me.coder.", a.name, w.name))
+		if err != nil {
+			return err
+		}
+		names[fqdn] = []netip.Addr{CoderServicePrefix.AddrFromUUID(a.id)}
+	}
+	// TODO: Possibly support <workspace>.coder. alias if there is only one agent
+	return nil
+}
+
 type agent struct {
 	id   uuid.UUID
 	name string
 }
 
 func (t *tunnelAllWorkspaceUpdatesController) New(client WorkspaceUpdatesClient) CloserWaiter {
 	updater := &tunnelUpdater{
-		client:       client,
-		errChan:      make(chan error, 1),
-		logger:       t.logger,
-		coordCtrl:    t.coordCtrl,
-		recvLoopDone: make(chan struct{}),
-		workspaces:   make(map[uuid.UUID]*workspace),
+		client:         client,
+		errChan:        make(chan error, 1),
+		logger:         t.logger,
+		coordCtrl:      t.coordCtrl,
+		dnsHostsSetter: t.dnsHostSetter,
+		recvLoopDone:   make(chan struct{}),
+		workspaces:     make(map[uuid.UUID]*workspace),
 	}
 	go updater.recvLoop()
 	return updater
 }
 
 type tunnelUpdater struct {
-	errChan      chan error
-	logger       slog.Logger
-	client       WorkspaceUpdatesClient
-	coordCtrl    *TunnelSrcCoordController
-	recvLoopDone chan struct{}
+	errChan        chan error
+	logger         slog.Logger
+	client         WorkspaceUpdatesClient
+	coordCtrl      *TunnelSrcCoordController
+	dnsHostsSetter DNSHostsSetter
+	recvLoopDone   chan struct{}
 
 	// don't need the mutex since only manipulated by the recvLoop
 	workspaces map[uuid.UUID]*workspace
@@ -991,6 +1018,16 @@ func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error {
 	}
 	allAgents := t.allAgentIDs()
 	t.coordCtrl.SyncDestinations(allAgents)
+	if t.dnsHostsSetter != nil {
+		t.logger.Debug(context.Background(), "updating dns hosts")
+		dnsNames := t.allDNSNames()
+		err := t.dnsHostsSetter.SetDNSHosts(dnsNames)
+		if err != nil {
+			return xerrors.Errorf("failed to set DNS hosts: %w", err)
+		}
+	} else {
+		t.logger.Debug(context.Background(), "skipping setting DNS names because we have no setter")
+	}
 	return nil
 }
 
@@ -1035,10 +1072,30 @@ func (t *tunnelUpdater) allAgentIDs() []uuid.UUID {
 	return out
 }
 
+func (t *tunnelUpdater) allDNSNames() map[dnsname.FQDN][]netip.Addr {
+	names := make(map[dnsname.FQDN][]netip.Addr)
+	for _, w := range t.workspaces {
+		err := w.addAllDNSNames(names)
+		if err != nil {
+			// This should never happen in production, because converting the FQDN only fails
+			// if names are too long, and we put strict length limits on agent, workspace, and user
+			// names.
+			t.logger.Critical(context.Background(),
+				"failed to include DNS name(s)",
+				slog.F("workspace_id", w.id),
+				slog.Error(err))
+		}
+	}
+	return names
+}
+
+// NewTunnelAllWorkspaceUpdatesController creates a WorkspaceUpdatesController that creates tunnels
+// (via the TunnelSrcCoordController) to all agents received over the WorkspaceUpdates RPC. If a
+// DNSHostSetter is provided, it also programs DNS hosts based on the agent and workspace names.
 func NewTunnelAllWorkspaceUpdatesController(
-	logger slog.Logger, c *TunnelSrcCoordController,
+	logger slog.Logger, c *TunnelSrcCoordController, d DNSHostsSetter,
 ) WorkspaceUpdatesController {
-	return &tunnelAllWorkspaceUpdatesController{logger: logger, coordCtrl: c}
+	return &tunnelAllWorkspaceUpdatesController{logger: logger, coordCtrl: c, dnsHostSetter: d}
 }
 
 // NewController creates a new Controller without running it
diff --git a/tailnet/controllers_test.go b/tailnet/controllers_test.go
