chore: add derpserver to wsproxy, add proxies to derpmap (#7311)

enterprise/wsproxy/wsproxy.go

@@ -2,9 +2,12 @@ package wsproxy
 
 import (
 	"context"
+	"crypto/tls"
+	"crypto/x509"
 	"fmt"
 	"net/http"
 	"net/url"
+	"os"
 	"reflect"
 	"regexp"
 	"strings"
@@ -12,9 +15,13 @@ import (
 
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
+	"github.com/hashicorp/go-multierror"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/xerrors"
+	"tailscale.com/derp"
+	"tailscale.com/derp/derphttp"
+	"tailscale.com/types/key"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/buildinfo"
@@ -25,9 +32,10 @@ import (
 	"github.com/coder/coder/coderd/workspaceapps"
 	"github.com/coder/coder/coderd/wsconncache"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/enterprise/derpmesh"
 	"github.com/coder/coder/enterprise/wsproxy/wsproxysdk"
 	"github.com/coder/coder/site"
-	agpl "github.com/coder/coder/tailnet"
+	"github.com/coder/coder/tailnet"
 )
 
 type Options struct {
@@ -52,14 +60,16 @@ type Options struct {
 	// options.AppHostname is set.
 	AppHostnameRegex *regexp.Regexp
 
-	RealIPConfig *httpmw.RealIPConfig
-
+	RealIPConfig       *httpmw.RealIPConfig
 	Tracing            trace.TracerProvider
 	PrometheusRegistry *prometheus.Registry
+	TLSCertificates    []tls.Certificate
 
-	APIRateLimit     int
-	SecureAuthCookie bool
-	DisablePathApps  bool
+	APIRateLimit           int
+	SecureAuthCookie       bool
+	DisablePathApps        bool
+	DERPEnabled            bool
+	DERPServerRelayAddress string
 
 	ProxySessionToken string
 	// AllowAllCors will set all CORs headers to '*'.
@@ -103,12 +113,14 @@ type Server struct {
 	// the moon's token.
 	SDKClient *wsproxysdk.Client
 
-	// TODO: Missing:
-	//		- derpserver
+	// DERP
+	derpMesh *derpmesh.Mesh
 
 	// Used for graceful shutdown. Required for the dialer.
-	ctx    context.Context
-	cancel context.CancelFunc
+	ctx           context.Context
+	cancel        context.CancelFunc
+	derpCloseFunc func()
+	registerDone  <-chan struct{}
 }
 
 // New creates a new workspace proxy server. This requires a primary coderd
@@ -143,21 +155,33 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		return nil, xerrors.Errorf("%q is a workspace proxy, not a primary coderd instance", opts.DashboardURL)
 	}
 
-	regResp, err := client.RegisterWorkspaceProxy(ctx, wsproxysdk.RegisterWorkspaceProxyRequest{
-		AccessURL:        opts.AccessURL.String(),
-		WildcardHostname: opts.AppHostname,
-	})
-	if err != nil {
-		return nil, xerrors.Errorf("register proxy: %w", err)
+	meshRootCA := x509.NewCertPool()
+	for _, certificate := range opts.TLSCertificates {
+		for _, certificatePart := range certificate.Certificate {
+			certificate, err := x509.ParseCertificate(certificatePart)
+			if err != nil {
+				return nil, xerrors.Errorf("parse certificate %s: %w", certificate.Subject.CommonName, err)
+			}
+			meshRootCA.AddCert(certificate)
+		}
+	}
+	// This TLS configuration spoofs access from the access URL hostname
+	// assuming that the certificates provided will cover that hostname.
+	//
+	// Replica sync and DERP meshing require accessing replicas via their
+	// internal IP addresses, and if TLS is configured we use the same
+	// certificates.
+	meshTLSConfig := &tls.Config{
+		MinVersion:   tls.VersionTLS12,
+		Certificates: opts.TLSCertificates,
+		RootCAs:      meshRootCA,
+		ServerName:   opts.AccessURL.Hostname(),
 	}
 
-	secKey, err := workspaceapps.KeyFromString(regResp.AppSecurityKey)
-	if err != nil {
-		return nil, xerrors.Errorf("parse app security key: %w", err)
-	}
+	derpServer := derp.NewServer(key.NewNode(), tailnet.Logger(opts.Logger.Named("derp")))
 
-	r := chi.NewRouter()
 	ctx, cancel := context.WithCancel(context.Background())
+	r := chi.NewRouter()
 	s := &Server{
 		Options:            opts,
 		Handler:            r,
@@ -166,11 +190,50 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		TracerProvider:     opts.Tracing,
 		PrometheusRegistry: opts.PrometheusRegistry,
 		SDKClient:          client,
+		derpMesh:           derpmesh.New(opts.Logger.Named("derpmesh"), derpServer, meshTLSConfig),
 		ctx:                ctx,
 		cancel:             cancel,
 	}
 
-	connInfo, err := client.SDKClient.WorkspaceAgentConnectionInfo(ctx)
+	// Register the workspace proxy with the primary coderd instance and start a
+	// goroutine to periodically re-register.
+	replicaID := uuid.New()
+	osHostname, err := os.Hostname()
+	if err != nil {
+		return nil, xerrors.Errorf("get OS hostname: %w", err)
+	}
+	regResp, registerDone, err := client.RegisterWorkspaceProxyLoop(ctx, wsproxysdk.RegisterWorkspaceProxyLoopOpts{
+		Logger: opts.Logger,
+		Request: wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           opts.AccessURL.String(),
+			WildcardHostname:    opts.AppHostname,
+			DerpEnabled:         opts.DERPEnabled,
+			ReplicaID:           replicaID,
+			ReplicaHostname:     osHostname,
+			ReplicaError:        "",
+			ReplicaRelayAddress: opts.DERPServerRelayAddress,
+			Version:             buildinfo.Version(),
+		},
+		MutateFn:   s.mutateRegister,
+		CallbackFn: s.handleRegister,
+		FailureFn:  s.handleRegisterFailure,
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("register proxy: %w", err)
+	}
+	s.registerDone = registerDone
+	err = s.handleRegister(ctx, regResp)
+	if err != nil {
+		return nil, xerrors.Errorf("handle register: %w", err)
+	}
+	derpServer.SetMeshKey(regResp.DERPMeshKey)
+
+	secKey, err := workspaceapps.KeyFromString(regResp.AppSecurityKey)
+	if err != nil {
+		return nil, xerrors.Errorf("parse app security key: %w", err)
+	}
+
+	connInfo, err := client.SDKClient.WorkspaceAgentConnectionInfoGeneric(ctx)
 	if err != nil {
 		return nil, xerrors.Errorf("get derpmap: %w", err)
 	}
@@ -216,6 +279,9 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		SecureAuthCookie: opts.SecureAuthCookie,
 	}
 
+	derpHandler := derphttp.Handler(derpServer)
+	derpHandler, s.derpCloseFunc = tailnet.WithWebsocketSupport(derpServer, derpHandler)
+
 	// The primary coderd dashboard needs to make some GET requests to
 	// the workspace proxies to check latency.
 	corsMW := httpmw.Cors(opts.AllowAllCors, opts.DashboardURL.String())
@@ -266,6 +332,14 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		s.AppServer.Attach(r)
 	})
 
+	r.Route("/derp", func(r chi.Router) {
+		r.Get("/", derpHandler.ServeHTTP)
+		// This is used when UDP is blocked, and latency must be checked via HTTP(s).
+		r.Get("/latency-check", func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		})
+	})
+
 	r.Get("/api/v2/buildinfo", s.buildInfo)
 	r.Get("/healthz", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("OK")) })
 	// TODO: @emyrk should this be authenticated or debounced?
