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

enterprise/audit/table.go

@@ -197,6 +197,8 @@ var auditableResourcesTypes = map[any]map[string]Action{
 		"updated_at":          ActionIgnore,
 		"deleted":             ActionIgnore,
 		"token_hashed_secret": ActionSecret,
+		"derp_enabled":        ActionTrack,
+		"region_id":           ActionTrack,
 	},
 }
 

enterprise/cli/proxyserver.go

@@ -220,21 +220,23 @@ func (*RootCmd) proxyServer() *clibase.Cmd {
 			}
 
 			proxy, err := wsproxy.New(ctx, &wsproxy.Options{
-				Logger:             logger,
-				Experiments:        coderd.ReadExperiments(logger, cfg.Experiments.Value()),
-				HTTPClient:         httpClient,
-				DashboardURL:       primaryAccessURL.Value(),
-				AccessURL:          cfg.AccessURL.Value(),
-				AppHostname:        appHostname,
-				AppHostnameRegex:   appHostnameRegex,
-				RealIPConfig:       realIPConfig,
-				Tracing:            tracer,
-				PrometheusRegistry: prometheusRegistry,
-				APIRateLimit:       int(cfg.RateLimit.API.Value()),
-				SecureAuthCookie:   cfg.SecureAuthCookie.Value(),
-				DisablePathApps:    cfg.DisablePathApps.Value(),
-				ProxySessionToken:  proxySessionToken.Value(),
-				AllowAllCors:       cfg.Dangerous.AllowAllCors.Value(),
+				Logger:                 logger,
+				Experiments:            coderd.ReadExperiments(logger, cfg.Experiments.Value()),
+				HTTPClient:             httpClient,
+				DashboardURL:           primaryAccessURL.Value(),
+				AccessURL:              cfg.AccessURL.Value(),
+				AppHostname:            appHostname,
+				AppHostnameRegex:       appHostnameRegex,
+				RealIPConfig:           realIPConfig,
+				Tracing:                tracer,
+				PrometheusRegistry:     prometheusRegistry,
+				APIRateLimit:           int(cfg.RateLimit.API.Value()),
+				SecureAuthCookie:       cfg.SecureAuthCookie.Value(),
+				DisablePathApps:        cfg.DisablePathApps.Value(),
+				ProxySessionToken:      proxySessionToken.Value(),
+				AllowAllCors:           cfg.Dangerous.AllowAllCors.Value(),
+				DERPEnabled:            cfg.DERP.Server.Enable.Value(),
+				DERPServerRelayAddress: cfg.DERP.Server.RelayURL.String(),
 			})
 			if err != nil {
 				return xerrors.Errorf("create workspace proxy: %w", err)

enterprise/coderd/coderd.go

@@ -5,11 +5,17 @@ import (
 	"crypto/ed25519"
 	"crypto/tls"
 	"crypto/x509"
+	"fmt"
+	"math"
 	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
 	"sync"
 	"time"
 
 	"golang.org/x/xerrors"
+	"tailscale.com/tailcfg"
 
 	"github.com/cenkalti/backoff/v4"
 	"github.com/go-chi/chi/v5"
@@ -158,7 +164,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 				r.Get("/coordinate", api.workspaceProxyCoordinate)
 				r.Post("/issue-signed-app-token", api.workspaceProxyIssueSignedAppToken)
 				r.Post("/register", api.workspaceProxyRegister)
-				r.Post("/goingaway", api.workspaceProxyGoingAway)
+				r.Post("/deregister", api.workspaceProxyDeregister)
 			})
 			r.Route("/{workspaceproxy}", func(r chi.Router) {
 				r.Use(
@@ -294,10 +300,11 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 		ServerName:   options.AccessURL.Hostname(),
 	}
 	api.replicaManager, err = replicasync.New(ctx, options.Logger, options.Database, options.Pubsub, &replicasync.Options{
-		ID:           api.AGPL.ID,
-		RelayAddress: options.DERPServerRelayAddress,
-		RegionID:     int32(options.DERPServerRegionID),
-		TLSConfig:    meshTLSConfig,
+		ID:             api.AGPL.ID,
+		RelayAddress:   options.DERPServerRelayAddress,
+		RegionID:       int32(options.DERPServerRegionID),
+		TLSConfig:      meshTLSConfig,
+		UpdateInterval: options.ReplicaSyncUpdateInterval,
 	})
 	if err != nil {
 		return nil, xerrors.Errorf("initialize replica: %w", err)
@@ -345,8 +352,9 @@ type Options struct {
 	SCIMAPIKey  []byte
 
 	// Used for high availability.
-	DERPServerRelayAddress string
-	DERPServerRegionID     int
+	ReplicaSyncUpdateInterval time.Duration
+	DERPServerRelayAddress    string
+	DERPServerRegionID        int
 
 	// Used for user quiet hours schedules.
 	DefaultQuietHoursSchedule string // cron schedule, if empty user quiet hours schedules are disabled
@@ -394,7 +402,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 
 	entitlements, err := license.Entitlements(
 		ctx, api.Database,
-		api.Logger, len(api.replicaManager.All()), len(api.GitAuthConfigs), api.Keys, map[codersdk.FeatureName]bool{
+		api.Logger, len(api.replicaManager.AllPrimary()), len(api.GitAuthConfigs), api.Keys, map[codersdk.FeatureName]bool{
 			codersdk.FeatureAuditLog:                   api.AuditLogging,
 			codersdk.FeatureBrowserOnly:                api.BrowserOnly,
 			codersdk.FeatureSCIM:                       len(api.SCIMAPIKey) != 0,
@@ -568,6 +576,15 @@ 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)
+			api.AGPL.DERPMapper.Store(&fn)
+		} else {
+			api.AGPL.DERPMapper.Store(nil)
+		}
+	}
+
 	api.entitlementsMu.Lock()
 	defer api.entitlementsMu.Unlock()
 	api.entitlements = entitlements
@@ -576,6 +593,167 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 	return nil
 }
 
+// getProxyDERPStartingRegionID returns the starting region ID that should be
+// used for workspace proxies. A proxy's actual region ID is the return value
+// from this function + it's RegionID field.
+//
+// Two ints are returned, the first is the starting region ID for proxies, and
+// the second is the maximum region ID that already exists in the DERP map.
+func getProxyDERPStartingRegionID(derpMap *tailcfg.DERPMap) (sID int64, mID int64) {
+	var maxRegionID int64
+	for _, region := range derpMap.Regions {
+		rid := int64(region.RegionID)
+		if rid > maxRegionID {
+			maxRegionID = rid
+		}
+	}
+	if maxRegionID < 0 {
+		maxRegionID = 0
+	}
+
+	// Round to the nearest 10,000 with a sufficient buffer of at least 2,000.
+	// The buffer allows for future "fixed" regions to be added to the base DERP
+	// map without conflicting with proxy region IDs (standard DERP maps usually
+	// use incrementing IDs for new regions).
+	//
+	// Example:
+	//  maxRegionID = -2_000 -> startingRegionID = 10_000
+	//  maxRegionID = 8_000 -> startingRegionID = 10_000
+	//  maxRegionID = 8_500 -> startingRegionID = 20_000
+	//  maxRegionID = 12_000 -> startingRegionID = 20_000
+	//  maxRegionID = 20_000 -> startingRegionID = 30_000
+	const roundStartingRegionID = 10_000
+	const startingRegionIDBuffer = 2_000
+	// Add the buffer first.
+	startingRegionID := maxRegionID + startingRegionIDBuffer
+	// Round UP to the nearest 10,000. Go's math.Ceil rounds up to the nearest
+	// integer, so we need to divide by 10,000 first and then multiply by
+	// 10,000.
+	startingRegionID = int64(math.Ceil(float64(startingRegionID)/roundStartingRegionID) * roundStartingRegionID)
+	// This should never be hit but it's here just in case.
+	if startingRegionID < roundStartingRegionID {
+		startingRegionID = roundStartingRegionID
+	}
+
+	return startingRegionID, maxRegionID
+}
+
+var (
+	lastDerpConflictMutex sync.Mutex
+	lastDerpConflictLog   time.Time
+)
+
+func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*tailcfg.DERPMap) *tailcfg.DERPMap {
+	return func(derpMap *tailcfg.DERPMap) *tailcfg.DERPMap {
+		derpMap = derpMap.Clone()
+
+		// Find the starting region ID that we'll use for proxies. This must be
+		// deterministic based on the derp map.
+		startingRegionID, largestRegionID := getProxyDERPStartingRegionID(derpMap)
+		if largestRegionID >= 1<<32 {
+			// Enforce an upper bound on the region ID. This shouldn't be hit in
+			// practice, but it's a good sanity check.
+			lastDerpConflictMutex.Lock()
+			shouldLog := lastDerpConflictLog.IsZero() || time.Since(lastDerpConflictLog) > time.Minute
+			if shouldLog {
+				lastDerpConflictLog = time.Now()
+			}
+			lastDerpConflictMutex.Unlock()
+			if shouldLog {
+				logger.Warn(
+					context.Background(),
+					"existing DERP region IDs are too large, proxy region IDs will not be populated in the derp map. Please ensure that all DERP region IDs are less than 2^32",
+					slog.F("largest_region_id", largestRegionID),
+					slog.F("max_region_id", 1<<32-1),
+				)
+				return derpMap
+			}
+		}
+
+		// Add all healthy proxies to the DERP map.
+		statusMap := proxyHealth.HealthStatus()
+	statusLoop:
+		for _, status := range statusMap {
+			if status.Status != proxyhealth.Healthy || !status.Proxy.DerpEnabled {
+				// Only add healthy proxies with DERP enabled to the DERP map.
+				continue
+			}
+
+			u, err := url.Parse(status.Proxy.Url)
+			if err != nil {
+				// Not really any need to log, the proxy should be unreachable
+				// anyways and filtered out by the above condition.
+				continue
+			}
+			port := u.Port()
+			if port == "" {
+				port = "80"
+				if u.Scheme == "https" {
+					port = "443"
+				}
+			}
+			portInt, err := strconv.Atoi(port)
+			if err != nil {
+				// Not really any need to log, the proxy should be unreachable
+				// anyways and filtered out by the above condition.
+				continue
+			}
+
+			// Sanity check that the region ID and code is unique.
+			//
+			// This should be impossible to hit as the IDs are enforced to be
+			// unique by the database and the computed ID is greater than any
+			// existing ID in the DERP map.
+			regionID := int(startingRegionID) + int(status.Proxy.RegionID)
+			regionCode := fmt.Sprintf("coder_%s", strings.ToLower(status.Proxy.Name))
+			for _, r := range derpMap.Regions {
+				if r.RegionID == regionID || r.RegionCode == regionCode {
+					// 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.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",
+							slog.F("proxy_id", status.Proxy.ID),
+							slog.F("proxy_name", status.Proxy.Name),
+							slog.F("proxy_display_name", status.Proxy.DisplayName),
+							slog.F("proxy_url", status.Proxy.Url),
+							slog.F("proxy_region_id", status.Proxy.RegionID),
+							slog.F("proxy_computed_region_id", regionID),
+							slog.F("proxy_computed_region_code", regionCode),
+						)
+					}
+
+					continue statusLoop
+				}
+			}
+
+			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",
+				}},
+			}
+		}
+
+		return derpMap
+	}
+}
+
 // @Summary Get entitlements
 // @ID get-entitlements
 // @Security CoderSessionToken

