fix: change servertailnet to register the DERP dialer before setting DERP map (coder#12137)

spikecurtis · web-flow · commit 2d0b9106c0d1 · 2024-02-15T10:51:12.000+04:00
I noticed a possible race where tailnet.Conn can try to dial the embedded region before we've set our custom dialer that send the DERP in-memory.  This closes that race and adds a test case for servertailnet with no STUN and an embedded relay
diff --git a/coderd/tailnet.go b/coderd/tailnet.go
@@ -52,19 +52,41 @@ func NewServerTailnet(
 	traceProvider trace.TracerProvider,
 ) (*ServerTailnet, error) {
 	logger = logger.Named("servertailnet")
-	originalDerpMap := derpMapFn()
 	conn, err := tailnet.NewConn(&tailnet.Options{
 		Addresses:           []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
-		DERPMap:             originalDerpMap,
 		DERPForceWebSockets: derpForceWebSockets,
 		Logger:              logger,
 	})
 	if err != nil {
 		return nil, xerrors.Errorf("create tailnet conn: %w", err)
 	}
-
 	serverCtx, cancel := context.WithCancel(ctx)
+
+	// This is set to allow local DERP traffic to be proxied through memory
+	// instead of needing to hit the external access URL. Don't use the ctx
+	// given in this callback, it's only valid while connecting.
+	if derpServer != nil {
+		conn.SetDERPRegionDialer(func(_ context.Context, region *tailcfg.DERPRegion) net.Conn {
+			if !region.EmbeddedRelay {
+				return nil
+			}
+			logger.Debug(ctx, "connecting to embedded DERP via in-memory pipe")
+			left, right := net.Pipe()
+			go func() {
+				defer left.Close()
+				defer right.Close()
+				brw := bufio.NewReadWriter(bufio.NewReader(right), bufio.NewWriter(right))
+				derpServer.Accept(ctx, right, brw, "internal")
+			}()
+			return left
+		})
+	}
+
 	derpMapUpdaterClosed := make(chan struct{})
+	originalDerpMap := derpMapFn()
+	// it's important to set the DERPRegionDialer above _before_ we set the DERP map so that if
+	// there is an embedded relay, we use the local in-memory dialer.
+	conn.SetDERPMap(originalDerpMap)
 	go func() {
 		defer close(derpMapUpdaterClosed)
 
@@ -139,25 +161,6 @@ func NewServerTailnet(
 	// registering the callback also triggers send of the initial node
 	tn.coordinatee.SetNodeCallback(tn.nodeCallback)
 
-	// This is set to allow local DERP traffic to be proxied through memory
-	// instead of needing to hit the external access URL. Don't use the ctx
-	// given in this callback, it's only valid while connecting.
-	if derpServer != nil {
-		conn.SetDERPRegionDialer(func(_ context.Context, region *tailcfg.DERPRegion) net.Conn {
-			if !region.EmbeddedRelay {
-				return nil
-			}
-			left, right := net.Pipe()
-			go func() {
-				defer left.Close()
-				defer right.Close()
-				brw := bufio.NewReadWriter(bufio.NewReader(right), bufio.NewWriter(right))
-				derpServer.Accept(ctx, right, brw, "internal")
-			}()
-			return left
-		})
-	}
-
 	go tn.watchAgentUpdates()
 	go tn.expireOldAgents()
 	return tn, nil
diff --git a/coderd/tailnet_test.go b/coderd/tailnet_test.go
@@ -50,6 +50,24 @@ func TestServerTailnet_AgentConn_OK(t *testing.T) {
 	assert.True(t, conn.AwaitReachable(ctx))
 }
 
+func TestServerTailnet_AgentConn_NoSTUN(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
+	defer cancel()
+
+	// Connect through the ServerTailnet
+	agents, serverTailnet := setupServerTailnetAgent(t, 1,
+		tailnettest.DisableSTUN, tailnettest.DERPIsEmbedded)
+	a := agents[0]
+
+	conn, release, err := serverTailnet.AgentConn(ctx, a.id)
+	require.NoError(t, err)
+	defer release()
+
+	assert.True(t, conn.AwaitReachable(ctx))
+}
+
 func TestServerTailnet_ReverseProxy(t *testing.T) {
 	t.Parallel()
 
@@ -311,9 +329,9 @@ type agentWithID struct {
 	agent.Agent
 }
 
-func setupServerTailnetAgent(t *testing.T, agentNum int) ([]agentWithID, *coderd.ServerTailnet) {
+func setupServerTailnetAgent(t *testing.T, agentNum int, opts ...tailnettest.DERPAndStunOption) ([]agentWithID, *coderd.ServerTailnet) {
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-	derpMap, derpServer := tailnettest.RunDERPAndSTUN(t)
+	derpMap, derpServer := tailnettest.RunDERPAndSTUN(t, opts...)
 
 	coord := tailnet.NewCoordinator(logger)
 	t.Cleanup(func() {
diff --git a/tailnet/configmaps.go b/tailnet/configmaps.go
@@ -204,7 +204,8 @@ func (c *configMaps) netMapLocked() *netmap.NetworkMap {
 	nm.Addresses = make([]netip.Prefix, len(c.addresses))
 	copy(nm.Addresses, c.addresses)
 
-	nm.DERPMap = c.derpMap.Clone()
+	// we don't need to set the DERPMap in the network map because we separately
+	// send the DERPMap directly via SetDERPMap
 	nm.Peers = c.peerConfigLocked()
 	nm.SelfNode.Addresses = nm.Addresses
 	nm.SelfNode.AllowedIPs = nm.Addresses
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -116,9 +116,6 @@ func NewConn(options *Options) (conn *Conn, err error) {
 	if len(options.Addresses) == 0 {
 		return nil, xerrors.New("At least one IP range must be provided")
 	}
-	if options.DERPMap == nil {
-		return nil, xerrors.New("DERPMap must be provided")
-	}
 
 	nodePrivateKey := key.NewNode()
 	var nodeID tailcfg.NodeID
@@ -219,7 +216,9 @@ func NewConn(options *Options) (conn *Conn, err error) {
 		magicConn.DiscoPublicKey(),
 	)
 	cfgMaps.setAddresses(options.Addresses)
-	cfgMaps.setDERPMap(options.DERPMap)
+	if options.DERPMap != nil {
+		cfgMaps.setDERPMap(options.DERPMap)
+	}
 	cfgMaps.setBlockEndpoints(options.BlockEndpoints)
 
 	nodeUp := newNodeUpdater(
diff --git a/tailnet/tailnettest/tailnettest.go b/tailnet/tailnettest/tailnettest.go
@@ -32,22 +32,47 @@ import (
 //go:generate mockgen -destination ./coordinatormock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinator
 //go:generate mockgen -destination ./coordinateemock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinatee
 
+type derpAndSTUNCfg struct {
+	DisableSTUN    bool
+	DERPIsEmbedded bool
+}
+
+type DERPAndStunOption func(cfg *derpAndSTUNCfg)
+
+func DisableSTUN(cfg *derpAndSTUNCfg) {
+	cfg.DisableSTUN = true
+}
+
+func DERPIsEmbedded(cfg *derpAndSTUNCfg) {
+	cfg.DERPIsEmbedded = true
+}
+
 // RunDERPAndSTUN creates a DERP mapping for tests.
-func RunDERPAndSTUN(t *testing.T) (*tailcfg.DERPMap, *derp.Server) {
+func RunDERPAndSTUN(t *testing.T, opts ...DERPAndStunOption) (*tailcfg.DERPMap, *derp.Server) {
+	cfg := new(derpAndSTUNCfg)
+	for _, o := range opts {
+		o(cfg)
+	}
 	logf := tailnet.Logger(slogtest.Make(t, nil))
 	d := derp.NewServer(key.NewNode(), logf)
 	server := httptest.NewUnstartedServer(derphttp.Handler(d))
 	server.Config.ErrorLog = tslogger.StdLogger(logf)
 	server.Config.TLSNextProto = make(map[string]func(*http.Server, *tls.Conn, http.Handler))
 	server.StartTLS()
-
-	stunAddr, stunCleanup := stuntest.ServeWithPacketListener(t, nettype.Std{})
 	t.Cleanup(func() {
 		server.CloseClientConnections()
 		server.Close()
 		d.Close()
-		stunCleanup()
 	})
+
+	stunPort := -1
+	if !cfg.DisableSTUN {
+		stunAddr, stunCleanup := stuntest.ServeWithPacketListener(t, nettype.Std{})
+		t.Cleanup(func() {
+			stunCleanup()
+		})
+		stunPort = stunAddr.Port
+	}
 	tcpAddr, ok := server.Listener.Addr().(*net.TCPAddr)
 	if !ok {
 		t.FailNow()
@@ -65,11 +90,12 @@ func RunDERPAndSTUN(t *testing.T) (*tailcfg.DERPMap, *derp.Server) {
 						RegionID:         1,
 						IPv4:             "127.0.0.1",
 						IPv6:             "none",
-						STUNPort:         stunAddr.Port,
+						STUNPort:         stunPort,
 						DERPPort:         tcpAddr.Port,
 						InsecureForTests: true,
 					},
 				},
+				EmbeddedRelay: cfg.DERPIsEmbedded,
 			},
 		},
 	}, d
