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

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()