enterprise/coderd/coderdenttest/coderdenttest.go

@@ -55,6 +55,7 @@ type Options struct {
 	NoDefaultQuietHoursSchedule bool
 	DontAddLicense              bool
 	DontAddFirstUser            bool
+	ReplicaSyncUpdateInterval   time.Duration
 }
 
 // New constructs a codersdk client connected to an in-memory Enterprise API instance.
@@ -86,7 +87,8 @@ func NewWithAPI(t *testing.T, options *Options) (
 		BrowserOnly:                options.BrowserOnly,
 		SCIMAPIKey:                 options.SCIMAPIKey,
 		DERPServerRelayAddress:     oop.AccessURL.String(),
-		DERPServerRegionID:         oop.DERPMap.RegionIDs()[0],
+		DERPServerRegionID:         oop.BaseDERPMap.RegionIDs()[0],
+		ReplicaSyncUpdateInterval:  options.ReplicaSyncUpdateInterval,
 		Options:                    oop,
 		EntitlementsUpdateInterval: options.EntitlementsUpdateInterval,
 		Keys:                       Keys,

enterprise/coderd/coderdenttest/proxytest.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/moby/moby/pkg/namesgenerator"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"cdr.dev/slog"
@@ -132,9 +133,15 @@ func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Clie
 		DisablePathApps:   options.DisablePathApps,
 		// We need a new registry to not conflict with the coderd internal
 		// proxy metrics.
-		PrometheusRegistry: prometheus.NewRegistry(),
+		PrometheusRegistry:     prometheus.NewRegistry(),
+		DERPEnabled:            true,
+		DERPServerRelayAddress: accessURL.String(),
 	})
 	require.NoError(t, err)
