feat: add wsproxy endpoint for fetching keys

coder · sreya · Sep 26, 2024 · Sep 24, 2024 · Sep 24, 2024 · Sep 24, 2024
commit 808aa32caf7dd1942145d5b47c3bd42d4426b068
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -1405,6 +1405,13 @@ func (q *querier) GetCryptoKeys(ctx context.Context) ([]database.CryptoKey, erro
 	return q.db.GetCryptoKeys(ctx)
 }
 
+func (q *querier) GetCryptoKeysByFeature(ctx context.Context, feature database.CryptoKeyFeature) ([]database.CryptoKey, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceCryptoKey); err != nil {
+		return nil, err
+	}
+	return q.db.GetCryptoKeysByFeature(ctx, feature)
+}
+
 func (q *querier) GetDBCryptKeys(ctx context.Context) ([]database.DBCryptKey, error) {
 	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
 		return nil, err

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -2302,6 +2302,10 @@ func (s *MethodTestSuite) TestCryptoKeys() {
 			DeletesAt: sql.NullTime{Time: time.Now(), Valid: true},
 		}).Asserts(rbac.ResourceCryptoKey, policy.ActionUpdate)
 	}))
+	s.Run("GetCryptoKeysByFeature", s.Subtest(func(db database.Store, check *expects) {
+		check.Args(database.CryptoKeyFeatureWorkspaceApps).
+			Asserts(rbac.ResourceCryptoKey, policy.ActionRead)
+	}))
 }
 
 func (s *MethodTestSuite) TestSystemFunctions() {

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -2429,6 +2429,23 @@ func (q *FakeQuerier) GetCryptoKeys(_ context.Context) ([]database.CryptoKey, er
 	return keys, nil
 }
 
+func (q *FakeQuerier) GetCryptoKeysByFeature(_ context.Context, feature database.CryptoKeyFeature) ([]database.CryptoKey, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	keys := make([]database.CryptoKey, 0)
+	for _, key := range q.cryptoKeys {
+		if key.Feature == feature {
+			keys = append(keys, key)
+		}
+	}
+	// We want to return the highest sequence number first.
+	slices.SortFunc(keys, func(i, j database.CryptoKey) int {
+		return int(j.Sequence - i.Sequence)
+	})
+	return keys, nil
+}
+
 func (q *FakeQuerier) GetDBCryptKeys(_ context.Context) ([]database.DBCryptKey, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/crypto_keys.sql b/coderd/database/queries/crypto_keys.sql
@@ -3,6 +3,13 @@ SELECT *
 FROM crypto_keys
 WHERE secret IS NOT NULL;
 
+-- name: GetCryptoKeysByFeature :many
+SELECT *
+FROM crypto_keys
+WHERE feature = $1
+AND secret IS NOT NULL
+ORDER BY sequence DESC;
+
 -- name: GetLatestCryptoKeyByFeature :one
 SELECT *
 FROM crypto_keys

diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -243,6 +243,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 				r.Post("/app-stats", api.workspaceProxyReportAppStats)
 				r.Post("/register", api.workspaceProxyRegister)
 				r.Post("/deregister", api.workspaceProxyDeregister)
+				r.Get("/crypto-keys", api.workspaceProxyCryptoKeys)
 			})
 			r.Route("/{workspaceproxy}", func(r chi.Router) {
 				r.Use(

diff --git a/enterprise/coderd/workspaceproxy.go b/enterprise/coderd/workspaceproxy.go
@@ -710,6 +710,20 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 	go api.forceWorkspaceProxyHealthUpdate(api.ctx)
 }
 
+func (api *API) workspaceProxyCryptoKeys(rw http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
+
+	keys, err := api.Database.GetCryptoKeysByFeature(ctx, database.CryptoKeyFeatureWorkspaceApps)
+	if err != nil {
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	httpapi.Write(ctx, rw, http.StatusOK, wsproxysdk.CryptoKeysResponse{
+		CryptoKeys: fromDBCryptoKeys(keys),
+	})
+}
+
 // @Summary Deregister workspace proxy
 // @ID deregister-workspace-proxy
 // @Security CoderSessionToken
@@ -967,3 +981,17 @@ func (w *workspaceProxiesFetchUpdater) Fetch(ctx context.Context) (codersdk.Regi
 func (w *workspaceProxiesFetchUpdater) Update(ctx context.Context) error {
 	return w.updateFunc(ctx)
 }
+
+func fromDBCryptoKeys(keys []database.CryptoKey) []wsproxysdk.CryptoKey {
+	wskeys := make([]wsproxysdk.CryptoKey, 0, len(keys))
+	for _, key := range keys {
+		wskeys = append(wskeys, wsproxysdk.CryptoKey{
+			Feature:   wsproxysdk.CryptoKeyFeature(key.Feature),
+			Secret:    key.Secret.String,
+			DeletesAt: key.DeletesAt.Time,
+			Sequence:  key.Sequence,
+			StartsAt:  key.StartsAt,
+		})
+	}
+	return wskeys
+}
diff --git a/enterprise/dbcrypt/dbcrypt.go b/enterprise/dbcrypt/dbcrypt.go
@@ -321,6 +321,21 @@ func (db *dbCrypt) UpdateCryptoKeyDeletesAt(ctx context.Context, arg database.Up
 	return key, nil
 }
 
+func (db *dbCrypt) GetCryptoKeysByFeature(ctx context.Context, feature database.CryptoKeyFeature) ([]database.CryptoKey, error) {
+	keys, err := db.Store.GetCryptoKeysByFeature(ctx, feature)
+	if err != nil {
+		return nil, err
+	}
+
+	for i := range keys {
+		if err := db.decryptField(&keys[i].Secret.String, keys[i].SecretKeyID); err != nil {
+			return nil, err
+		}
+	}
+
+	return keys, nil
+}
+
 func (db *dbCrypt) encryptField(field *string, digest *sql.NullString) error {
 	// If no cipher is loaded, then we can't encrypt anything!
 	if db.ciphers == nil || db.primaryCipherDigest == "" {

diff --git a/enterprise/dbcrypt/dbcrypt_internal_test.go b/enterprise/dbcrypt/dbcrypt_internal_test.go
@@ -450,6 +450,35 @@ func TestCryptoKeys(t *testing.T) {
 		require.Equal(t, ciphers[0].HexDigest(), key.SecretKeyID.String)
 	})
 
+	t.Run("GetCryptoKeysByFeature", func(t *testing.T) {
+		t.Parallel()
+		db, crypt, ciphers := setup(t)
+		expected := dbgen.CryptoKey(t, crypt, database.CryptoKey{
+			Sequence: 2,
+			Feature:  database.CryptoKeyFeatureTailnetResume,
+			Secret:   sql.NullString{String: "test", Valid: true},
+		})
+		_ = dbgen.CryptoKey(t, crypt, database.CryptoKey{
+			Feature:  database.CryptoKeyFeatureWorkspaceApps,
+			Sequence: 43,
+		})
+		keys, err := crypt.GetCryptoKeysByFeature(ctx, database.CryptoKeyFeatureTailnetResume)
+		require.NoError(t, err)
+		require.Len(t, keys, 1)
+		require.Equal(t, "test", keys[0].Secret.String)
+		require.Equal(t, ciphers[0].HexDigest(), keys[0].SecretKeyID.String)
+		require.Equal(t, expected.Sequence, keys[0].Sequence)
+		require.Equal(t, expected.Feature, keys[0].Feature)
+
+		keys, err = db.GetCryptoKeysByFeature(ctx, database.CryptoKeyFeatureTailnetResume)
+		require.NoError(t, err)
+		require.Len(t, keys, 1)
+		requireEncryptedEquals(t, ciphers[0], keys[0].Secret.String, "test")
+		require.Equal(t, ciphers[0].HexDigest(), keys[0].SecretKeyID.String)
+		require.Equal(t, expected.Sequence, keys[0].Sequence)
+		require.Equal(t, expected.Feature, keys[0].Feature)
+	})
+
 	t.Run("DecryptErr", func(t *testing.T) {
 		t.Parallel()
 		db, crypt, ciphers := setup(t)

diff --git a/enterprise/wsproxy/wsproxysdk/wsproxysdk.go b/enterprise/wsproxy/wsproxysdk/wsproxysdk.go
@@ -204,7 +204,37 @@ type RegisterWorkspaceProxyRequest struct {
 	Version string `json:"version"`
 }
 
+type CryptoKeyFeature string
+
+const (
+	CryptoKeyFeatureWorkspaceApp  CryptoKeyFeature = "workspace_apps"
+	CryptoKeyFeatureOIDCConvert   CryptoKeyFeature = "oidc_convert"
+	CryptoKeyFeatureTailnetResume CryptoKeyFeature = "tailnet_resume"
+)
+
+type CryptoKey struct {
+	Feature   CryptoKeyFeature `json:"feature"`
+	Secret    string           `json:"secret"`
+	DeletesAt time.Time        `json:"deletes_at"`
+	Sequence  int32            `json:"sequence"`
+	StartsAt  time.Time        `json:"starts_at"`
+}
+
+func (c CryptoKey) Active(now time.Time) bool {
+	now = now.UTC()
+	isAfterStartsAt := !c.StartsAt.IsZero() && !now.Before(c.StartsAt)
+	return isAfterStartsAt && !c.Invalid(now)
+}
+
+func (c CryptoKey) Invalid(now time.Time) bool {
+	now = now.UTC()
+	noSecret := c.Secret == ""
+	afterDelete := !c.DeletesAt.IsZero() && !now.Before(c.DeletesAt.UTC())
+	return noSecret || afterDelete
+}
+
 type RegisterWorkspaceProxyResponse struct {
+	Keys                []CryptoKey      `json:"keys"`
 	AppSecurityKey      string           `json:"app_security_key"`
 	DERPMeshKey         string           `json:"derp_mesh_key"`
 	DERPRegionID        int32            `json:"derp_region_id"`
@@ -371,11 +401,6 @@ func (l *RegisterWorkspaceProxyLoop) Start(ctx context.Context) (RegisterWorkspa
 			}
 			failedAttempts = 0
 
-			// Check for consistency.
-			if originalRes.AppSecurityKey != resp.AppSecurityKey {
-				l.failureFn(xerrors.New("app security key has changed, proxy must be restarted"))
-				return
-			}
 			if originalRes.DERPMeshKey != resp.DERPMeshKey {
 				l.failureFn(xerrors.New("DERP mesh key has changed, proxy must be restarted"))
 				return
@@ -580,6 +605,27 @@ func (c *Client) DialCoordinator(ctx context.Context) (agpl.MultiAgentConn, erro
 	return ma, nil
 }
 
+type CryptoKeysResponse struct {
+	CryptoKeys []CryptoKey `json:"crypto_keys"`
+}
+
+func (c *Client) CryptoKeys(ctx context.Context) (CryptoKeysResponse, error) {
+	res, err := c.Request(ctx, http.MethodGet,
+		"/api/v2/workspaceproxies/me/crypto-keys",
+		nil,
+	)
+	if err != nil {
+		return CryptoKeysResponse{}, xerrors.Errorf("make request: %w", err)
+	}
+	defer res.Body.Close()
+
+	if res.StatusCode != http.StatusOK {
+		return CryptoKeysResponse{}, codersdk.ReadBodyAsError(res)
+	}
+	var resp CryptoKeysResponse
+	return resp, json.NewDecoder(res.Body).Decode(&resp)
+}
+
 type remoteMultiAgentHandler struct {
 	sdk      *Client
 	logger   slog.Logger