@@ -295,20 +369,56 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 func (s *Server) Close() error {
 	s.cancel()
 
-	// A timeout to prevent the SDK from blocking the server shutdown.
-	tmp, cancel := context.WithTimeout(context.Background(), 3*time.Second)
-	defer cancel()
-	_ = s.SDKClient.WorkspaceProxyGoingAway(tmp)
-	_ = s.AppServer.AgentProvider.Close()
-
-	return s.AppServer.Close()
+	var err error
+	registerDoneWaitTicker := time.NewTicker(11 * time.Second) // the attempt timeout is 10s
+	select {
+	case <-registerDoneWaitTicker.C:
+		err = multierror.Append(err, xerrors.New("timed out waiting for registerDone"))
+	case <-s.registerDone:
+	}
+	s.derpCloseFunc()
+	appServerErr := s.AppServer.Close()
+	if appServerErr != nil {
+		err = multierror.Append(err, appServerErr)
+	}
+	agentProviderErr := s.AppServer.AgentProvider.Close()
+	if agentProviderErr != nil {
+		err = multierror.Append(err, agentProviderErr)
+	}
+	s.SDKClient.SDKClient.HTTPClient.CloseIdleConnections()
+	return err
 }
 
 func (s *Server) DialWorkspaceAgent(id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
 	return s.SDKClient.DialWorkspaceAgent(s.ctx, id, nil)
 }
 