+	t.Cleanup(func() {
+		err := wssrv.Close()
+		assert.NoError(t, err)
+	})
 
 	mutex.Lock()
 	handler = wssrv.Handler

enterprise/coderd/replicas.go

@@ -24,7 +24,7 @@ func (api *API) replicas(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	replicas := api.replicaManager.All()
+	replicas := api.replicaManager.AllPrimary()
 	res := make([]codersdk.Replica, 0, len(replicas))
 	for _, replica := range replicas {
 		res = append(res, convertReplica(replica))

enterprise/coderd/workspaceproxy.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/sha256"
 	"database/sql"
+	"flag"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -14,6 +15,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/buildinfo"
 	agpl "github.com/coder/coder/coderd"
 	"github.com/coder/coder/coderd/audit"
 	"github.com/coder/coder/coderd/database"
@@ -25,6 +27,7 @@ import (
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/cryptorand"
 	"github.com/coder/coder/enterprise/coderd/proxyhealth"
+	"github.com/coder/coder/enterprise/replicasync"
 	"github.com/coder/coder/enterprise/wsproxy/wsproxysdk"
 )
 
@@ -347,10 +350,13 @@ func (api *API) postWorkspaceProxy(rw http.ResponseWriter, r *http.Request) {
 		DisplayName:       req.DisplayName,
 		Icon:              req.Icon,
 		TokenHashedSecret: hashedSecret[:],
-		CreatedAt:         database.Now(),
-		UpdatedAt:         database.Now(),
+		// Enabled by default, but will be disabled on register if the proxy has
+		// it disabled.
+		DerpEnabled: true,
+		CreatedAt:   database.Now(),
+		UpdatedAt:   database.Now(),
 	})
-	if database.IsUniqueViolation(err) {
+	if database.IsUniqueViolation(err, database.UniqueWorkspaceProxiesLowerNameIndex) {
 		httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
 			Message: fmt.Sprintf("Workspace proxy with name %q already exists.", req.Name),
 		})
@@ -489,13 +495,17 @@ func (api *API) workspaceProxyIssueSignedAppToken(rw http.ResponseWriter, r *htt
 // in the database and returns a signed token that can be used to authenticate
 // tokens.
 //
+// This is called periodically by the proxy in the background (every 30s per
+// replica) to ensure that the proxy is still registered and the corresponding
+// replica table entry is refreshed.
+//
 // @Summary Register workspace proxy
 // @ID register-workspace-proxy
 // @Security CoderSessionToken
 // @Accept json
 // @Produce json
 // @Tags Enterprise
-// @Param request body wsproxysdk.RegisterWorkspaceProxyRequest true "Issue signed app token request"
+// @Param request body wsproxysdk.RegisterWorkspaceProxyRequest true "Register workspace proxy request"
 // @Success 201 {object} wsproxysdk.RegisterWorkspaceProxyResponse
 // @Router /workspaceproxies/me/register [post]
 // @x-apidocgen {"skip": true}
@@ -523,6 +533,17 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
+	// Version check should be forced in non-dev builds and when running in
+	// tests.
+	shouldForceVersion := !buildinfo.IsDev() || flag.Lookup("test.v") != nil
+	if shouldForceVersion && req.Version != buildinfo.Version() {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Version mismatch.",
+			Detail:  fmt.Sprintf("Proxy version %q does not match primary server version %q", req.Version, buildinfo.Version()),
+		})
+		return
+	}
+
 	if err := validateProxyURL(req.AccessURL); err != nil {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
 			Message: "URL is invalid.",
@@ -541,11 +562,80 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 		}
 	}
 
