feat: enable key rotation (#15066)

This PR contains the remaining logic necessary to hook up key rotation
to the product.

enterprise/coderd/coderdenttest/proxytest.go

@@ -65,6 +65,8 @@ type WorkspaceProxy struct {
 // owner client. If a token is provided, the proxy will become a replica of the
 // existing proxy region.
 func NewWorkspaceProxyReplica(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Client, options *ProxyOptions) WorkspaceProxy {
+	t.Helper()
+
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	t.Cleanup(cancelFunc)
 
@@ -142,8 +144,10 @@ func NewWorkspaceProxyReplica(t *testing.T, coderdAPI *coderd.API, owner *coders
 		statsCollectorOptions.Flush = options.FlushStats
 	}
 
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug).With(slog.F("server_url", serverURL.String()))
+
 	wssrv, err := wsproxy.New(ctx, &wsproxy.Options{
-		Logger:            slogtest.Make(t, nil).Leveled(slog.LevelDebug).With(slog.F("server_url", serverURL.String())),
+		Logger:            logger,
 		Experiments:       options.Experiments,
 		DashboardURL:      coderdAPI.AccessURL,
 		AccessURL:         accessURL,

enterprise/coderd/workspaceproxy.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"slices"
 	"strings"
 	"time"
 
@@ -33,6 +34,13 @@ import (
 	"github.com/coder/coder/v2/enterprise/wsproxy/wsproxysdk"
 )
 
+// whitelistedCryptoKeyFeatures is a list of crypto key features that are
+// allowed to be queried with workspace proxies.
+var whitelistedCryptoKeyFeatures = []database.CryptoKeyFeature{
+	database.CryptoKeyFeatureWorkspaceAppsToken,
+	database.CryptoKeyFeatureWorkspaceAppsAPIKey,
+}
+
 // forceWorkspaceProxyHealthUpdate forces an update of the proxy health.
 // This is useful when a proxy is created or deleted. Errors will be logged.
 func (api *API) forceWorkspaceProxyHealthUpdate(ctx context.Context) {
@@ -700,7 +708,6 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 	}
 
 	httpapi.Write(ctx, rw, http.StatusCreated, wsproxysdk.RegisterWorkspaceProxyResponse{
-		AppSecurityKey:      api.AppSecurityKey.String(),
 		DERPMeshKey:         api.DERPServer.MeshKey(),
 		DERPRegionID:        regionID,
 		DERPMap:             api.AGPL.DERPMap(),
@@ -721,13 +728,29 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 // @Security CoderSessionToken
 // @Produce json
 // @Tags Enterprise
+// @Param feature query string true "Feature key"
 // @Success 200 {object} wsproxysdk.CryptoKeysResponse
 // @Router /workspaceproxies/me/crypto-keys [get]
 // @x-apidocgen {"skip": true}
 func (api *API) workspaceProxyCryptoKeys(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
-	keys, err := api.Database.GetCryptoKeysByFeature(ctx, database.CryptoKeyFeatureWorkspaceApps)
+	feature := database.CryptoKeyFeature(r.URL.Query().Get("feature"))
+	if feature == "" {
+		httpapi.Write(r.Context(), rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Missing feature query parameter.",
+		})
+		return
+	}
+
+	if !slices.Contains(whitelistedCryptoKeyFeatures, feature) {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: fmt.Sprintf("Invalid feature: %q", feature),
+		})
+		return
+	}
+
+	keys, err := api.Database.GetCryptoKeysByFeature(ctx, feature)
 	if err != nil {
 		httpapi.InternalServerError(rw, err)
 		return

enterprise/coderd/workspaceproxy_test.go

@@ -320,7 +320,6 @@ func TestProxyRegisterDeregister(t *testing.T) {
 		}
 		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)
@@ -609,11 +608,8 @@ func TestProxyRegisterDeregister(t *testing.T) {
 func TestIssueSignedAppToken(t *testing.T) {
 	t.Parallel()
 
-	db, pubsub := dbtestutil.NewDB(t)
 	client, user := coderdenttest.New(t, &coderdenttest.Options{
 		Options: &coderdtest.Options{
-			Database:                 db,
-			Pubsub:                   pubsub,
 			IncludeProvisionerDaemon: true,
 		},
 		LicenseOptions: &coderdenttest.LicenseOptions{
@@ -716,6 +712,10 @@ func TestReconnectingPTYSignedToken(t *testing.T) {
 		closer.Close()
 	})
 
+	_ = dbgen.CryptoKey(t, db, database.CryptoKey{
+		Feature: database.CryptoKeyFeatureWorkspaceAppsToken,
+	})
+
 	// Create a workspace + apps
 	authToken := uuid.NewString()
 	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
@@ -915,51 +915,86 @@ func TestGetCryptoKeys(t *testing.T) {
 		now := time.Now()
 
 		expectedKey1 := dbgen.CryptoKey(t, db, database.CryptoKey{
-			Feature:  database.CryptoKeyFeatureWorkspaceApps,
+			Feature:  database.CryptoKeyFeatureWorkspaceAppsAPIKey,
 			StartsAt: now.Add(-time.Hour),
 			Sequence: 2,
 		})
-		key1 := db2sdk.CryptoKey(expectedKey1)
+		encryptionKey := db2sdk.CryptoKey(expectedKey1)
 
 		expectedKey2 := dbgen.CryptoKey(t, db, database.CryptoKey{
-			Feature:  database.CryptoKeyFeatureWorkspaceApps,
+			Feature:  database.CryptoKeyFeatureWorkspaceAppsToken,
 			StartsAt: now,
 			Sequence: 3,
 		})
-		key2 := db2sdk.CryptoKey(expectedKey2)
+		signingKey := db2sdk.CryptoKey(expectedKey2)
 
 		// Create a deleted key.
 		_ = dbgen.CryptoKey(t, db, database.CryptoKey{
-			Feature:  database.CryptoKeyFeatureWorkspaceApps,
+			Feature:  database.CryptoKeyFeatureWorkspaceAppsAPIKey,
 			StartsAt: now.Add(-time.Hour),
 			Secret: sql.NullString{
 				String: "secret1",
 				Valid:  false,
 			},
-			Sequence: 1,
-		})
-
-		// Create a key with different features.
-		_ = dbgen.CryptoKey(t, db, database.CryptoKey{
-			Feature:  database.CryptoKeyFeatureTailnetResume,
-			StartsAt: now.Add(-time.Hour),
-			Sequence: 1,
-		})
-		_ = dbgen.CryptoKey(t, db, database.CryptoKey{
-			Feature:  database.CryptoKeyFeatureOidcConvert,
-			StartsAt: now.Add(-time.Hour),
-			Sequence: 1,
+			Sequence: 4,
 		})
 
 		proxy := coderdenttest.NewWorkspaceProxyReplica(t, api, cclient, &coderdenttest.ProxyOptions{
 			Name: testutil.GetRandomName(t),
 		})
 
-		keys, err := proxy.SDKClient.CryptoKeys(ctx)
+		keys, err := proxy.SDKClient.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey)
 		require.NoError(t, err)
 		require.NotEmpty(t, keys)
+		// 1 key is generated on startup, the other we manually generated.
 		require.Equal(t, 2, len(keys.CryptoKeys))
-		requireContainsKeys(t, keys.CryptoKeys, key1, key2)
+		requireContainsKeys(t, keys.CryptoKeys, encryptionKey)
+		requireNotContainsKeys(t, keys.CryptoKeys, signingKey)
+
+		keys, err = proxy.SDKClient.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppsToken)
+		require.NoError(t, err)
+		require.NotEmpty(t, keys)
+		// 1 key is generated on startup, the other we manually generated.
+		require.Equal(t, 2, len(keys.CryptoKeys))
+		requireContainsKeys(t, keys.CryptoKeys, signingKey)
+		requireNotContainsKeys(t, keys.CryptoKeys, encryptionKey)
+	})
+
+	t.Run("InvalidFeature", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		db, pubsub := dbtestutil.NewDB(t)
+		cclient, _, api, _ := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database:                 db,
+				Pubsub:                   pubsub,
+				IncludeProvisionerDaemon: true,
+			},
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureWorkspaceProxy: 1,
+				},
+			},
+		})
+
+		proxy := coderdenttest.NewWorkspaceProxyReplica(t, api, cclient, &coderdenttest.ProxyOptions{
+			Name: testutil.GetRandomName(t),
+		})
+
+		_, err := proxy.SDKClient.CryptoKeys(ctx, codersdk.CryptoKeyFeatureOIDCConvert)
+		require.Error(t, err)
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+		_, err = proxy.SDKClient.CryptoKeys(ctx, codersdk.CryptoKeyFeatureTailnetResume)
+		require.Error(t, err)
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+		_, err = proxy.SDKClient.CryptoKeys(ctx, "invalid")
+		require.Error(t, err)
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
 	})
 
 	t.Run("Unauthorized", func(t *testing.T) {
@@ -987,7 +1022,7 @@ func TestGetCryptoKeys(t *testing.T) {
 		client := wsproxysdk.New(cclient.URL)
 		client.SetSessionToken(cclient.SessionToken())
 
-		_, err := client.CryptoKeys(ctx)
+		_, err := client.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey)
 		require.Error(t, err)
 		var sdkErr *codersdk.Error
 		require.ErrorAs(t, err, &sdkErr)
@@ -995,6 +1030,18 @@ func TestGetCryptoKeys(t *testing.T) {
 	})
 }
 
+func requireNotContainsKeys(t *testing.T, keys []codersdk.CryptoKey, unexpected ...codersdk.CryptoKey) {
+	t.Helper()
+
+	for _, unexpectedKey := range unexpected {
+		for _, key := range keys {
+			if key.Feature == unexpectedKey.Feature && key.Sequence == unexpectedKey.Sequence {
+				t.Fatalf("unexpected key %+v found", unexpectedKey)
+			}
+		}
+	}
+}
+
 func requireContainsKeys(t *testing.T, keys []codersdk.CryptoKey, expected ...codersdk.CryptoKey) {
 	t.Helper()