-func (s *Server) DialCoordinator(ctx context.Context) (agpl.MultiAgentConn, error) {
+func (*Server) mutateRegister(_ *wsproxysdk.RegisterWorkspaceProxyRequest) {
+	// TODO: we should probably ping replicas similarly to the replicasync
+	// package in the primary and update req.ReplicaError accordingly.
+}
+
+func (s *Server) handleRegister(_ context.Context, res wsproxysdk.RegisterWorkspaceProxyResponse) error {
+	addresses := make([]string, len(res.SiblingReplicas))
+	for i, replica := range res.SiblingReplicas {
+		addresses[i] = replica.RelayAddress
+	}
+	s.derpMesh.SetAddresses(addresses, false)
+
+	return nil
+}
+
+func (s *Server) handleRegisterFailure(err error) {
+	if s.ctx.Err() != nil {
+		return
+	}
+	s.Logger.Fatal(s.ctx,
+		"failed to periodically re-register workspace proxy with primary Coder deployment",
+		slog.Error(err),
+	)
+}
+
+func (s *Server) DialCoordinator(ctx context.Context) (tailnet.MultiAgentConn, error) {
 	return s.SDKClient.DialCoordinator(ctx)
 }
 

enterprise/wsproxy/wsproxy_test.go

@@ -1,19 +1,376 @@
 package wsproxy_test
 
 import (
+	"fmt"
 	"net"
 	"testing"
 
+	"github.com/davecgh/go-spew/spew"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"tailscale.com/tailcfg"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/agent"
 	"github.com/coder/coder/cli/clibase"
 	"github.com/coder/coder/coderd"
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/coderd/healthcheck"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/workspaceapps/apptest"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/codersdk/agentsdk"
 	"github.com/coder/coder/enterprise/coderd/coderdenttest"
 	"github.com/coder/coder/enterprise/coderd/license"
+	"github.com/coder/coder/provisioner/echo"
+	"github.com/coder/coder/testutil"
 )
 
+func TestDERP(t *testing.T) {
+	t.Parallel()
+
+	deploymentValues := coderdtest.DeploymentValues(t)
+	deploymentValues.Experiments = []string{
+		string(codersdk.ExperimentMoons),
+		"*",
+	}
+
+	client, closer, api, user := 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 two running external proxies.
+	proxyAPI1 := coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
+		Name: "best-proxy",
+	})
+	proxyAPI2 := coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
+		Name: "worst-proxy",
+	})
+
+	// Create a proxy that is never started.
+	createProxyCtx := testutil.Context(t, testutil.WaitLong)
+	_, err := client.CreateWorkspaceProxy(createProxyCtx, codersdk.CreateWorkspaceProxyRequest{
+		Name: "never-started-proxy",
+	})
+	require.NoError(t, err)
+
+	// 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, 4) {
+			return false
+		}
+
+		// The first 3 regions should be healthy.
+		for _, r := range regions[:3] {
+			if !r.Healthy {
+				return false
+			}
+		}
+
+		// The last region should never be healthy.
+		assert.False(t, regions[3].Healthy)
+		return true
+	}, testutil.WaitLong, testutil.IntervalMedium)
+
+	// Create a workspace + apps
+	authToken := uuid.NewString()
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse:          echo.ParseComplete,
+		ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
+	})
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+	build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+	workspace.LatestBuild = build
+
+	agentID := uuid.Nil
+resourceLoop:
+	for _, res := range build.Resources {
+		for _, agnt := range res.Agents {
+			agentID = agnt.ID
+			break resourceLoop
+		}
+	}
+	require.NotEqual(t, uuid.Nil, agentID)
+
+	// Connect an agent to the workspace
+	agentClient := agentsdk.New(client.URL)
+	agentClient.SetSessionToken(authToken)
+	agentCloser := agent.New(agent.Options{
+		Client: agentClient,
+		Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
+	})
+	t.Cleanup(func() {
+		_ = agentCloser.Close()
+	})
+	coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+	t.Run("ReturnedInDERPMap", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		connInfo, err := client.WorkspaceAgentConnectionInfo(ctx, agentID)
+		require.NoError(t, err)
+
+		// There should be three DERP servers in the map: the primary, and each
+		// of the two running proxies.
+		require.NotNil(t, connInfo.DERPMap)
+		require.Len(t, connInfo.DERPMap.Regions, 3)
+
+		var (
+			primaryRegion *tailcfg.DERPRegion
+			proxy1Region  *tailcfg.DERPRegion
+			proxy2Region  *tailcfg.DERPRegion
+		)
+		for _, r := range connInfo.DERPMap.Regions {
+			if r.EmbeddedRelay {
+				primaryRegion = r
+				continue
+			}
+			if r.RegionName == "best-proxy" {
+				proxy1Region = r
+				continue
+			}
+			if r.RegionName == "worst-proxy" {
+				proxy2Region = 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)
+
+		// The first proxy region:
+		require.Equal(t, "best-proxy", proxy1Region.RegionName)
+		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)
+
+		// 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)
+	})
+
+	t.Run("ConnectDERP", func(t *testing.T) {
+		t.Parallel()
+
+		connInfo, err := client.WorkspaceAgentConnectionInfo(testutil.Context(t, testutil.WaitLong), agentID)
+		require.NoError(t, err)
+		require.NotNil(t, connInfo.DERPMap)
+		require.Len(t, connInfo.DERPMap.Regions, 3)
+
+		// Connect to each region.
+		for _, r := range connInfo.DERPMap.Regions {
+			r := r
+
+			t.Run(r.RegionName, func(t *testing.T) {
+				t.Parallel()
+
+				derpMap := &tailcfg.DERPMap{
+					Regions: map[int]*tailcfg.DERPRegion{
+						r.RegionID: r,
+					},
+					OmitDefaultRegions: true,
+				}
+
+				ctx := testutil.Context(t, testutil.WaitLong)
+				report := healthcheck.DERPReport{}
+				report.Run(ctx, &healthcheck.DERPReportOptions{
+					DERPMap: derpMap,
+				})
+
+				t.Log("healthcheck report: " + spew.Sdump(&report))
+				require.True(t, report.Healthy, "healthcheck failed, see report dump")
+			})
+		}
+	})
+}
+
+func TestDERPEndToEnd(t *testing.T) {
+	t.Parallel()
+
+	deploymentValues := coderdtest.DeploymentValues(t)
+	deploymentValues.Experiments = []string{
+		string(codersdk.ExperimentMoons),
+		"*",
+	}
+
+	client, closer, api, user := 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()
+	})
+
+	coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
+		Name: "best-proxy",
+	})
+
+	// Wait for the proxy 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
+		}
+		for _, r := range regions {
+			if !r.Healthy {
+				return false
+			}
+		}
+		return true
+	}, testutil.WaitLong, testutil.IntervalMedium)
+
+	// Swap out the DERPMapper for a fake one that only returns the proxy. This
+	// allows us to force the agent to pick the proxy as its preferred region.
+	oldDERPMapper := *api.AGPL.DERPMapper.Load()
+	newDERPMapper := func(derpMap *tailcfg.DERPMap) *tailcfg.DERPMap {
+		derpMap = oldDERPMapper(derpMap)
+		// Strip everything but the proxy, which is region ID 10001.
+		derpMap.Regions = map[int]*tailcfg.DERPRegion{
+			10001: derpMap.Regions[10001],
+		}
+		derpMap.OmitDefaultRegions = true
+		return derpMap
+	}
+	api.AGPL.DERPMapper.Store(&newDERPMapper)
+
+	// Create a workspace + apps
+	authToken := uuid.NewString()
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse:          echo.ParseComplete,
+		ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
+	})
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+	build := coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+	workspace.LatestBuild = build
+
+	agentID := uuid.Nil
+resourceLoop:
+	for _, res := range build.Resources {
+		for _, agnt := range res.Agents {
+			agentID = agnt.ID
+			break resourceLoop
+		}
+	}
+	require.NotEqual(t, uuid.Nil, agentID)
+
+	// Connect an agent to the workspace
+	agentClient := agentsdk.New(client.URL)
+	agentClient.SetSessionToken(authToken)
+	agentCloser := agent.New(agent.Options{
+		Client: agentClient,
+		Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
+	})
+	t.Cleanup(func() {
+		_ = agentCloser.Close()
+	})
+	coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+	// Connect to the workspace agent.
+	ctx := testutil.Context(t, testutil.WaitLong)
+	conn, err := client.DialWorkspaceAgent(ctx, agentID, &codersdk.DialWorkspaceAgentOptions{
+		Logger: slogtest.Make(t, &slogtest.Options{
+			IgnoreErrors: true,
+		}).Named("client").Leveled(slog.LevelDebug),
+		// Force DERP.
+		BlockEndpoints: true,
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		err := conn.Close()
+		assert.NoError(t, err)
+	})
+
+	ok := conn.AwaitReachable(ctx)
+	require.True(t, ok)
+
+	_, p2p, _, err := conn.Ping(ctx)
+	require.NoError(t, err)
+	require.False(t, p2p)
+}
+
 func TestWorkspaceProxyWorkspaceApps_Wsconncache(t *testing.T) {
 	t.Parallel()
 
@@ -27,7 +384,7 @@ func TestWorkspaceProxyWorkspaceApps_Wsconncache(t *testing.T) {
 			"*",
 		}
 
-		client, _, api, user := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+		client, closer, api, user := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
 			Options: &coderdtest.Options{
 				DeploymentValues:         deploymentValues,
 				AppHostname:              "*.primary.test.coder.com",
@@ -48,6 +405,9 @@ func TestWorkspaceProxyWorkspaceApps_Wsconncache(t *testing.T) {
 				},
 			},
 		})
