Refactor cryptokeys Fetcher to include feature param

Enhances flexibility by making the `Fetcher` interface receive a `CryptoKeyFeature` parameter. This change aligns various call sites that implement or utilize `Fetcher`, allowing for more granular queries.
coder · sreya · Oct 25, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
commit 33cdb96baa0f67e0b4d2f7f34326a436b8cc0f9c
diff --git a/cli/server.go b/cli/server.go
@@ -746,9 +746,18 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				return xerrors.Errorf("set deployment id: %w", err)
 			}
 
-			resumeKeycache, err := cryptokeys.NewSigningCache(logger, options.Database, database.CryptoKeyFeatureTailnetResume)
+			fetcher := &cryptokeys.DBFetcher{
+				DB: options.Database,
+			}
+
+			// TODO(JonA): The instantiation of this cache + coordinator seems like it should be done inside coderd so that it uses the correct context.
+			resumeKeycache, err := cryptokeys.NewSigningCache(ctx,
+				logger,
+				fetcher,
+				codersdk.CryptoKeyFeatureTailnetResume,
+			)
 			if err != nil {
-				return xerrors.Errorf("create resume token key cache: %w", err)
+				return xerrors.Errorf("create tailnet resume key cache: %w", err)
 			}
 
 			options.CoordinatorResumeTokenProvider = tailnet.NewResumeTokenKeyProvider(resumeKeycache, quartz.NewReal(), tailnet.DefaultResumeTokenExpiry)

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -448,6 +448,44 @@ func New(options *Options) *API {
 	if err != nil {
 		panic(xerrors.Errorf("get deployment ID: %w", err))
 	}
+
+	// Start a background process that rotates keys.
+	err = cryptokeys.StartRotator(ctx, options.Logger.Named("keyrotator"), options.Database)
+	if err != nil {
+		options.Logger.Fatal(ctx, "start key rotator", slog.Error(err))
+	}
+
+	fetcher := &cryptokeys.DBFetcher{
+		DB: options.Database,
+	}
+
+	if options.OIDCConvertKeyCache == nil {
+		options.OIDCConvertKeyCache, err = cryptokeys.NewSigningCache(ctx,
+			options.Logger.Named("oidc_convert_keycache"),
+			fetcher,
+			codersdk.CryptoKeyFeatureOIDCConvert,
+		)
+		must(options.Logger, "start oidc convert key cache", err)
+	}
+
+	if options.AppSigningKeyCache == nil {
+		options.AppSigningKeyCache, err = cryptokeys.NewSigningCache(ctx,
+			options.Logger.Named("app_signing_keycache"),
+			fetcher,
+			codersdk.CryptoKeyFeatureWorkspaceAppsToken,
+		)
+		must(options.Logger, "start app signing key cache", err)
+	}
+
+	if options.AppEncryptionKeyCache == nil {
+		options.AppEncryptionKeyCache, err = cryptokeys.NewEncryptionCache(ctx,
+			options.Logger.Named("app_encryption_keycache"),
+			fetcher,
+			codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey,
+		)
+		must(options.Logger, "start app encryption key cache", err)
+	}
+
 	api := &API{
 		ctx:          ctx,
 		cancel:       cancel,
@@ -484,7 +522,7 @@ func New(options *Options) *API {
 			options.Database,
 			options.Pubsub,
 		),
-		dbRolluper: options.DatabaseRolluper,
+		dbRolluper:          options.DatabaseRolluper,
 	}
 
 	f := appearance.NewDefaultFetcher(api.DeploymentValues.DocsURL.String())
@@ -613,12 +651,6 @@ func New(options *Options) *API {
 		api.Logger.Fatal(api.ctx, "failed to initialize tailnet client service", slog.Error(err))
 	}
 
-	// Start a background process that rotates keys.
-	err = cryptokeys.StartRotator(api.ctx, api.Logger.Named("keyrotator"), api.Database)
-	if err != nil {
-		api.Logger.Fatal(api.ctx, "start key rotator", slog.Error(err))
-	}
-
 	api.statsReporter = workspacestats.NewReporter(workspacestats.ReporterOptions{
 		Database:              options.Database,
 		Logger:                options.Logger.Named("workspacestats"),
@@ -1612,3 +1644,9 @@ func ReadExperiments(log slog.Logger, raw []string) codersdk.Experiments {
 	}
 	return exps
 }
+
+func must(logger slog.Logger, msg string, err error) {
+	if err != nil {
+		logger.Fatal(context.Background(), msg, slog.Error(err))
+	}
+}
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -55,7 +55,6 @@ import (
 	"github.com/coder/coder/v2/coderd/audit"
 	"github.com/coder/coder/v2/coderd/autobuild"
 	"github.com/coder/coder/v2/coderd/awsidentity"
-	"github.com/coder/coder/v2/coderd/cryptokeys"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/db2sdk"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
@@ -326,13 +325,6 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	}
 	auditor.Store(&options.Auditor)
 
-	oidcConvertKeyCache, err := cryptokeys.NewSigningCache(options.Logger.Named("oidc_convert_keycache"), options.Database, database.CryptoKeyFeatureOIDCConvert)
-	require.NoError(t, err)
-	appSigningKeyCache, err := cryptokeys.NewSigningCache(options.Logger.Named("app_signing_keycache"), options.Database, database.CryptoKeyFeatureWorkspaceAppsToken)
-	require.NoError(t, err)
-	appEncryptionKeyCache, err := cryptokeys.NewEncryptionCache(options.Logger.Named("app_encryption_keycache"), options.Database, database.CryptoKeyFeatureWorkspaceAppsAPIKey)
-	require.NoError(t, err)
-
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	lifecycleExecutor := autobuild.NewExecutor(
 		ctx,
@@ -540,9 +532,6 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			WorkspaceUsageTracker:              wuTracker,
 			NotificationsEnqueuer:              options.NotificationsEnqueuer,
 			OneTimePasscodeValidityPeriod:      options.OneTimePasscodeValidityPeriod,
-			AppSigningKeyCache:                 appSigningKeyCache,
-			AppEncryptionKeyCache:              appEncryptionKeyCache,
-			OIDCConvertKeyCache:                oidcConvertKeyCache,
 		}
 }
 

diff --git a/coderd/cryptokeys/cache.go b/coderd/cryptokeys/cache.go
@@ -25,7 +25,7 @@ var (
 )
 
 type Fetcher interface {
-	Fetch(ctx context.Context) ([]codersdk.CryptoKey, error)
+	Fetch(ctx context.Context, feature codersdk.CryptoKeyFeature) ([]codersdk.CryptoKey, error)
 }
 
 type EncryptionKeycache interface {
@@ -62,12 +62,11 @@ const (
 )
 
 type DBFetcher struct {
-	DB      database.Store
-	Feature database.CryptoKeyFeature
+	DB database.Store
 }
 
-func (d *DBFetcher) Fetch(ctx context.Context) ([]codersdk.CryptoKey, error) {
-	keys, err := d.DB.GetCryptoKeysByFeature(ctx, d.Feature)
+func (d *DBFetcher) Fetch(ctx context.Context, feature codersdk.CryptoKeyFeature) ([]codersdk.CryptoKey, error) {
+	keys, err := d.DB.GetCryptoKeysByFeature(ctx, database.CryptoKeyFeature(feature))
 	if err != nil {
 		return nil, xerrors.Errorf("get crypto keys by feature: %w", err)
 	}
@@ -198,7 +197,7 @@ func (c *cache) VerifyingKey(ctx context.Context, id string) (interface{}, error
 }
 
 func isEncryptionKeyFeature(feature codersdk.CryptoKeyFeature) bool {
-	return feature == codersdk.CryptoKeyFeatureWorkspaceApp
+	return feature == codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey
 }
 
 func isSigningKeyFeature(feature codersdk.CryptoKeyFeature) bool {
@@ -332,7 +331,7 @@ func (c *cache) refresh() {
 // cryptoKeys queries the control plane for the crypto keys.
 // Outside of initialization, this should only be called by fetch.
 func (c *cache) cryptoKeys(ctx context.Context) (map[int32]codersdk.CryptoKey, error) {
-	keys, err := c.fetcher.Fetch(ctx)
+	keys, err := c.fetcher.Fetch(ctx, c.feature)
 	if err != nil {
 		return nil, xerrors.Errorf("crypto keys: %w", err)
 	}

diff --git a/coderd/cryptokeys/cache_test.go b/coderd/cryptokeys/cache_test.go
@@ -488,7 +488,7 @@ type fakeFetcher struct {
 	called int
 }
 
-func (f *fakeFetcher) Fetch(_ context.Context) ([]codersdk.CryptoKey, error) {
+func (f *fakeFetcher) Fetch(_ context.Context, _ codersdk.CryptoKeyFeature) ([]codersdk.CryptoKey, error) {
 	f.called++
 	return f.keys, nil
 }

diff --git a/coderd/jwtutils/jwt_test.go b/coderd/jwtutils/jwt_test.go
@@ -240,7 +240,7 @@ func TestJWS(t *testing.T) {
 				StartsAt: time.Now(),
 			})
 			log     = slogtest.Make(t, nil)
-			fetcher = &cryptokeys.DBFetcher{DB: db, Feature: database.CryptoKeyFeatureOidcConvert}
+			fetcher = &cryptokeys.DBFetcher{DB: db}
 		)
 
 		cache, err := cryptokeys.NewSigningCache(ctx, log, fetcher, codersdk.CryptoKeyFeatureOIDCConvert)
@@ -331,10 +331,10 @@ func TestJWE(t *testing.T) {
 			})
 			log = slogtest.Make(t, nil)
 
-			fetcher = &cryptokeys.DBFetcher{DB: db, Feature: database.CryptoKeyFeatureWorkspaceApps}
+			fetcher = &cryptokeys.DBFetcher{DB: db}
 		)
 
-		cache, err := cryptokeys.NewEncryptionCache(ctx, log, fetcher, codersdk.CryptoKeyFeatureWorkspaceApp)
+		cache, err := cryptokeys.NewEncryptionCache(ctx, log, fetcher, codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey)
 		require.NoError(t, err)
 
 		claims := testClaims{

diff --git a/codersdk/deployment.go b/codersdk/deployment.go
@@ -3109,7 +3109,7 @@ func (c *Client) SSHConfiguration(ctx context.Context) (SSHConfigResponse, error
 type CryptoKeyFeature string
 
 const (
-	CryptoKeyFeatureWorkspaceAppAPIKey CryptoKeyFeature = "workspace_apps_api_key"
+	CryptoKeyFeatureWorkspaceAppsAPIKey CryptoKeyFeature = "workspace_apps_api_key"
 	//nolint:gosec // This denotes a type of key, not a literal.
 	CryptoKeyFeatureWorkspaceAppsToken CryptoKeyFeature = "workspace_apps_token"
 	CryptoKeyFeatureOIDCConvert        CryptoKeyFeature = "oidc_convert"

diff --git a/enterprise/coderd/workspaceproxy_test.go b/enterprise/coderd/workspaceproxy_test.go
@@ -954,7 +954,7 @@ func TestGetCryptoKeys(t *testing.T) {
 			Name: testutil.GetRandomName(t),
 		})
 
-		keys, err := proxy.SDKClient.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppAPIKey)
+		keys, err := proxy.SDKClient.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey)
 		require.NoError(t, err)
 		require.NotEmpty(t, keys)
 		require.Equal(t, 2, len(keys.CryptoKeys))
@@ -986,7 +986,7 @@ func TestGetCryptoKeys(t *testing.T) {
 		client := wsproxysdk.New(cclient.URL)
 		client.SetSessionToken(cclient.SessionToken())
 
-		_, err := client.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppAPIKey)
+		_, err := client.CryptoKeys(ctx, codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey)
 		require.Error(t, err)
 		var sdkErr *codersdk.Error
 		require.ErrorAs(t, err, &sdkErr)

diff --git a/enterprise/wsproxy/keyfetcher.go b/enterprise/wsproxy/keyfetcher.go
@@ -12,12 +12,11 @@ import (
 var _ cryptokeys.Fetcher = &ProxyFetcher{}
 
 type ProxyFetcher struct {
-	Client  *wsproxysdk.Client
-	Feature codersdk.CryptoKeyFeature
+	Client *wsproxysdk.Client
 }
 
-func (p *ProxyFetcher) Fetch(ctx context.Context) ([]codersdk.CryptoKey, error) {
-	keys, err := p.Client.CryptoKeys(ctx)
+func (p *ProxyFetcher) Fetch(ctx context.Context, feature codersdk.CryptoKeyFeature) ([]codersdk.CryptoKey, error) {
+	keys, err := p.Client.CryptoKeys(ctx, feature)
 	if err != nil {
 		return nil, xerrors.Errorf("crypto keys: %w", err)
 	}