coder · deansheather · Aug 1, 2023 · Jul 31, 2023 · Jul 31, 2023 · Jul 31, 2023
diff --git a/agent/agent.go b/agent/agent.go
@@ -85,6 +85,8 @@ type Client interface {
 
 type Agent interface {
 	HTTPDebug() http.Handler
+	// TailnetConn may be nil.
+	TailnetConn() *tailnet.Conn
 	io.Closer
 }
 
@@ -200,6 +202,10 @@ type agent struct {
 	metrics            *agentMetrics
 }
 
+func (a *agent) TailnetConn() *tailnet.Conn {
+	return a.network
+}
+
 func (a *agent) init(ctx context.Context) {
 	sshSrv, err := agentssh.NewServer(ctx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, a.sshMaxTimeout, "")
 	if err != nil {

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -1818,6 +1818,15 @@ func TestAgent_UpdatedDERP(t *testing.T) {
 	})
 	require.NoError(t, err)
 
+	require.Eventually(t, func() bool {
+		conn := closer.TailnetConn()
+		if conn == nil {
+			return false
+		}
+		regionIDs := conn.DERPMap().RegionIDs()
+		return len(regionIDs) == 1 && regionIDs[0] == 2 && conn.Node().PreferredDERP == 2
+	}, testutil.WaitLong, testutil.IntervalFast)
+
 	// 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())

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -303,9 +303,21 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 		accessURL = serverURL
 	}
 
-	stunAddr, stunCleanup := stuntest.ServeWithPacketListener(t, nettype.Std{})
-	stunAddr.IP = net.ParseIP("127.0.0.1")
-	t.Cleanup(stunCleanup)
+	// If the STUNAddresses setting is empty or the default, start a STUN
+	// server. Otherwise, use the value as is.
+	var (
+		stunAddresses   []string
+		dvStunAddresses = options.DeploymentValues.DERP.Server.STUNAddresses.Value()
+	)
+	if len(dvStunAddresses) == 0 || (len(dvStunAddresses) == 1 && dvStunAddresses[0] == "stun.l.google.com:19302") {
+		stunAddr, stunCleanup := stuntest.ServeWithPacketListener(t, nettype.Std{})
+		stunAddr.IP = net.ParseIP("127.0.0.1")
+		t.Cleanup(stunCleanup)
+		stunAddresses = []string{stunAddr.String()}
+		options.DeploymentValues.DERP.Server.STUNAddresses = stunAddresses
+	} else if dvStunAddresses[0] != "disable" {
+		stunAddresses = options.DeploymentValues.DERP.Server.STUNAddresses.Value()
+	}
 
 	derpServer := derp.NewServer(key.NewNode(), tailnet.Logger(slogtest.Make(t, nil).Named("derp").Leveled(slog.LevelDebug)))
 	derpServer.SetMeshKey("test-key")
@@ -346,7 +358,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	if !options.DeploymentValues.DERP.Server.Enable.Value() {
 		region = nil
 	}
-	derpMap, err := tailnet.NewDERPMap(ctx, region, []string{stunAddr.String()}, "", "", options.DeploymentValues.DERP.Config.BlockDirect.Value())
+	derpMap, err := tailnet.NewDERPMap(ctx, region, stunAddresses, "", "", options.DeploymentValues.DERP.Config.BlockDirect.Value())
 	require.NoError(t, err)
 
 	return func(h http.Handler) {

diff --git a/coderd/healthcheck/derp.go b/coderd/healthcheck/derp.go
@@ -172,7 +172,7 @@ func (r *DERPNodeReport) derpURL() *url.URL {
 	if r.Node.HostName == "" {
 		derpURL.Host = r.Node.IPv4
 	}
-	if r.Node.DERPPort != 0 {
+	if r.Node.DERPPort != 0 && !(r.Node.DERPPort == 443 && derpURL.Scheme == "https") && !(r.Node.DERPPort == 80 && derpURL.Scheme == "http") {
 		derpURL.Host = fmt.Sprintf("%s:%d", derpURL.Host, r.Node.DERPPort)
 	}
 

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -866,15 +866,14 @@ func (api *API) derpMapUpdates(rw http.ResponseWriter, r *http.Request) {
 			lastDERPMap = derpMap
 		}
 
+		ticker.Reset(api.Options.DERPMapUpdateFrequency)
 		select {
 		case <-ctx.Done():
 			return
 		case <-api.ctx.Done():
 			return
 		case <-ticker.C:
 		}
-
-		ticker.Reset(api.Options.DERPMapUpdateFrequency)
 	}
 }
 

diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -1270,11 +1270,9 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) {
 	defer closer.Close()
 	user := coderdtest.CreateFirstUser(t, client)
 
-	originalDerpMap := api.DERPMap()
-	require.NotNil(t, originalDerpMap)
-
 	// Change the DERP mapper to our custom one.
 	var currentDerpMap atomic.Pointer[tailcfg.DERPMap]
+	originalDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
 	currentDerpMap.Store(originalDerpMap)
 	derpMapFn := func(_ *tailcfg.DERPMap) *tailcfg.DERPMap {
 		return currentDerpMap.Load().Clone()
@@ -1304,10 +1302,8 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) {
 	resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
 	agentID := resources[0].Agents[0].ID
 
-	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
-	defer cancel()
-
 	// Connect from a client.
+	ctx := testutil.Context(t, testutil.WaitLong)
 	conn1, err := client.DialWorkspaceAgent(ctx, agentID, &codersdk.DialWorkspaceAgentOptions{
 		Logger: logger.Named("client1"),
 	})
@@ -1328,7 +1324,14 @@ func TestWorkspaceAgent_UpdatedDERP(t *testing.T) {
 	currentDerpMap.Store(newDerpMap)
 
 	// Wait for the agent's DERP map to be updated.
-	// TODO: this
+	require.Eventually(t, func() bool {
+		conn := agentCloser.TailnetConn()
+		if conn == nil {
+			return false
+		}
+		regionIDs := conn.DERPMap().RegionIDs()
+		return len(regionIDs) == 1 && regionIDs[0] == 2 && conn.Node().PreferredDERP == 2
+	}, testutil.WaitLong, testutil.IntervalFast)
 
 	// Wait for the DERP map to be updated on the existing client.
 	require.Eventually(t, func() bool {

diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -578,7 +578,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 
 	if initial, changed, enabled := featureChanged(codersdk.FeatureWorkspaceProxy); shouldUpdate(initial, changed, enabled) {
 		if enabled {
-			fn := derpMapper(api.Logger, api.ProxyHealth)
+			fn := derpMapper(api.Logger, api.DeploymentValues, api.ProxyHealth)
 			api.AGPL.DERPMapper.Store(&fn)
 		} else {
 			api.AGPL.DERPMapper.Store(nil)
@@ -643,7 +643,7 @@ var (
 	lastDerpConflictLog   time.Time
 )
 
-func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*tailcfg.DERPMap) *tailcfg.DERPMap {
+func derpMapper(logger slog.Logger, cfg *codersdk.DeploymentValues, proxyHealth *proxyhealth.ProxyHealth) func(*tailcfg.DERPMap) *tailcfg.DERPMap {
 	return func(derpMap *tailcfg.DERPMap) *tailcfg.DERPMap {
 		derpMap = derpMap.Clone()
 
@@ -718,7 +718,7 @@ func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*
 					lastDerpConflictMutex.Unlock()
 					if shouldLog {
 						logger.Warn(context.Background(),
-							"proxy region ID or code conflict, ignoring workspace proxy for DERP map. Please change the flags on the affected proxy to use a different region ID and code",
+							"proxy region ID or code conflict, ignoring workspace proxy for DERP map",
 							slog.F("proxy_id", status.Proxy.ID),
 							slog.F("proxy_name", status.Proxy.Name),
 							slog.F("proxy_display_name", status.Proxy.DisplayName),
@@ -733,20 +733,43 @@ func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*
 				}
 			}
 
+			var stunNodes []*tailcfg.DERPNode
+			if !cfg.DERP.Config.BlockDirect.Value() {
+				stunNodes, err = agpltailnet.STUNNodes(regionID, cfg.DERP.Server.STUNAddresses)
+				if err != nil {
+					// Log a warning if we haven't logged one in the last
+					// minute.
+					lastDerpConflictMutex.Lock()
+					shouldLog := lastDerpConflictLog.IsZero() || time.Since(lastDerpConflictLog) > time.Minute
+					if shouldLog {
+						lastDerpConflictLog = time.Now()
+					}
+					lastDerpConflictMutex.Unlock()
+					if shouldLog {
+						logger.Error(context.Background(), "failed to calculate STUN nodes", slog.Error(err))
+					}
+
+					// No continue because we can keep going.
+					stunNodes = []*tailcfg.DERPNode{}
+				}
+			}
+
+			nodes := append(stunNodes, &tailcfg.DERPNode{
+				Name:      fmt.Sprintf("%da", regionID),
+				RegionID:  regionID,
+				HostName:  u.Hostname(),
+				DERPPort:  portInt,
+				STUNPort:  -1,
+				ForceHTTP: u.Scheme == "http",
+			})
+
 			derpMap.Regions[regionID] = &tailcfg.DERPRegion{
 				// EmbeddedRelay ONLY applies to the primary.
 				EmbeddedRelay: false,
 				RegionID:      regionID,
 				RegionCode:    regionCode,
 				RegionName:    status.Proxy.Name,
-				Nodes: []*tailcfg.DERPNode{{
-					Name:      fmt.Sprintf("%da", regionID),
-					RegionID:  regionID,
-					HostName:  u.Hostname(),
-					DERPPort:  portInt,
-					STUNPort:  -1,
-					ForceHTTP: u.Scheme == "http",
-				}},
+				Nodes:         nodes,
 			}
 		}
 

diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go
@@ -25,6 +25,7 @@ import (
 	"github.com/coder/coder/enterprise/coderd/coderdenttest"
 	"github.com/coder/coder/enterprise/coderd/license"
 	"github.com/coder/coder/provisioner/echo"
+	"github.com/coder/coder/tailnet"
 	"github.com/coder/coder/testutil"
 )
 
@@ -184,24 +185,24 @@ resourceLoop:
 		require.Equal(t, "coder_best-proxy", proxy1Region.RegionCode)
 		require.Equal(t, 10001, proxy1Region.RegionID)
 		require.False(t, proxy1Region.EmbeddedRelay)
-		require.Len(t, proxy1Region.Nodes, 1)
-		require.Equal(t, "10001a", proxy1Region.Nodes[0].Name)
-		require.Equal(t, 10001, proxy1Region.Nodes[0].RegionID)
-		require.Equal(t, proxyAPI1.Options.AccessURL.Hostname(), proxy1Region.Nodes[0].HostName)
-		require.Equal(t, proxyAPI1.Options.AccessURL.Port(), fmt.Sprint(proxy1Region.Nodes[0].DERPPort))
-		require.Equal(t, proxyAPI1.Options.AccessURL.Scheme == "http", proxy1Region.Nodes[0].ForceHTTP)
+		require.Len(t, proxy1Region.Nodes, 2) // proxy + stun
+		require.Equal(t, "10001a", proxy1Region.Nodes[1].Name)
+		require.Equal(t, 10001, proxy1Region.Nodes[1].RegionID)
+		require.Equal(t, proxyAPI1.Options.AccessURL.Hostname(), proxy1Region.Nodes[1].HostName)
+		require.Equal(t, proxyAPI1.Options.AccessURL.Port(), fmt.Sprint(proxy1Region.Nodes[1].DERPPort))
+		require.Equal(t, proxyAPI1.Options.AccessURL.Scheme == "http", proxy1Region.Nodes[1].ForceHTTP)
 
 		// The second proxy region:
 		require.Equal(t, "worst-proxy", proxy2Region.RegionName)
 		require.Equal(t, "coder_worst-proxy", proxy2Region.RegionCode)
 		require.Equal(t, 10002, proxy2Region.RegionID)
 		require.False(t, proxy2Region.EmbeddedRelay)
-		require.Len(t, proxy2Region.Nodes, 1)
-		require.Equal(t, "10002a", proxy2Region.Nodes[0].Name)
-		require.Equal(t, 10002, proxy2Region.Nodes[0].RegionID)
-		require.Equal(t, proxyAPI2.Options.AccessURL.Hostname(), proxy2Region.Nodes[0].HostName)
-		require.Equal(t, proxyAPI2.Options.AccessURL.Port(), fmt.Sprint(proxy2Region.Nodes[0].DERPPort))
-		require.Equal(t, proxyAPI2.Options.AccessURL.Scheme == "http", proxy2Region.Nodes[0].ForceHTTP)
+		require.Len(t, proxy2Region.Nodes, 2) // proxy + stun
+		require.Equal(t, "10002a", proxy2Region.Nodes[1].Name)
+		require.Equal(t, 10002, proxy2Region.Nodes[1].RegionID)
+		require.Equal(t, proxyAPI2.Options.AccessURL.Hostname(), proxy2Region.Nodes[1].HostName)
+		require.Equal(t, proxyAPI2.Options.AccessURL.Port(), fmt.Sprint(proxy2Region.Nodes[1].DERPPort))
+		require.Equal(t, proxyAPI2.Options.AccessURL.Scheme == "http", proxy2Region.Nodes[1].ForceHTTP)
 	})
 
 	t.Run("ConnectDERP", func(t *testing.T) {
@@ -239,6 +240,132 @@ resourceLoop:
 	})
 }
 
+func TestDERPMapStunNodes(t *testing.T) {
+	t.Parallel()
+
+	deploymentValues := coderdtest.DeploymentValues(t)
+	deploymentValues.Experiments = []string{
+		string(codersdk.ExperimentMoons),
+		"*",
+	}
+	stunAddresses := []string{
+		"stun.l.google.com:19302",
+		"stun1.l.google.com:19302",
+		"stun2.l.google.com:19302",
+		"stun3.l.google.com:19302",
+		"stun4.l.google.com:19302",
+	}
+	deploymentValues.DERP.Server.STUNAddresses = stunAddresses
+
+	client, closer, api, _ := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+		Options: &coderdtest.Options{
+			DeploymentValues:         deploymentValues,
+			AppHostname:              "*.primary.test.coder.com",
+			IncludeProvisionerDaemon: true,
+			RealIPConfig: &httpmw.RealIPConfig{
+				TrustedOrigins: []*net.IPNet{{
+					IP:   net.ParseIP("127.0.0.1"),
+					Mask: net.CIDRMask(8, 32),
+				}},
+				TrustedHeaders: []string{
+					"CF-Connecting-IP",
+				},
+			},
+		},
+		LicenseOptions: &coderdenttest.LicenseOptions{
+			Features: license.Features{
+				codersdk.FeatureWorkspaceProxy: 1,
+			},
+		},
+	})
+	t.Cleanup(func() {
+		_ = closer.Close()
+	})
+
+	// Create a running external proxy.
+	_ = coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
+		Name: "cool-proxy",
+	})
+
+	// Wait for both running proxies to become healthy.
+	require.Eventually(t, func() bool {
+		healthCtx := testutil.Context(t, testutil.WaitLong)
+		err := api.ProxyHealth.ForceUpdate(healthCtx)
+		if !assert.NoError(t, err) {
+			return false
+		}
+
+		regions, err := client.Regions(healthCtx)
+		if !assert.NoError(t, err) {
+			return false
+		}
+		if !assert.Len(t, regions, 2) {
+			return false
+		}
+
+		// All regions should be healthy.
+		for _, r := range regions {
+			if !r.Healthy {
+				return false
+			}
+		}
+		return true
+	}, testutil.WaitLong, testutil.IntervalMedium)
+
+	// Get the DERP map and ensure that the built-in region and the proxy region
+	// both have the STUN nodes.
+	ctx := testutil.Context(t, testutil.WaitLong)
+	connInfo, err := client.WorkspaceAgentConnectionInfoGeneric(ctx)
+	require.NoError(t, err)
+
+	// There should be two DERP servers in the map: the primary and the
+	// proxy.
+	require.NotNil(t, connInfo.DERPMap)
+	require.Len(t, connInfo.DERPMap.Regions, 2)
+
+	var (
+		primaryRegion *tailcfg.DERPRegion
+		proxyRegion   *tailcfg.DERPRegion
+	)
+	for _, r := range connInfo.DERPMap.Regions {
+		if r.EmbeddedRelay {
+			primaryRegion = r
+			continue
+		}
+		if r.RegionName == "cool-proxy" {
+			proxyRegion = r
+			continue
+		}
+
+		t.Fatalf("unexpected region: %+v", r)
+	}
+
+	// The primary region:
+	require.Equal(t, "Coder Embedded Relay", primaryRegion.RegionName)
+	require.Equal(t, "coder", primaryRegion.RegionCode)
+	require.Equal(t, 999, primaryRegion.RegionID)
+	require.True(t, primaryRegion.EmbeddedRelay)
+	require.Len(t, primaryRegion.Nodes, len(stunAddresses)+1)
+
+	// The proxy region:
+	require.Equal(t, "cool-proxy", proxyRegion.RegionName)
+	require.Equal(t, "coder_cool-proxy", proxyRegion.RegionCode)
+	require.Equal(t, 10001, proxyRegion.RegionID)
+	require.False(t, proxyRegion.EmbeddedRelay)
+	require.Len(t, proxyRegion.Nodes, len(stunAddresses)+1)
+
+	for _, region := range []*tailcfg.DERPRegion{primaryRegion, proxyRegion} {
+		stunNodes, err := tailnet.STUNNodes(region.RegionID, stunAddresses)
+		require.NoError(t, err)
+		require.Len(t, stunNodes, len(stunAddresses))
+
+		require.Equal(t, stunNodes, region.Nodes[:len(stunNodes)])
+
+		// The last node should be the Coder server.
+		require.NotZero(t, region.Nodes[len(region.Nodes)-1].DERPPort)
+	}
+}
+
 func TestDERPEndToEnd(t *testing.T) {
 	t.Parallel()