-	_, err := api.Database.RegisterWorkspaceProxy(ctx, database.RegisterWorkspaceProxyParams{
-		ID:               proxy.ID,
-		Url:              req.AccessURL,
-		WildcardHostname: req.WildcardHostname,
-	})
+	if req.ReplicaID == uuid.Nil {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Replica ID is invalid.",
+		})
+		return
+	}
+
+	startingRegionID, _ := getProxyDERPStartingRegionID(api.Options.BaseDERPMap)
+	regionID := int32(startingRegionID) + proxy.RegionID
+
+	err := api.Database.InTx(func(db database.Store) error {
+		// First, update the proxy's values in the database.
+		_, err := db.RegisterWorkspaceProxy(ctx, database.RegisterWorkspaceProxyParams{
+			ID:               proxy.ID,
+			Url:              req.AccessURL,
+			DerpEnabled:      req.DerpEnabled,
+			WildcardHostname: req.WildcardHostname,
+		})
+		if err != nil {
+			return xerrors.Errorf("register workspace proxy: %w", err)
+		}
+
+		// Second, find the replica that corresponds to this proxy and refresh
+		// it if it exists. If it doesn't exist, create it.
+		now := time.Now()
+		replica, err := db.GetReplicaByID(ctx, req.ReplicaID)
+		if err == nil {
+			// Replica exists, update it.
+			if replica.StoppedAt.Valid && !replica.StartedAt.IsZero() {
+				// If the replica deregistered, it shouldn't be able to
+				// re-register before restarting.
+				// TODO: sadly this results in 500 when it should be 400
+				return xerrors.Errorf("replica %s is marked stopped", replica.ID)
+			}
+
+			replica, err = db.UpdateReplica(ctx, database.UpdateReplicaParams{
+				ID:              replica.ID,
+				UpdatedAt:       now,
+				StartedAt:       replica.StartedAt,
+				StoppedAt:       replica.StoppedAt,
+				RelayAddress:    req.ReplicaRelayAddress,
+				RegionID:        regionID,
+				Hostname:        req.ReplicaHostname,
+				Version:         req.Version,
+				Error:           req.ReplicaError,
+				DatabaseLatency: 0,
+				Primary:         false,
+			})
+			if err != nil {
+				return xerrors.Errorf("update replica: %w", err)
+			}
+		} else if xerrors.Is(err, sql.ErrNoRows) {
+			// Replica doesn't exist, create it.
+			replica, err = db.InsertReplica(ctx, database.InsertReplicaParams{
+				ID:              req.ReplicaID,
+				CreatedAt:       now,
+				StartedAt:       now,
+				UpdatedAt:       now,
+				Hostname:        req.ReplicaHostname,
+				RegionID:        regionID,
+				RelayAddress:    req.ReplicaRelayAddress,
+				Version:         req.Version,
+				DatabaseLatency: 0,
+				Primary:         false,
+			})
+			if err != nil {
+				return xerrors.Errorf("insert replica: %w", err)
+			}
+		} else if err != nil {
+			return xerrors.Errorf("get replica: %w", err)
+		}
+
+		return nil
+	}, nil)
 	if httpapi.Is404Error(err) {
 		httpapi.ResourceNotFound(rw)
 		return
@@ -555,39 +645,112 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
+	// Update replica sync and notify all other replicas to update their
+	// replica list.
+	err = api.replicaManager.PublishUpdate()
+	if err != nil {
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+	replicaUpdateCtx, replicaUpdateCancel := context.WithTimeout(ctx, 5*time.Second)
+	defer replicaUpdateCancel()
+	err = api.replicaManager.UpdateNow(replicaUpdateCtx)
+	if err != nil {
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	// Find sibling regions to respond with for derpmesh.
+	siblings := api.replicaManager.InRegion(regionID)
+	siblingsRes := make([]codersdk.Replica, 0, len(siblings))
+	for _, replica := range siblings {
+		if replica.ID == req.ReplicaID {
+			continue
+		}
+		siblingsRes = append(siblingsRes, convertReplica(replica))
+	}
+
 	// aReq.New = updatedProxy
 	httpapi.Write(ctx, rw, http.StatusCreated, wsproxysdk.RegisterWorkspaceProxyResponse{
-		AppSecurityKey: api.AppSecurityKey.String(),
+		AppSecurityKey:  api.AppSecurityKey.String(),
+		DERPMeshKey:     api.DERPServer.MeshKey(),
+		DERPRegionID:    regionID,
+		SiblingReplicas: siblingsRes,
 	})
 
 	go api.forceWorkspaceProxyHealthUpdate(api.ctx)
 }
 
-// workspaceProxyGoingAway is used to tell coderd that the workspace proxy is
-// shutting down and going away. The main purpose of this function is for the
-// health status of the workspace proxy to be more quickly updated when we know
-// that the proxy is going to be unhealthy. This does not delete the workspace
-// or cause any other side effects.
-// If the workspace proxy comes back online, even without a register, it will
-// be found healthy again by the normal checks.
-// @Summary Workspace proxy going away
-// @ID workspace-proxy-going-away
+// @Summary Deregister workspace proxy
+// @ID deregister-workspace-proxy
 // @Security CoderSessionToken
-// @Produce json
+// @Accept json
 // @Tags Enterprise
-// @Success 201 {object} codersdk.Response
-// @Router /workspaceproxies/me/goingaway [post]
+// @Param request body wsproxysdk.DeregisterWorkspaceProxyRequest true "Deregister workspace proxy request"
+// @Success 204
+// @Router /workspaceproxies/me/deregister [post]
 // @x-apidocgen {"skip": true}
-func (api *API) workspaceProxyGoingAway(rw http.ResponseWriter, r *http.Request) {
+func (api *API) workspaceProxyDeregister(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
-	// Force a health update to happen immediately. The proxy should
-	// not return a successful response if it is going away.
-	go api.forceWorkspaceProxyHealthUpdate(api.ctx)
+	var req wsproxysdk.DeregisterWorkspaceProxyRequest
+	if !httpapi.Read(ctx, rw, r, &req) {
+		return
+	}
 
-	httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
-		Message: "OK",
-	})
+	err := api.Database.InTx(func(db database.Store) error {
+		now := time.Now()
+		replica, err := db.GetReplicaByID(ctx, req.ReplicaID)
+		if err != nil {
+			return xerrors.Errorf("get replica: %w", err)
+		}
+
+		if replica.StoppedAt.Valid && !replica.StartedAt.IsZero() {
+			// TODO: sadly this results in 500 when it should be 400
+			return xerrors.Errorf("replica %s is already marked stopped", replica.ID)
+		}
+
+		replica, err = db.UpdateReplica(ctx, database.UpdateReplicaParams{
+			ID:        replica.ID,
+			UpdatedAt: now,
+			StartedAt: replica.StartedAt,
+			StoppedAt: sql.NullTime{
+				Valid: true,
+				Time:  now,
+			},
+			RelayAddress:    replica.RelayAddress,
+			RegionID:        replica.RegionID,
+			Hostname:        replica.Hostname,
+			Version:         replica.Version,
+			Error:           replica.Error,
+			DatabaseLatency: replica.DatabaseLatency,
+			Primary:         replica.Primary,
+		})
+		if err != nil {
+			return xerrors.Errorf("update replica: %w", err)
+		}
+
+		return nil
+	}, nil)
+	if httpapi.Is404Error(err) {
+		httpapi.ResourceNotFound(rw)
+		return
+	}
+	if err != nil {
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	// Publish a replicasync event with a nil ID so every replica (yes, even the
+	// current replica) will refresh its replicas list.
+	err = api.Pubsub.Publish(replicasync.PubsubEvent, []byte(uuid.Nil.String()))
+	if err != nil {
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	rw.WriteHeader(http.StatusNoContent)
+	go api.forceWorkspaceProxyHealthUpdate(api.ctx)
 }
 
 // reconnectingPTYSignedToken issues a signed app token for use when connecting
@@ -670,7 +833,8 @@ func (api *API) reconnectingPTYSignedToken(rw http.ResponseWriter, r *http.Reque
 		},
 		SessionToken: httpmw.APITokenFromRequest(r),
 		// The following fields aren't required as long as the request is authed
-		// with a valid API key.
+		// with a valid API key, which we know since this endpoint is protected
+		// by auth middleware already.
 		PathAppBaseURL: "",
 		AppHostname:    "",
 		// The following fields are empty for terminal apps.
@@ -733,10 +897,11 @@ func convertProxy(p database.WorkspaceProxy, status proxyhealth.ProxyStatus) cod
 		status.Status = proxyhealth.Unknown
 	}
 	return codersdk.WorkspaceProxy{
-		Region:    convertRegion(p, status),
-		CreatedAt: p.CreatedAt,
-		UpdatedAt: p.UpdatedAt,
-		Deleted:   p.Deleted,
+		Region:      convertRegion(p, status),
+		DerpEnabled: p.DerpEnabled,
+		CreatedAt:   p.CreatedAt,
+		UpdatedAt:   p.UpdatedAt,
+		Deleted:     p.Deleted,
 		Status: codersdk.WorkspaceProxyStatus{
 			Status:    codersdk.ProxyHealthStatus(status.Status),
 			Report:    status.Report,

enterprise/coderd/workspaceproxy_test.go

@@ -17,7 +17,9 @@ import (
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/agent"
+	"github.com/coder/coder/buildinfo"
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/dbtestutil"
 	"github.com/coder/coder/coderd/workspaceapps"
 	"github.com/coder/coder/codersdk"
@@ -167,69 +169,6 @@ func TestRegions(t *testing.T) {
 		require.Error(t, err)
 		require.Empty(t, regions)
 	})
-
-	t.Run("GoingAway", func(t *testing.T) {
-		t.Skip("This is flakey in CI because it relies on internal go routine timing. Should refactor.")
-		t.Parallel()
-
-		dv := coderdtest.DeploymentValues(t)
-		dv.Experiments = []string{
-			string(codersdk.ExperimentMoons),
-			"*",
-		}
-
-		db, pubsub := dbtestutil.NewDB(t)
-
-		ctx := testutil.Context(t, testutil.WaitLong)
-
-		client, closer, api, _ := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
-			Options: &coderdtest.Options{
-				AppHostname:      appHostname,
-				Database:         db,
-				Pubsub:           pubsub,
-				DeploymentValues: dv,
-			},
-			// The interval is set to 1 hour so the proxy health
-			// check will never happen manually. All checks will be
-			// forced updates.
-			ProxyHealthInterval: time.Hour,
-			LicenseOptions: &coderdenttest.LicenseOptions{
-				Features: license.Features{
-					codersdk.FeatureWorkspaceProxy: 1,
-				},
-			},
-		})
-		t.Cleanup(func() {
-			_ = closer.Close()
-		})
-
-		const proxyName = "testproxy"
-		proxy := coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
-			Name: proxyName,
-		})
-		_ = proxy
-
-		require.Eventuallyf(t, func() bool {
-			proxy, err := client.WorkspaceProxyByName(ctx, proxyName)
-			if err != nil {
-				// We are testing the going away, not the initial healthy.
-				// Just force an update to change this to healthy.
-				_ = api.ProxyHealth.ForceUpdate(ctx)
-				return false
-			}
-			return proxy.Status.Status == codersdk.ProxyHealthy
-		}, testutil.WaitShort, testutil.IntervalFast, "proxy never became healthy")
-
-		_ = proxy.Close()
-		// The proxy should tell the primary on close that is is no longer healthy.
-		require.Eventuallyf(t, func() bool {
-			proxy, err := client.WorkspaceProxyByName(ctx, proxyName)
-			if err != nil {
-				return false
-			}
-			return proxy.Status.Status == codersdk.ProxyUnhealthy
-		}, testutil.WaitShort, testutil.IntervalFast, "proxy never became unhealthy after close")
-	})
 }
 
 func TestWorkspaceProxyCRUD(t *testing.T) {
@@ -321,6 +260,386 @@ func TestWorkspaceProxyCRUD(t *testing.T) {
 	})
 }
 