+		t.Cleanup(func() {
+			_ = closer.Close()
+		})
 
 		// Create the external proxy
 		if opts.DisableSubdomainApps {

enterprise/wsproxy/wsproxysdk/wsproxysdk.go

@@ -21,7 +21,7 @@ import (
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/workspaceapps"
 	"github.com/coder/coder/codersdk"
-	agpl "github.com/coder/coder/tailnet"
+	"github.com/coder/coder/tailnet"
 )
 
 // Client is a HTTP client for a subset of Coder API routes that external
@@ -157,10 +157,39 @@ type RegisterWorkspaceProxyRequest struct {
 	AccessURL string `json:"access_url"`
 	// WildcardHostname that the workspace proxy api is serving for subdomain apps.
 	WildcardHostname string `json:"wildcard_hostname"`
+	// DerpEnabled indicates whether the proxy should be included in the DERP
+	// map or not.
+	DerpEnabled bool `json:"derp_enabled"`
+
+	// ReplicaID is a unique identifier for the replica of the proxy that is
+	// registering. It should be generated by the client on startup and
+	// persisted (in memory only) until the process is restarted.
+	ReplicaID uuid.UUID `json:"replica_id"`
+	// ReplicaHostname is the OS hostname of the machine that the proxy is running
+	// on.  This is only used for tracking purposes in the replicas table.
+	ReplicaHostname string `json:"hostname"`
+	// ReplicaError is the error that the replica encountered when trying to
+	// dial it's peers. This is stored in the replicas table for debugging
+	// purposes but does not affect the proxy's ability to register.
+	//
+	// This value is only stored on subsequent requests to the register
+	// endpoint, not the first request.
+	ReplicaError string `json:"replica_error"`
+	// ReplicaRelayAddress is the DERP address of the replica that other
+	// replicas may use to connect internally for DERP meshing.
+	ReplicaRelayAddress string `json:"replica_relay_address"`
+
+	// Version is the Coder version of the proxy.
+	Version string `json:"version"`
 }
 
 type RegisterWorkspaceProxyResponse struct {
 	AppSecurityKey string `json:"app_security_key"`
+	DERPMeshKey    string `json:"derp_mesh_key"`
+	DERPRegionID   int32  `json:"derp_region_id"`
+	// SiblingReplicas is a list of all other replicas of the proxy that have
+	// not timed out.
+	SiblingReplicas []codersdk.Replica `json:"sibling_replicas"`
 }
 
 func (c *Client) RegisterWorkspaceProxy(ctx context.Context, req RegisterWorkspaceProxyRequest) (RegisterWorkspaceProxyResponse, error) {
@@ -180,22 +209,183 @@ func (c *Client) RegisterWorkspaceProxy(ctx context.Context, req RegisterWorkspa
 	return resp, json.NewDecoder(res.Body).Decode(&resp)
 }
 
-func (c *Client) WorkspaceProxyGoingAway(ctx context.Context) error {
+type DeregisterWorkspaceProxyRequest struct {
+	// ReplicaID is a unique identifier for the replica of the proxy that is
+	// deregistering. It should be generated by the client on startup and
+	// should've already been passed to the register endpoint.
+	ReplicaID uuid.UUID `json:"replica_id"`
+}
+
+func (c *Client) DeregisterWorkspaceProxy(ctx context.Context, req DeregisterWorkspaceProxyRequest) error {
 	res, err := c.Request(ctx, http.MethodPost,
-		"/api/v2/workspaceproxies/me/goingaway",
-		nil,
+		"/api/v2/workspaceproxies/me/deregister",
+		req,
 	)
 	if err != nil {
 		return xerrors.Errorf("make request: %w", err)
 	}
 	defer res.Body.Close()
 
-	if res.StatusCode != http.StatusOK {
+	if res.StatusCode != http.StatusNoContent {
 		return codersdk.ReadBodyAsError(res)
 	}
 	return nil
 }
 
+type RegisterWorkspaceProxyLoopOpts struct {
+	Logger  slog.Logger
+	Request RegisterWorkspaceProxyRequest
+
+	// Interval between registration attempts. Defaults to 30 seconds. Note that
+	// the initial registration is not delayed by this interval.
+	Interval time.Duration
+	// MaxFailureCount is the maximum amount of attempts that the loop will
+	// retry registration before giving up. Defaults to 10 (for ~5 minutes).
+	MaxFailureCount int
+	// AttemptTimeout is the maximum amount of time that the loop will wait for
+	// a response from the server before considering the attempt a failure.
+	// Defaults to 10 seconds.
+	AttemptTimeout time.Duration
+
+	// MutateFn is called before each request to mutate the request struct. This
+	// can be used to update fields like ReplicaError.
+	MutateFn func(req *RegisterWorkspaceProxyRequest)
+	// CallbackFn is called with the response from the server after each
+	// successful registration, except the first. The callback function is
+	// called in a blocking manner, so it should avoid blocking for too long. If
+	// the callback returns an error, the loop will stop immediately and the
+	// error will be returned to the FailureFn.
+	CallbackFn func(ctx context.Context, res RegisterWorkspaceProxyResponse) error
+	// FailureFn is called with the last error returned from the server if the
+	// context is canceled, registration fails for more than MaxFailureCount,
+	// or if any permanent values in the response change.
+	FailureFn func(err error)
+}
+
+// RegisterWorkspaceProxyLoop will register the workspace proxy and then start a
+// goroutine to keep registering periodically in the background.
+//
+// The first response is returned immediately, and subsequent responses will be
+// notified to the given CallbackFn. When the context is canceled the loop will
+// stop immediately and the context error will be returned to the FailureFn.
+//
+// The returned channel will be closed when the loop stops and can be used to
+// ensure the loop is dead before continuing. When a fatal error is encountered,
+// the proxy will be deregistered (with the same ReplicaID and AttemptTimeout)
+// before calling the FailureFn.
+func (c *Client) RegisterWorkspaceProxyLoop(ctx context.Context, opts RegisterWorkspaceProxyLoopOpts) (RegisterWorkspaceProxyResponse, <-chan struct{}, error) {
+	if opts.Interval == 0 {
+		opts.Interval = 30 * time.Second
+	}
+	if opts.MaxFailureCount == 0 {
+		opts.MaxFailureCount = 10
+	}
+	if opts.AttemptTimeout == 0 {
+		opts.AttemptTimeout = 10 * time.Second
+	}
+	if opts.MutateFn == nil {
+		opts.MutateFn = func(_ *RegisterWorkspaceProxyRequest) {}
+	}
+	if opts.CallbackFn == nil {
+		opts.CallbackFn = func(_ context.Context, _ RegisterWorkspaceProxyResponse) error {
+			return nil
+		}
+	}
+
+	failureFn := func(err error) {
+		// We have to use background context here because the original context
+		// may be canceled.
+		deregisterCtx, cancel := context.WithTimeout(context.Background(), opts.AttemptTimeout)
+		defer cancel()
+		deregisterErr := c.DeregisterWorkspaceProxy(deregisterCtx, DeregisterWorkspaceProxyRequest{
+			ReplicaID: opts.Request.ReplicaID,
+		})
+		if deregisterErr != nil {
+			opts.Logger.Error(ctx,
+				"failed to deregister workspace proxy with Coder primary (it will be automatically deregistered shortly)",
+				slog.Error(deregisterErr),
+			)
+		}
+
+		if opts.FailureFn != nil {
+			opts.FailureFn(err)
+		}
+	}
+
+	originalRes, err := c.RegisterWorkspaceProxy(ctx, opts.Request)
+	if err != nil {
+		return RegisterWorkspaceProxyResponse{}, nil, xerrors.Errorf("register workspace proxy: %w", err)
+	}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+
+		var (
+			failedAttempts = 0
+			ticker         = time.NewTicker(opts.Interval)
+		)
+		for {
+			select {
+			case <-ctx.Done():
+				failureFn(ctx.Err())
+				return
+			case <-ticker.C:
+			}
+
+			opts.Logger.Debug(ctx,
+				"re-registering workspace proxy with Coder primary",
+				slog.F("req", opts.Request),
+				slog.F("timeout", opts.AttemptTimeout),
+				slog.F("failed_attempts", failedAttempts),
+			)
+			opts.MutateFn(&opts.Request)
+			registerCtx, cancel := context.WithTimeout(ctx, opts.AttemptTimeout)
+			res, err := c.RegisterWorkspaceProxy(registerCtx, opts.Request)
+			cancel()
+			if err != nil {
+				failedAttempts++
+				opts.Logger.Warn(ctx,
+					"failed to re-register workspace proxy with Coder primary",
+					slog.F("req", opts.Request),
+					slog.F("timeout", opts.AttemptTimeout),
+					slog.F("failed_attempts", failedAttempts),
+					slog.Error(err),
+				)
+
+				if failedAttempts > opts.MaxFailureCount {
+					failureFn(xerrors.Errorf("exceeded re-registration failure count of %d: last error: %w", opts.MaxFailureCount, err))
+					return
+				}
+				continue
+			}
+			failedAttempts = 0
+
+			if res.AppSecurityKey != originalRes.AppSecurityKey {
+				failureFn(xerrors.New("app security key has changed, proxy must be restarted"))
+				return
+			}
+			if res.DERPMeshKey != originalRes.DERPMeshKey {
+				failureFn(xerrors.New("DERP mesh key has changed, proxy must be restarted"))
+				return
+			}
+			if res.DERPRegionID != originalRes.DERPRegionID {
+				failureFn(xerrors.New("DERP region ID has changed, proxy must be restarted"))
+			}
+
+			err = opts.CallbackFn(ctx, res)
+			if err != nil {
+				failureFn(xerrors.Errorf("callback fn returned error: %w", err))
+				return
+			}
+
+			ticker.Reset(opts.Interval)
+		}
+	}()
+
+	return originalRes, done, nil
+}
+
 type CoordinateMessageType int
 
 const (
@@ -207,14 +397,14 @@ const (
 type CoordinateMessage struct {
 	Type    CoordinateMessageType `json:"type"`
 	AgentID uuid.UUID             `json:"agent_id"`
-	Node    *agpl.Node            `json:"node"`
+	Node    *tailnet.Node         `json:"node"`
 }
 
 type CoordinateNodes struct {
-	Nodes []*agpl.Node
+	Nodes []*tailnet.Node
 }
 
-func (c *Client) DialCoordinator(ctx context.Context) (agpl.MultiAgentConn, error) {
+func (c *Client) DialCoordinator(ctx context.Context) (tailnet.MultiAgentConn, error) {
 	ctx, cancel := context.WithCancel(ctx)
 
 	coordinateURL, err := c.SDKClient.URL.Parse("/api/v2/workspaceproxies/me/coordinate")
@@ -248,7 +438,7 @@ func (c *Client) DialCoordinator(ctx context.Context) (agpl.MultiAgentConn, erro
 		legacyAgentCache: map[uuid.UUID]bool{},
 	}
 
-	ma := (&agpl.MultiAgent{
+	ma := (&tailnet.MultiAgent{
 		ID:                uuid.New(),
 		AgentIsLegacyFunc: rma.AgentIsLegacy,
 		OnSubscribe:       rma.OnSubscribe,
@@ -300,7 +490,7 @@ func (a *remoteMultiAgentHandler) writeJSON(v interface{}) error {
 
 	// Set a deadline so that hung connections don't put back pressure on the system.
 	// Node updates are tiny, so even the dinkiest connection can handle them if it's not hung.
-	err = a.nc.SetWriteDeadline(time.Now().Add(agpl.WriteTimeout))
+	err = a.nc.SetWriteDeadline(time.Now().Add(tailnet.WriteTimeout))
 	if err != nil {
 		return xerrors.Errorf("set write deadline: %w", err)
 	}
@@ -322,21 +512,21 @@ func (a *remoteMultiAgentHandler) writeJSON(v interface{}) error {
 	return nil
 }
 
-func (a *remoteMultiAgentHandler) OnNodeUpdate(_ uuid.UUID, node *agpl.Node) error {
+func (a *remoteMultiAgentHandler) OnNodeUpdate(_ uuid.UUID, node *tailnet.Node) error {
 	return a.writeJSON(CoordinateMessage{
 		Type: CoordinateMessageTypeNodeUpdate,
 		Node: node,
 	})
 }
 
-func (a *remoteMultiAgentHandler) OnSubscribe(_ agpl.Queue, agentID uuid.UUID) (*agpl.Node, error) {
+func (a *remoteMultiAgentHandler) OnSubscribe(_ tailnet.Queue, agentID uuid.UUID) (*tailnet.Node, error) {
 	return nil, a.writeJSON(CoordinateMessage{
 		Type:    CoordinateMessageTypeSubscribe,
 		AgentID: agentID,
 	})
 }
 
-func (a *remoteMultiAgentHandler) OnUnsubscribe(_ agpl.Queue, agentID uuid.UUID) error {
+func (a *remoteMultiAgentHandler) OnUnsubscribe(_ tailnet.Queue, agentID uuid.UUID) error {
 	return a.writeJSON(CoordinateMessage{
 		Type:    CoordinateMessageTypeUnsubscribe,
 		AgentID: agentID,