make coderd and wsproxy use websockets for derp

deansheather · deansheather · commit 68986b0fa28b · 2023-08-23T11:34:48.000Z
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -405,6 +405,7 @@ func New(options *Options) *API {
 			options.Logger,
 			options.DERPServer,
 			api.DERPMap,
+			options.DeploymentValues.DERP.Config.ForceWebSockets.Value(),
 			func(context.Context) (tailnet.MultiAgentConn, error) {
 				return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New()), nil
 			},
diff --git a/coderd/tailnet.go b/coderd/tailnet.go
@@ -45,16 +45,18 @@ func NewServerTailnet(
 	logger slog.Logger,
 	derpServer *derp.Server,
 	derpMapFn func() *tailcfg.DERPMap,
+	derpForceWebSockets bool,
 	getMultiAgent func(context.Context) (tailnet.MultiAgentConn, error),
 	cache *wsconncache.Cache,
 	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,
-		Logger:    logger,
+		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)
diff --git a/coderd/tailnet_test.go b/coderd/tailnet_test.go
@@ -232,6 +232,7 @@ func setupAgent(t *testing.T, agentAddresses []netip.Prefix) (uuid.UUID, agent.A
 		logger,
 		derpServer,
 		func() *tailcfg.DERPMap { return manifest.DERPMap },
+		false,
 		func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()), nil },
 		cache,
 		trace.NewNoopTracerProvider(),
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -734,10 +734,11 @@ func (api *API) _dialWorkspaceAgentTailnet(agentID uuid.UUID) (*codersdk.Workspa
 
 	derpMap := api.DERPMap()
 	conn, err := tailnet.NewConn(&tailnet.Options{
-		Addresses:      []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
-		DERPMap:        api.DERPMap(),
-		Logger:         api.Logger.Named("net.tailnet"),
-		BlockEndpoints: api.DeploymentValues.DERP.Config.BlockDirect.Value(),
+		Addresses:           []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
+		DERPMap:             api.DERPMap(),
+		DERPForceWebSockets: api.DeploymentValues.DERP.Config.ForceWebSockets.Value(),
+		Logger:              api.Logger.Named("net.tailnet"),
+		BlockEndpoints:      api.DeploymentValues.DERP.Config.BlockDirect.Value(),
 	})
 	if err != nil {
 		_ = clientConn.Close()
diff --git a/coderd/wsconncache/wsconncache_test.go b/coderd/wsconncache/wsconncache_test.go
@@ -179,9 +179,10 @@ func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Durati
 		_ = closer.Close()
 	})
 	conn, err := tailnet.NewConn(&tailnet.Options{
-		Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
-		DERPMap:   manifest.DERPMap,
-		Logger:    slogtest.Make(t, nil).Named("tailnet").Leveled(slog.LevelDebug),
+		Addresses:           []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
+		DERPMap:             manifest.DERPMap,
+		DERPForceWebSockets: manifest.DERPForceWebSockets,
+		Logger:              slogtest.Make(t, nil).Named("tailnet").Leveled(slog.LevelDebug),
 	})
 	require.NoError(t, err)
 	clientConn, serverConn := net.Pipe()
diff --git a/enterprise/coderd/workspaceproxy.go b/enterprise/coderd/workspaceproxy.go
@@ -717,11 +717,12 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 
 	// aReq.New = updatedProxy
 	httpapi.Write(ctx, rw, http.StatusCreated, wsproxysdk.RegisterWorkspaceProxyResponse{
-		AppSecurityKey:  api.AppSecurityKey.String(),
-		DERPMeshKey:     api.DERPServer.MeshKey(),
-		DERPRegionID:    regionID,
-		DERPMap:         api.AGPL.DERPMap(),
-		SiblingReplicas: siblingsRes,
+		AppSecurityKey:      api.AppSecurityKey.String(),
+		DERPMeshKey:         api.DERPServer.MeshKey(),
+		DERPRegionID:        regionID,
+		DERPMap:             api.AGPL.DERPMap(),
+		DERPForceWebSockets: api.DeploymentValues.DERP.Config.ForceWebSockets.Value(),
+		SiblingReplicas:     siblingsRes,
 	})
 
 	go api.forceWorkspaceProxyHealthUpdate(api.ctx)
diff --git a/enterprise/wsproxy/wsproxy.go b/enterprise/wsproxy/wsproxy.go
@@ -11,6 +11,7 @@ import (
 	"reflect"
 	"regexp"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/go-chi/chi/v5"
@@ -121,7 +122,7 @@ type Server struct {
 
 	// DERP
 	derpMesh      *derpmesh.Mesh
-	latestDERPMap *tailcfg.DERPMap
+	latestDERPMap atomic.Pointer[tailcfg.DERPMap]
 
 	// Used for graceful shutdown. Required for the dialer.
 	ctx           context.Context
@@ -247,8 +248,9 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 			s.Logger,
 			nil,
 			func() *tailcfg.DERPMap {
-				return s.latestDERPMap
+				return s.latestDERPMap.Load()
 			},
+			regResp.DERPForceWebSockets,
 			s.DialCoordinator,
 			wsconncache.New(s.DialWorkspaceAgent, 0),
 			s.TracerProvider,
@@ -455,7 +457,7 @@ func (s *Server) handleRegister(_ context.Context, res wsproxysdk.RegisterWorksp
 	}
 	s.derpMesh.SetAddresses(addresses, false)
 
-	s.latestDERPMap = res.DERPMap
+	s.latestDERPMap.Store(res.DERPMap)
 
 	return nil
 }
diff --git a/enterprise/wsproxy/wsproxysdk/wsproxysdk.go b/enterprise/wsproxy/wsproxysdk/wsproxysdk.go
@@ -207,10 +207,11 @@ type RegisterWorkspaceProxyRequest struct {
 }
 
 type RegisterWorkspaceProxyResponse struct {
-	AppSecurityKey string           `json:"app_security_key"`
-	DERPMeshKey    string           `json:"derp_mesh_key"`
-	DERPRegionID   int32            `json:"derp_region_id"`
-	DERPMap        *tailcfg.DERPMap `json:"derp_map"`
+	AppSecurityKey      string           `json:"app_security_key"`
+	DERPMeshKey         string           `json:"derp_mesh_key"`
+	DERPRegionID        int32            `json:"derp_region_id"`
+	DERPMap             *tailcfg.DERPMap `json:"derp_map"`
+	DERPForceWebSockets bool             `json:"derp_force_websockets"`
 	// SiblingReplicas is a list of all other replicas of the proxy that have
 	// not timed out.
 	SiblingReplicas []codersdk.Replica `json:"sibling_replicas"`