+func TestProxyRegisterDeregister(t *testing.T) {
+	t.Parallel()
+
+	setup := func(t *testing.T) (*codersdk.Client, database.Store) {
+		dv := coderdtest.DeploymentValues(t)
+		dv.Experiments = []string{
+			string(codersdk.ExperimentMoons),
+			"*",
+		}
+
+		db, pubsub := dbtestutil.NewDB(t)
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				DeploymentValues:         dv,
+				Database:                 db,
+				Pubsub:                   pubsub,
+				IncludeProvisionerDaemon: true,
+			},
+			ReplicaSyncUpdateInterval: time.Minute,
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureWorkspaceProxy: 1,
+				},
+			},
+		})
+
+		return client, db
+	}
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		client, db := setup(t)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		const (
+			proxyName        = "hello"
+			proxyDisplayName = "Hello World"
+			proxyIcon        = "/emojis/flag.png"
+		)
+		createRes, err := client.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
+			Name:        proxyName,
+			DisplayName: proxyDisplayName,
+			Icon:        proxyIcon,
+		})
+		require.NoError(t, err)
+
+		proxyClient := wsproxysdk.New(client.URL)
+		proxyClient.SetSessionToken(createRes.ProxyToken)
+
+		// Register
+		req := wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           "https://proxy.coder.test",
+			WildcardHostname:    "*.proxy.coder.test",
+			DerpEnabled:         true,
+			ReplicaID:           uuid.New(),
+			ReplicaHostname:     "mars",
+			ReplicaError:        "",
+			ReplicaRelayAddress: "http://127.0.0.1:8080",
+			Version:             buildinfo.Version(),
+		}
+		registerRes1, err := proxyClient.RegisterWorkspaceProxy(ctx, req)
+		require.NoError(t, err)
+		require.NotEmpty(t, registerRes1.AppSecurityKey)
+		require.NotEmpty(t, registerRes1.DERPMeshKey)
+		require.EqualValues(t, 10001, registerRes1.DERPRegionID)
+		require.Empty(t, registerRes1.SiblingReplicas)
+
+		proxy, err := client.WorkspaceProxyByID(ctx, createRes.Proxy.ID)
+		require.NoError(t, err)
+		require.Equal(t, createRes.Proxy.ID, proxy.ID)
+		require.Equal(t, proxyName, proxy.Name)
+		require.Equal(t, proxyDisplayName, proxy.DisplayName)
+		require.Equal(t, proxyIcon, proxy.IconURL)
+		require.Equal(t, req.AccessURL, proxy.PathAppURL)
+		require.Equal(t, req.AccessURL, proxy.PathAppURL)
+		require.Equal(t, req.WildcardHostname, proxy.WildcardHostname)
+		require.Equal(t, req.DerpEnabled, proxy.DerpEnabled)
+		require.False(t, proxy.Deleted)
+
+		// Get the replica from the DB.
+		replica, err := db.GetReplicaByID(ctx, req.ReplicaID)
+		require.NoError(t, err)
+		require.Equal(t, req.ReplicaID, replica.ID)
+		require.Equal(t, req.ReplicaHostname, replica.Hostname)
+		require.Equal(t, req.ReplicaError, replica.Error)
+		require.Equal(t, req.ReplicaRelayAddress, replica.RelayAddress)
+		require.Equal(t, req.Version, replica.Version)
+		require.EqualValues(t, 10001, replica.RegionID)
+		require.False(t, replica.StoppedAt.Valid)
+		require.Zero(t, replica.DatabaseLatency)
+		require.False(t, replica.Primary)
+
+		// Re-register with most fields changed.
+		req = wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           "https://cool.proxy.coder.test",
+			WildcardHostname:    "*.cool.proxy.coder.test",
+			DerpEnabled:         false,
+			ReplicaID:           req.ReplicaID,
+			ReplicaHostname:     "venus",
+			ReplicaError:        "error",
+			ReplicaRelayAddress: "http://127.0.0.1:9090",
+			Version:             buildinfo.Version(),
+		}
+		registerRes2, err := proxyClient.RegisterWorkspaceProxy(ctx, req)
+		require.NoError(t, err)
+		require.Equal(t, registerRes1, registerRes2)
+
+		// Get the proxy to ensure nothing has changed except updated_at.
+		proxyNew, err := client.WorkspaceProxyByID(ctx, createRes.Proxy.ID)
+		require.NoError(t, err)
+		require.Equal(t, createRes.Proxy.ID, proxyNew.ID)
+		require.Equal(t, proxyName, proxyNew.Name)
+		require.Equal(t, proxyDisplayName, proxyNew.DisplayName)
+		require.Equal(t, proxyIcon, proxyNew.IconURL)
+		require.Equal(t, req.AccessURL, proxyNew.PathAppURL)
+		require.Equal(t, req.AccessURL, proxyNew.PathAppURL)
+		require.Equal(t, req.WildcardHostname, proxyNew.WildcardHostname)
+		require.Equal(t, req.DerpEnabled, proxyNew.DerpEnabled)
+		require.False(t, proxyNew.Deleted)
+
+		// Get the replica from the DB and ensure the fields have been updated,
+		// especially the updated_at.
+		replica, err = db.GetReplicaByID(ctx, req.ReplicaID)
+		require.NoError(t, err)
+		require.Equal(t, req.ReplicaID, replica.ID)
+		require.Equal(t, req.ReplicaHostname, replica.Hostname)
+		require.Equal(t, req.ReplicaError, replica.Error)
+		require.Equal(t, req.ReplicaRelayAddress, replica.RelayAddress)
+		require.Equal(t, req.Version, replica.Version)
+		require.EqualValues(t, 10001, replica.RegionID)
+		require.False(t, replica.StoppedAt.Valid)
+		require.Zero(t, replica.DatabaseLatency)
+		require.False(t, replica.Primary)
+
+		// Deregister
+		err = proxyClient.DeregisterWorkspaceProxy(ctx, wsproxysdk.DeregisterWorkspaceProxyRequest{
+			ReplicaID: req.ReplicaID,
+		})
+		require.NoError(t, err)
+
+		// Ensure the replica has been fully stopped.
+		replica, err = db.GetReplicaByID(ctx, req.ReplicaID)
+		require.NoError(t, err)
+		require.Equal(t, req.ReplicaID, replica.ID)
+		require.True(t, replica.StoppedAt.Valid)
+
+		// Re-register should fail
+		_, err = proxyClient.RegisterWorkspaceProxy(ctx, wsproxysdk.RegisterWorkspaceProxyRequest{})
+		require.Error(t, err)
+	})
+
+	t.Run("BlockMismatchingVersion", func(t *testing.T) {
+		t.Parallel()
+
+		client, _ := setup(t)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		createRes, err := client.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
+			Name: "hi",
+		})
+		require.NoError(t, err)
+
+		proxyClient := wsproxysdk.New(client.URL)
+		proxyClient.SetSessionToken(createRes.ProxyToken)
+
+		_, err = proxyClient.RegisterWorkspaceProxy(ctx, wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           "https://proxy.coder.test",
+			WildcardHostname:    "*.proxy.coder.test",
+			DerpEnabled:         true,
+			ReplicaID:           uuid.New(),
+			ReplicaHostname:     "mars",
+			ReplicaError:        "",
+			ReplicaRelayAddress: "http://127.0.0.1:8080",
+			Version:             "v0.0.0",
+		})
+		require.Error(t, err)
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+		require.Contains(t, sdkErr.Response.Message, "Version mismatch")
+	})
+
+	t.Run("ReregisterUpdateReplica", func(t *testing.T) {
+		t.Parallel()
+
+		client, db := setup(t)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		createRes, err := client.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
+			Name: "hi",
+		})
+		require.NoError(t, err)
+
+		proxyClient := wsproxysdk.New(client.URL)
+		proxyClient.SetSessionToken(createRes.ProxyToken)
+
+		req := wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           "https://proxy.coder.test",
+			WildcardHostname:    "*.proxy.coder.test",
+			DerpEnabled:         true,
+			ReplicaID:           uuid.New(),
+			ReplicaHostname:     "mars",
+			ReplicaError:        "",
+			ReplicaRelayAddress: "http://127.0.0.1:8080",
+			Version:             buildinfo.Version(),
+		}
+		_, err = proxyClient.RegisterWorkspaceProxy(ctx, req)
+		require.NoError(t, err)
+
+		// Get the replica from the DB.
+		replica, err := db.GetReplicaByID(ctx, req.ReplicaID)
+		require.NoError(t, err)
+		require.Equal(t, req.ReplicaID, replica.ID)
+
+		time.Sleep(time.Millisecond)
+
+		// Re-register with no changed fields.
+		_, err = proxyClient.RegisterWorkspaceProxy(ctx, req)
+		require.NoError(t, err)
+
+		// Get the replica from the DB and make sure updated_at has changed.
+		replica, err = db.GetReplicaByID(ctx, req.ReplicaID)
+		require.NoError(t, err)
+		require.Equal(t, req.ReplicaID, replica.ID)
+		require.Greater(t, replica.UpdatedAt.UnixNano(), replica.CreatedAt.UnixNano())
+	})
+
+	t.Run("DeregisterNonExistentReplica", func(t *testing.T) {
+		t.Parallel()
+
+		client, _ := setup(t)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		createRes, err := client.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
+			Name: "hi",
+		})
+		require.NoError(t, err)
+
+		proxyClient := wsproxysdk.New(client.URL)
+		proxyClient.SetSessionToken(createRes.ProxyToken)
+
+		err = proxyClient.DeregisterWorkspaceProxy(ctx, wsproxysdk.DeregisterWorkspaceProxyRequest{
+			ReplicaID: uuid.New(),
+		})
+		require.Error(t, err)
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
+	})
+
+	t.Run("ReturnSiblings", func(t *testing.T) {
+		t.Parallel()
+
+		client, _ := setup(t)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		createRes1, err := client.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
+			Name: "one",
+		})
+		require.NoError(t, err)
+		createRes2, err := client.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
+			Name: "two",
+		})
+		require.NoError(t, err)
+
+		// Register a replica on proxy 2. This shouldn't be returned by replicas
+		// for proxy 1.
+		proxyClient2 := wsproxysdk.New(client.URL)
+		proxyClient2.SetSessionToken(createRes2.ProxyToken)
+		_, err = proxyClient2.RegisterWorkspaceProxy(ctx, wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           "https://other.proxy.coder.test",
+			WildcardHostname:    "*.other.proxy.coder.test",
+			DerpEnabled:         true,
+			ReplicaID:           uuid.New(),
+			ReplicaHostname:     "venus",
+			ReplicaError:        "",
+			ReplicaRelayAddress: "http://127.0.0.1:9090",
+			Version:             buildinfo.Version(),
+		})
+		require.NoError(t, err)
+
+		// Register replica 1.
+		proxyClient1 := wsproxysdk.New(client.URL)
+		proxyClient1.SetSessionToken(createRes1.ProxyToken)
+		req1 := wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           "https://one.proxy.coder.test",
+			WildcardHostname:    "*.one.proxy.coder.test",
+			DerpEnabled:         true,
+			ReplicaID:           uuid.New(),
+			ReplicaHostname:     "mars1",
+			ReplicaError:        "",
+			ReplicaRelayAddress: "http://127.0.0.1:8081",
+			Version:             buildinfo.Version(),
+		}
+		registerRes1, err := proxyClient1.RegisterWorkspaceProxy(ctx, req1)
+		require.NoError(t, err)
+		require.Empty(t, registerRes1.SiblingReplicas)
+
+		// Register replica 2 and expect to get replica 1 as a sibling.
+		req2 := wsproxysdk.RegisterWorkspaceProxyRequest{
+			AccessURL:           "https://two.proxy.coder.test",
+			WildcardHostname:    "*.two.proxy.coder.test",
+			DerpEnabled:         true,
+			ReplicaID:           uuid.New(),
+			ReplicaHostname:     "mars2",
+			ReplicaError:        "",
+			ReplicaRelayAddress: "http://127.0.0.1:8082",
+			Version:             buildinfo.Version(),
+		}
+		registerRes2, err := proxyClient1.RegisterWorkspaceProxy(ctx, req2)
+		require.NoError(t, err)
+		require.Len(t, registerRes2.SiblingReplicas, 1)
+		require.Equal(t, req1.ReplicaID, registerRes2.SiblingReplicas[0].ID)
+		require.Equal(t, req1.ReplicaHostname, registerRes2.SiblingReplicas[0].Hostname)
+		require.Equal(t, req1.ReplicaRelayAddress, registerRes2.SiblingReplicas[0].RelayAddress)
+		require.EqualValues(t, 10001, registerRes2.SiblingReplicas[0].RegionID)
+
+		// Re-register replica 1 and expect to get replica 2 as a sibling.
+		registerRes1, err = proxyClient1.RegisterWorkspaceProxy(ctx, req1)
+		require.NoError(t, err)
+		require.Len(t, registerRes1.SiblingReplicas, 1)
+		require.Equal(t, req2.ReplicaID, registerRes1.SiblingReplicas[0].ID)
+		require.Equal(t, req2.ReplicaHostname, registerRes1.SiblingReplicas[0].Hostname)
+		require.Equal(t, req2.ReplicaRelayAddress, registerRes1.SiblingReplicas[0].RelayAddress)
+		require.EqualValues(t, 10001, registerRes1.SiblingReplicas[0].RegionID)
+	})
+
+	// ReturnSiblings2 tries to create 100 proxy replicas and ensures that they
+	// all return the correct number of siblings.
+	t.Run("ReturnSiblings2", func(t *testing.T) {
+		t.Parallel()
+
+		client, _ := setup(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		createRes, err := client.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
+			Name: "proxy",
+		})
+		require.NoError(t, err)
+
+		proxyClient := wsproxysdk.New(client.URL)
+		proxyClient.SetSessionToken(createRes.ProxyToken)
+
+		for i := 0; i < 100; i++ {
+			ok := false
+			for j := 0; j < 2; j++ {
+				registerRes, err := proxyClient.RegisterWorkspaceProxy(ctx, wsproxysdk.RegisterWorkspaceProxyRequest{
+					AccessURL:           "https://proxy.coder.test",
+					WildcardHostname:    "*.proxy.coder.test",
+					DerpEnabled:         true,
+					ReplicaID:           uuid.New(),
+					ReplicaHostname:     "venus",
+					ReplicaError:        "",
+					ReplicaRelayAddress: fmt.Sprintf("http://127.0.0.1:%d", 8080+i),
+					Version:             buildinfo.Version(),
+				})
+				require.NoErrorf(t, err, "register proxy %d", i)
+
+				// If the sibling replica count is wrong, try again. The impact
+				// of this not being immediate is that proxies may not function
+				// as DERP relays until they register again in 30 seconds.
+				//
+				// In the real world, replicas will not be registering this
+				// quickly. Kubernetes rolls out gradually in practice.
+				if len(registerRes.SiblingReplicas) != i {
+					t.Logf("%d: expected %d siblings, got %d", i, i, len(registerRes.SiblingReplicas))
+					time.Sleep(100 * time.Millisecond)
+					continue
+				}
+
+				ok = true
+				break
+			}
+
+			require.True(t, ok, "expected to register replica %d", i)
+		}
+	})
+}
+
 func TestIssueSignedAppToken(t *testing.T) {
 	t.Parallel()
 

enterprise/replicasync/replicasync.go

@@ -74,6 +74,7 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, ps pubsub.P
 		RelayAddress:    options.RelayAddress,
 		Version:         buildinfo.Version(),
 		DatabaseLatency: int32(databaseLatency.Microseconds()),
+		Primary:         true,
 	})
 	if err != nil {
 		return nil, xerrors.Errorf("insert replica: %w", err)
@@ -125,6 +126,20 @@ type Manager struct {
 	callback func()
 }
 
+func (m *Manager) ID() uuid.UUID {
+	return m.id
+}
+
+// UpdateNow synchronously updates replicas.
+func (m *Manager) UpdateNow(ctx context.Context) error {
+	return m.syncReplicas(ctx)
+}
+
+// PublishUpdate notifies all other replicas to update.
+func (m *Manager) PublishUpdate() error {
+	return m.pubsub.Publish(PubsubEvent, []byte(m.id.String()))
+}
+
 // updateInterval is used to determine a replicas state.
 // If the replica was updated > the time, it's considered healthy.
 // If the replica was updated < the time, it's considered stale.
@@ -299,13 +314,14 @@ func (m *Manager) syncReplicas(ctx context.Context) error {
 		Version:         m.self.Version,
 		Error:           replicaError,
 		DatabaseLatency: int32(databaseLatency.Microseconds()),
+		Primary:         m.self.Primary,
 	})
 	if err != nil {
 		return xerrors.Errorf("update replica: %w", err)
 	}
 	if m.self.Error != replica.Error {
 		// Publish an update occurred!
-		err = m.pubsub.Publish(PubsubEvent, []byte(m.self.ID.String()))
+		err = m.PublishUpdate()
 		if err != nil {
 			return xerrors.Errorf("publish replica update: %w", err)
 		}
@@ -324,12 +340,17 @@ func (m *Manager) Self() database.Replica {
 	return m.self
 }
 
-// All returns every replica, including itself.
-func (m *Manager) All() []database.Replica {
+// AllPrimary returns every primary replica (not workspace proxy replicas),
+// including itself.
+func (m *Manager) AllPrimary() []database.Replica {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 	replicas := make([]database.Replica, 0, len(m.peers))
 	for _, replica := range append(m.peers, m.self) {
+		if !replica.Primary {
+			continue
+		}
+
 		// When we assign the non-pointer to a
 		// variable it loses the reference.
 		replica := replica
@@ -338,13 +359,13 @@ func (m *Manager) All() []database.Replica {
 	return replicas
 }
 
-// Regional returns all replicas in the same region excluding itself.
-func (m *Manager) Regional() []database.Replica {
+// InRegion returns every replica in the given DERP region excluding itself.
+func (m *Manager) InRegion(regionID int32) []database.Replica {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 	replicas := make([]database.Replica, 0)
 	for _, replica := range m.peers {
-		if replica.RegionID != m.self.RegionID {
+		if replica.RegionID != regionID {
 			continue
 		}
 		replicas = append(replicas, replica)
@@ -352,6 +373,11 @@ func (m *Manager) Regional() []database.Replica {
 	return replicas
 }
 
+// Regional returns all replicas in the same region excluding itself.
+func (m *Manager) Regional() []database.Replica {
+	return m.InRegion(m.self.RegionID)
+}
+
 // SetCallback sets a function to execute whenever new peers
 // are refreshed or updated.
 func (m *Manager) SetCallback(callback func()) {

enterprise/replicasync/replicasync_test.go

@@ -66,6 +66,7 @@ func TestReplica(t *testing.T) {
 			UpdatedAt:    database.Now(),
 			Hostname:     "something",
 			RelayAddress: srv.URL,
+			Primary:      true,
 		})
 		require.NoError(t, err)
 		ctx, cancelCtx := context.WithCancel(context.Background())
@@ -110,6 +111,7 @@ func TestReplica(t *testing.T) {
 			UpdatedAt:    database.Now(),
 			Hostname:     "something",
 			RelayAddress: srv.URL,
+			Primary:      true,
 		})
 		require.NoError(t, err)
 		ctx, cancelCtx := context.WithCancel(context.Background())
@@ -137,6 +139,7 @@ func TestReplica(t *testing.T) {
 			Hostname:  "something",
 			// Fake address to dial!
 			RelayAddress: "http://127.0.0.1:1",
+			Primary:      true,
 		})
 		require.NoError(t, err)
 		ctx, cancelCtx := context.WithCancel(context.Background())
@@ -171,6 +174,7 @@ func TestReplica(t *testing.T) {
 			ID:           uuid.New(),
 			RelayAddress: srv.URL,
 			UpdatedAt:    database.Now(),
+			Primary:      true,
 		})
 		require.NoError(t, err)
 		// Publish multiple times to ensure it can handle that case.
@@ -189,6 +193,7 @@ func TestReplica(t *testing.T) {
 		_, err := db.InsertReplica(context.Background(), database.InsertReplicaParams{
 			ID:        uuid.New(),
 			UpdatedAt: database.Now().Add(-time.Hour),
+			Primary:   true,
 		})
 		require.NoError(t, err)
 		ctx, cancelCtx := context.WithCancel(context.Background())
@@ -236,8 +241,7 @@ func TestReplica(t *testing.T) {
 			server.SetCallback(func() {
 				m.Lock()
 				defer m.Unlock()
-
-				if len(server.All()) != count {
+				if len(server.AllPrimary()) != count {
 					return
 				}
 				if done {

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,