Update cryptokey feature test and migration logic

- Enhance testing scenarios for cryptokey features including cases with no keys and specific key states. - Update tests to ensure new cryptokey features are handled correctly. - Remove outdated migration scripts for cryptokey features as they are not required anymore. - Refactor workspace proxy keys to only allow whitelisted cryptokey features, improving security and stability.
coder · sreya · Oct 25, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
commit fa9a75d9d3278faff14c7387dce1ca2285e89199
diff --git a/coderd/cryptokeys/rotate_internal_test.go b/coderd/cryptokeys/rotate_internal_test.go
@@ -365,9 +365,11 @@ func Test_rotateKeys(t *testing.T) {
 
 		now := dbnow(clock)
 
-		// We'll test a scenario where one feature has no valid keys.
-		// Another has a key that should be rotate. And one that
-		// has a valid key that shouldn't trigger an action.
+		// We'll test a scenario where:
+		// - One feature has no valid keys.
+		// - One has a key that should be rotated.
+		// - One has a valid key that shouldn't trigger an action.
+		// - One has no keys at all.
 		_ = dbgen.CryptoKey(t, db, database.CryptoKey{
 			Feature:  database.CryptoKeyFeatureTailnetResume,
 			StartsAt: now.Add(-keyDuration),
@@ -377,6 +379,7 @@ func Test_rotateKeys(t *testing.T) {
 				Valid:  false,
 			},
 		})
+		// Generate another deleted key to ensure we insert after the latest sequence.
 		deletedKey := dbgen.CryptoKey(t, db, database.CryptoKey{
 			Feature:  database.CryptoKeyFeatureTailnetResume,
 			StartsAt: now.Add(-keyDuration),
@@ -406,7 +409,7 @@ func Test_rotateKeys(t *testing.T) {
 
 		keys, err := db.GetCryptoKeys(ctx)
 		require.NoError(t, err)
-		require.Len(t, keys, 4)
+		require.Len(t, keys, 5)
 
 		kbf, err := keysByFeature(keys, database.AllCryptoKeyFeatureValues())
 		require.NoError(t, err)
@@ -418,12 +421,14 @@ func Test_rotateKeys(t *testing.T) {
 		// No existing key for tailnet resume should've
 		// caused a key to be inserted.
 		require.Len(t, kbf[database.CryptoKeyFeatureTailnetResume], 1)
+		require.Len(t, kbf[database.CryptoKeyFeatureWorkspaceAppsToken], 1)
 
 		oidcKey := kbf[database.CryptoKeyFeatureOIDCConvert][0]
 		tailnetKey := kbf[database.CryptoKeyFeatureTailnetResume][0]
+		appTokenKey := kbf[database.CryptoKeyFeatureWorkspaceAppsToken][0]
 		requireKey(t, oidcKey, database.CryptoKeyFeatureOIDCConvert, now, nullTime, validKey.Sequence)
 		requireKey(t, tailnetKey, database.CryptoKeyFeatureTailnetResume, now, nullTime, deletedKey.Sequence+1)
-
+		requireKey(t, appTokenKey, database.CryptoKeyFeatureWorkspaceAppsToken, now, nullTime, 1)
 		newKey := kbf[database.CryptoKeyFeatureWorkspaceAppsAPIKey][0]
 		oldKey := kbf[database.CryptoKeyFeatureWorkspaceAppsAPIKey][1]
 		if newKey.Sequence == rotatedKey.Sequence {
@@ -589,6 +594,8 @@ func requireKey(t *testing.T, key database.CryptoKey, feature database.CryptoKey
 	switch key.Feature {
 	case database.CryptoKeyFeatureOIDCConvert:
 		require.Len(t, secret, 64)
+	case database.CryptoKeyFeatureWorkspaceAppsToken:
+		require.Len(t, secret, 64)
 	case database.CryptoKeyFeatureWorkspaceAppsAPIKey:
 		require.Len(t, secret, 32)
 	case database.CryptoKeyFeatureTailnetResume:

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -1050,6 +1050,8 @@ func newCryptoKeySecret(feature database.CryptoKeyFeature) (string, error) {
 	switch feature {
 	case database.CryptoKeyFeatureWorkspaceAppsAPIKey:
 		return generateCryptoKey(32)
+	case database.CryptoKeyFeatureWorkspaceAppsToken:
+		return generateCryptoKey(64)
 	case database.CryptoKeyFeatureOIDCConvert:
 		return generateCryptoKey(64)
 	case database.CryptoKeyFeatureTailnetResume:

diff --git a/coderd/database/migrations/000264_cryptokey_features.down.sql b/coderd/database/migrations/000264_cryptokey_features.down.sql
diff --git a/coderd/database/migrations/000264_cryptokey_features.up.sql b/coderd/database/migrations/000264_cryptokey_features.up.sql
diff --git a/enterprise/coderd/workspaceproxy.go b/enterprise/coderd/workspaceproxy.go
@@ -34,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) {
@@ -736,7 +743,7 @@ func (api *API) workspaceProxyCryptoKeys(rw http.ResponseWriter, r *http.Request
 		return
 	}
 
-	if !slices.Contains(database.AllCryptoKeyFeatureValues(), feature) {
+	if !slices.Contains(whitelistedCryptoKeyFeatures, feature) {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
 			Message: fmt.Sprintf("Invalid feature: %q", feature),
 		})

diff --git a/enterprise/coderd/workspaceproxy_test.go b/enterprise/coderd/workspaceproxy_test.go
@@ -918,14 +918,14 @@ func TestGetCryptoKeys(t *testing.T) {
 			StartsAt: now.Add(-time.Hour),
 			Sequence: 2,
 		})
-		key1 := db2sdk.CryptoKey(expectedKey1)
+		encryptionKey := db2sdk.CryptoKey(expectedKey1)
 
 		expectedKey2 := dbgen.CryptoKey(t, db, database.CryptoKey{
-			Feature:  database.CryptoKeyFeatureWorkspaceAppsAPIKey,
+			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{
@@ -935,19 +935,7 @@ func TestGetCryptoKeys(t *testing.T) {
 				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{
@@ -957,8 +945,53 @@ func TestGetCryptoKeys(t *testing.T) {
 		keys, err := proxy.SDKClient.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey)
 		require.NoError(t, err)
 		require.NotEmpty(t, keys)
+		// 1 key is generated on startup, the other is the one we generated for our test.
 		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)
+		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) {
@@ -994,6 +1027,18 @@ func TestGetCryptoKeys(t *testing.T) {
 	})
 }
 
+func requireNotContainsKeys(t *testing.T, keys []codersdk.CryptoKey, unexpected ...codersdk.CryptoKey) {
+	t.Helper()
+
+	for _, expectedKey := range unexpected {
+		for _, key := range keys {
+			if key.Feature == expectedKey.Feature && key.Sequence == expectedKey.Sequence {
+				t.Fatalf("unexpected key %+v found", expectedKey)
+			}
+		}
+	}
+}
+
 func requireContainsKeys(t *testing.T, keys []codersdk.CryptoKey, expected ...codersdk.CryptoKey) {
 	t.Helper()