chore: support adding dns hosts to tailnet.Conn (#15419)

Relates to #14718.

The remaining changes (regarding the Tailscale DNS service) will need to
be made on `coder/tailscale`.

tailnet/configmaps.go

@@ -5,7 +5,9 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"maps"
 	"net/netip"
+	"slices"
 	"sync"
 	"time"
 
@@ -14,9 +16,11 @@ import (
 	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/net/dns"
 	"tailscale.com/tailcfg"
+	"tailscale.com/types/dnstype"
 	"tailscale.com/types/ipproto"
 	"tailscale.com/types/key"
 	"tailscale.com/types/netmap"
+	"tailscale.com/util/dnsname"
 	"tailscale.com/wgengine"
 	"tailscale.com/wgengine/filter"
 	"tailscale.com/wgengine/router"
@@ -30,6 +34,10 @@ import (
 
 const lostTimeout = 15 * time.Minute
 
+// CoderDNSSuffix is the default DNS suffix that we append to Coder DNS
+// records.
+const CoderDNSSuffix = "coder."
+
 // engineConfigurable is the subset of wgengine.Engine that we use for configuration.
 //
 // This allows us to test configuration code without faking the whole interface.
@@ -63,6 +71,7 @@ type configMaps struct {
 
 	engine         engineConfigurable
 	static         netmap.NetworkMap
+	hosts          map[dnsname.FQDN][]netip.Addr
 	peers          map[uuid.UUID]*peerLifecycle
 	addresses      []netip.Prefix
 	derpMap        *tailcfg.DERPMap
@@ -79,6 +88,7 @@ func newConfigMaps(logger slog.Logger, engine engineConfigurable, nodeID tailcfg
 		phased: phased{Cond: *(sync.NewCond(&sync.Mutex{}))},
 		logger: logger,
 		engine: engine,
+		hosts:  make(map[dnsname.FQDN][]netip.Addr),
 		static: netmap.NetworkMap{
 			SelfNode: &tailcfg.Node{
 				ID:       nodeID,
@@ -153,10 +163,11 @@ func (c *configMaps) configLoop() {
 		}
 		if c.netmapDirty {
 			nm := c.netMapLocked()
+			hosts := c.hostsLocked()
 			actions = append(actions, func() {
 				c.logger.Debug(context.Background(), "updating engine network map", slog.F("network_map", nm))
 				c.engine.SetNetworkMap(nm)
-				c.reconfig(nm)
+				c.reconfig(nm, hosts)
 			})
 		}
 		if c.filterDirty {
@@ -212,6 +223,11 @@ func (c *configMaps) netMapLocked() *netmap.NetworkMap {
 	return nm
 }
 
+// hostsLocked returns the current DNS hosts mapping.  c.L must be held.
+func (c *configMaps) hostsLocked() map[dnsname.FQDN][]netip.Addr {
+	return maps.Clone(c.hosts)
+}
+
 // peerConfigLocked returns the set of peer nodes we have.  c.L must be held.
 func (c *configMaps) peerConfigLocked() []*tailcfg.Node {
 	out := make([]*tailcfg.Node, 0, len(c.peers))
@@ -261,6 +277,37 @@ 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()
+	c.hosts = make(map[dnsname.FQDN][]netip.Addr)
+	for name, addrs := range hosts {
+		c.hosts[name] = slices.Clone(addrs)
+	}
+	c.netmapDirty = true
+	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
@@ -305,7 +352,15 @@ func (c *configMaps) derpMapLocked() *tailcfg.DERPMap {
 // reconfig computes the correct wireguard config and calls the engine.Reconfig
 // with the config we have.  It is not intended for this to be called outside of
 // the updateLoop()
-func (c *configMaps) reconfig(nm *netmap.NetworkMap) {
+func (c *configMaps) reconfig(nm *netmap.NetworkMap, hosts map[dnsname.FQDN][]netip.Addr) {
+	dnsCfg := &dns.Config{}
+	if len(hosts) > 0 {
+		dnsCfg.Hosts = hosts
+		dnsCfg.OnlyIPv6 = true
+		dnsCfg.Routes = map[dnsname.FQDN][]*dnstype.Resolver{
+			CoderDNSSuffix: nil,
+		}
+	}
 	cfg, err := nmcfg.WGCfg(nm, Logger(c.logger.Named("net.wgconfig")), netmap.AllowSingleHosts, "")
 	if err != nil {
 		// WGCfg never returns an error at the time this code was written.  If it starts, returning
@@ -314,8 +369,11 @@ func (c *configMaps) reconfig(nm *netmap.NetworkMap) {
 		return
 	}
 
-	rc := &router.Config{LocalAddrs: nm.Addresses}
-	err = c.engine.Reconfig(cfg, rc, &dns.Config{}, &tailcfg.Debug{})
+	rc := &router.Config{
+		LocalAddrs: nm.Addresses,
+		Routes:     []netip.Prefix{CoderServicePrefix.AsNetip()},
+	}
+	err = c.engine.Reconfig(cfg, rc, dnsCfg, &tailcfg.Debug{})
 	if err != nil {
 		if errors.Is(err, wgengine.ErrNoChanges) {
 			return

tailnet/configmaps_internal_test.go

@@ -10,11 +10,14 @@ 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"
+	"tailscale.com/types/dnstype"
 	"tailscale.com/types/key"
 	"tailscale.com/types/netmap"
+	"tailscale.com/util/dnsname"
 	"tailscale.com/wgengine/filter"
 	"tailscale.com/wgengine/router"
 	"tailscale.com/wgengine/wgcfg"
@@ -1157,6 +1160,127 @@ func TestConfigMaps_updatePeers_nonexist(t *testing.T) {
 	}
 }
 
+func TestConfigMaps_addRemoveHosts(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())
+	defer uut.close()
+
+	addr1 := CoderServicePrefix.AddrFromUUID(uuid.New())
+	addr2 := CoderServicePrefix.AddrFromUUID(uuid.New())
+	addr3 := CoderServicePrefix.AddrFromUUID(uuid.New())
+	addr4 := CoderServicePrefix.AddrFromUUID(uuid.New())
+
+	// WHEN: we add two hosts
+	uut.addHosts(map[dnsname.FQDN][]netip.Addr{
+		"agent.myws.me.coder.": {
+			addr1,
+		},
+		"dev.main.me.coder.": {
+			addr2,
+			addr3,
+		},
+	})
+
+	// THEN: the engine is reconfigured with those same 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,
+			},
+		},
+		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.": {
+			addr4,
+		},
+		"newagent2.main.me.coder.": {
+			addr1,
+		},
+	})
+
+	// THEN: The engine is reconfigured with only the 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{
+			"newagent.myws.me.coder.": {
+				addr4,
+			},
+			"newagent2.main.me.coder.": {
+				addr1,
+			},
+		},
+		OnlyIPv6: true,
+	})
+
+	// WHEN: we remove all the hosts, and a bad host
+	uut.removeHosts(append(maps.Keys(req.dnsCfg.Hosts), "badhostname"))
+	_ = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
+	req = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
+
+	// THEN: the engine is reconfigured with an empty config
+	require.Equal(t, req.dnsCfg, &dns.Config{})
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		uut.close()
+	}()
+	_ = testutil.RequireRecvCtx(ctx, t, done)
+}
+
 func newTestNode(id int) *Node {
 	return &Node{
 		ID:            tailcfg.NodeID(id),
@@ -1199,6 +1323,7 @@ func requireNeverConfigures(ctx context.Context, t *testing.T, uut *phased) {
 type reconfigCall struct {
 	wg     *wgcfg.Config
 	router *router.Config
+	dnsCfg *dns.Config
 }
 
 var _ engineConfigurable = &fakeEngineConfigurable{}
@@ -1235,8 +1360,8 @@ func (f fakeEngineConfigurable) SetNetworkMap(networkMap *netmap.NetworkMap) {
 	f.setNetworkMap <- networkMap
 }
 
-func (f fakeEngineConfigurable) Reconfig(wg *wgcfg.Config, r *router.Config, _ *dns.Config, _ *tailcfg.Debug) error {
-	f.reconfig <- reconfigCall{wg: wg, router: r}
+func (f fakeEngineConfigurable) Reconfig(wg *wgcfg.Config, r *router.Config, dnsCfg *dns.Config, _ *tailcfg.Debug) error {
+	f.reconfig <- reconfigCall{wg: wg, router: r, dnsCfg: dnsCfg}
 	return nil
 }
 

tailnet/conn.go

@@ -33,6 +33,7 @@ import (
 	tslogger "tailscale.com/types/logger"
 	"tailscale.com/types/netlogtype"
 	"tailscale.com/types/netmap"
+	"tailscale.com/util/dnsname"
 	"tailscale.com/wgengine"
 	"tailscale.com/wgengine/capture"
 	"tailscale.com/wgengine/magicsock"
@@ -290,6 +291,7 @@ func NewConn(options *Options) (conn *Conn, err error) {
 		configMaps:      cfgMaps,
 		nodeUpdater:     nodeUp,
 		telemetrySink:   options.TelemetrySink,
+		dnsConfigurator: options.DNSConfigurator,
 		telemetryStore:  telemetryStore,
 		createdAt:       time.Now(),
 		watchCtx:        ctx,
@@ -379,6 +381,12 @@ func (p ServicePrefix) RandomPrefix() netip.Prefix {
 	return netip.PrefixFrom(p.RandomAddr(), 128)
 }
 
+func (p ServicePrefix) AsNetip() netip.Prefix {
+	out := [16]byte{}
+	copy(out[:], p[:])
+	return netip.PrefixFrom(netip.AddrFrom16(out), 48)
+}
+
 // Conn is an actively listening Wireguard connection.
 type Conn struct {
 	// Unique ID used for telemetry.
@@ -396,6 +404,7 @@ type Conn struct {
 	wireguardMonitor *netmon.Monitor
 	wireguardRouter  *router.Config
 	wireguardEngine  wgengine.Engine
+	dnsConfigurator  dns.OSConfigurator
 	listeners        map[listenKey]*listener
 	clientType       proto.TelemetryEvent_ClientType
 	createdAt        time.Time
@@ -442,6 +451,31 @@ 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 {
+		return xerrors.New("no DNSConfigurator set")
+	}
+	c.configMaps.setHosts(hosts)
+	return nil
+}
+
 func (c *Conn) SetNodeCallback(callback func(node *Node)) {
 	c.nodeUpdater.setCallback(callback)
 }