Refactor key cache management for better clarity

- Consolidates key cache handling by replacing legacy key cache references with central key caches. - Enhances modularity and maintainability by using consistent key management methods. - Removes redundant `StaticKeyManager` implementation for streamlined code. - Adjusts cryptographic key generation and cache utilization across critical components.
coder · sreya · Oct 25, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
commit 0323f791865e40c6c4dc3e27ed84afcda9d6647e
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -250,6 +250,11 @@ type Options struct {
 
 	// OneTimePasscodeValidityPeriod specifies how long a one time passcode should be valid for.
 	OneTimePasscodeValidityPeriod time.Duration
+
+	// Keycaches
+	AppSigningKeyCache    cryptokeys.SigningKeycache
+	AppEncryptionKeyCache cryptokeys.EncryptionKeycache
+	OIDCConvertKeyCache   cryptokeys.SigningKeycache
 }
 
 // @title Coder API
@@ -622,12 +627,6 @@ func New(options *Options) *API {
 		api.Logger.Fatal(api.ctx, "start key rotator", slog.Error(err))
 	}
 
-	api.oauthConvertKeycache, err = cryptokeys.NewSigningCache(api.Logger.Named("oauth_convert_keycache"), api.Database, database.CryptoKeyFeatureOIDCConvert)
-	if err != nil {
-		api.Logger.Fatal(api.ctx, "failed to initialize oauth convert key cache", slog.Error(err))
-	}
-	api.workspaceAppsKeyCache = appEncryptingKeyCache
-
 	api.statsReporter = workspacestats.NewReporter(workspacestats.ReporterOptions{
 		Database:              options.Database,
 		Logger:                options.Logger.Named("workspacestats"),
@@ -1406,12 +1405,6 @@ type API struct {
 	// dbRolluper rolls up template usage stats from raw agent and app
 	// stats. This is used to provide insights in the WebUI.
 	dbRolluper *dbrollup.Rolluper
-
-	// resumeTokenKeycache is used to fetch and cache keys used for signing JWTs
-	// oauthConvertKeycache is used to fetch and cache keys used for signing JWTs
-	// during OAuth conversions. See userauth.go.convertUserToOauth.
-	oauthConvertKeycache  cryptokeys.SigningKeycache
-	workspaceAppsKeyCache cryptokeys.EncryptionKeycache
 }
 
 // Close waits for all WebSocket connections to drain before returning.

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -55,6 +55,7 @@ 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"
@@ -157,8 +158,7 @@ type Options struct {
 	DatabaseRolluper                   *dbrollup.Rolluper
 	WorkspaceUsageTrackerFlush         chan int
 	WorkspaceUsageTrackerTick          chan time.Time
-
-	NotificationsEnqueuer notifications.Enqueuer
+	NotificationsEnqueuer              notifications.Enqueuer
 }
 
 // New constructs a codersdk client connected to an in-memory API instance.
@@ -326,6 +326,13 @@ 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,
@@ -533,6 +540,9 @@ 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/rotate.go b/coderd/cryptokeys/rotate.go
@@ -231,6 +231,8 @@ func generateNewSecret(feature database.CryptoKeyFeature) (string, error) {
 	switch feature {
 	case database.CryptoKeyFeatureWorkspaceAppsAPIKey:
 		return generateKey(32)
+	case database.CryptoKeyFeatureWorkspaceAppsToken:
+		return generateKey(64)
 	case database.CryptoKeyFeatureOIDCConvert:
 		return generateKey(64)
 	case database.CryptoKeyFeatureTailnetResume:
@@ -252,6 +254,8 @@ func tokenDuration(feature database.CryptoKeyFeature) time.Duration {
 	switch feature {
 	case database.CryptoKeyFeatureWorkspaceAppsAPIKey:
 		return WorkspaceAppsTokenDuration
+	case database.CryptoKeyFeatureWorkspaceAppsToken:
+		return WorkspaceAppsTokenDuration
 	case database.CryptoKeyFeatureOIDCConvert:
 		return OIDCConvertTokenDuration
 	case database.CryptoKeyFeatureTailnetResume:

diff --git a/coderd/jwtutils/jwe.go b/coderd/jwtutils/jwe.go
@@ -130,30 +130,3 @@ func Decrypt(ctx context.Context, d DecryptKeyProvider, token string, claims Cla
 
 	return claims.Validate(options.RegisteredClaims)
 }
-
-type StaticKeyManager struct {
-	ID  string
-	Key interface{}
-}
-
-func (s StaticKeyManager) SigningKey(_ context.Context) (string, interface{}, error) {
-	return s.ID, s.Key, nil
-}
-
-func (s StaticKeyManager) VerifyingKey(_ context.Context, id string) (interface{}, error) {
-	if id != s.ID {
-		return nil, xerrors.Errorf("invalid id %q", id)
-	}
-	return s.Key, nil
-}
-
-func (s StaticKeyManager) EncryptingKey(_ context.Context) (string, interface{}, error) {
-	return s.ID, s.Key, nil
-}
-
-func (s StaticKeyManager) DecryptingKey(_ context.Context, id string) (interface{}, error) {
-	if id != s.ID {
-		return nil, xerrors.Errorf("invalid id %q", id)
-	}
-	return s.Key, nil
-}
diff --git a/coderd/userauth.go b/coderd/userauth.go
@@ -167,7 +167,7 @@ func (api *API) postConvertLoginType(rw http.ResponseWriter, r *http.Request) {
 		ToLoginType:   req.ToType,
 	}
 
-	token, err := jwtutils.Sign(dbauthz.AsKeyRotator(ctx), api.oauthConvertKeycache, claims)
+	token, err := jwtutils.Sign(dbauthz.AsKeyRotator(ctx), api.OIDCConvertKeyCache, claims)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error signing state jwt.",
@@ -1676,7 +1676,7 @@ func (api *API) convertUserToOauth(ctx context.Context, r *http.Request, db data
 		}
 	}
 	var claims OAuthConvertStateClaims
-	err = jwtutils.Verify(dbauthz.AsKeyRotator(ctx), api.oauthConvertKeycache, jwtCookie.Value, &claims)
+	err = jwtutils.Verify(dbauthz.AsKeyRotator(ctx), api.OIDCConvertKeyCache, jwtCookie.Value, &claims)
 	if xerrors.Is(err, cryptokeys.ErrKeyNotFound) || xerrors.Is(err, cryptokeys.ErrKeyInvalid) || xerrors.Is(err, jose.ErrCryptoFailure) {
 		// These errors are probably because the user is mixing 2 coder deployments.
 		return database.User{}, idpsync.HTTPError{

diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -560,7 +560,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
 	resumeTokenSigningKey, err := tailnet.GenerateResumeTokenSigningKey()
 	mgr := jwtutils.StaticKeyManager{
 		ID:  uuid.New().String(),
-		Key: resumeTokenSigningKey,
+		Key: resumeTokenSigningKey[:],
 	}
 	require.NoError(t, err)
 	resumeTokenProvider := newResumeTokenRecordingProvider(

diff --git a/coderd/workspaceapps.go b/coderd/workspaceapps.go
@@ -123,7 +123,7 @@ func (api *API) workspaceApplicationAuth(rw http.ResponseWriter, r *http.Request
 		return
 	}
 
-	encryptedAPIKey, err := jwtutils.Encrypt(ctx, api.workspaceAppsKeyCache, workspaceapps.EncryptedAPIKeyPayload{
+	encryptedAPIKey, err := jwtutils.Encrypt(ctx, api.AppEncryptionKeyCache, workspaceapps.EncryptedAPIKeyPayload{
 		APIKey: cookie.Value,
 	})
 	if err != nil {

diff --git a/coderd/workspaceapps/db.go b/coderd/workspaceapps/db.go
@@ -13,6 +13,8 @@ import (
 	"golang.org/x/exp/slices"
 	"golang.org/x/xerrors"
 
+	"github.com/go-jose/go-jose/v4/jwt"
+
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
@@ -22,7 +24,6 @@ import (
 	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/coderd/rbac/policy"
 	"github.com/coder/coder/v2/codersdk"
-	"github.com/go-jose/go-jose/v4/jwt"
 )
 
 // DBTokenProvider provides authentication and authorization for workspace apps

diff --git a/coderd/workspaceapps/db_test.go b/coderd/workspaceapps/db_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/httpmw"
+	"github.com/coder/coder/v2/coderd/jwtutils"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
 	"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
 	"github.com/coder/coder/v2/codersdk"
@@ -295,7 +296,8 @@ func Test_ResolveRequest(t *testing.T) {
 					require.Equal(t, codersdk.SignedAppTokenCookie, cookie.Name)
 					require.Equal(t, req.BasePath, cookie.Path)
 
-					parsedToken, err := api.workspaceAppsKeyCache.VerifySignedToken(cookie.Value)
+					var parsedToken workspaceapps.SignedToken
+					err := jwtutils.Verify(ctx, api.AppSigningKeyCache, cookie.Value, &parsedToken)
 					require.NoError(t, err)
 					// normalize expiry
 					require.WithinDuration(t, token.Expiry.Time(), parsedToken.Expiry.Time(), 2*time.Second)
@@ -551,7 +553,8 @@ func Test_ResolveRequest(t *testing.T) {
 			AgentID:     agentID,
 			AppURL:      appURL,
 		}
-		badTokenStr, err := api.AppSecurityKey.SignToken(badToken)
+
+		badTokenStr, err := jwtutils.Sign(ctx, api.AppSigningKeyCache, badToken)
 		require.NoError(t, err)
 
 		req := (workspaceapps.Request{
@@ -594,7 +597,8 @@ func Test_ResolveRequest(t *testing.T) {
 		require.Len(t, cookies, 1)
 		require.Equal(t, cookies[0].Name, codersdk.SignedAppTokenCookie)
 		require.NotEqual(t, cookies[0].Value, badTokenStr)
-		parsedToken, err := api.AppSecurityKey.VerifySignedToken(cookies[0].Value)
+		var parsedToken workspaceapps.SignedToken
+		err = jwtutils.Verify(ctx, api.AppSigningKeyCache, cookies[0].Value, &parsedToken)
 		require.NoError(t, err)
 		require.Equal(t, appNameOwner, parsedToken.AppSlugOrPort)
 	})

diff --git a/coderd/workspaceapps/token_test.go b/coderd/workspaceapps/token_test.go
@@ -7,9 +7,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-jose/go-jose/v4/jwt"
+
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/testutil"
-	"github.com/go-jose/go-jose/v4/jwt"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"

diff --git a/codersdk/workspacesdk/connector_internal_test.go b/codersdk/workspacesdk/connector_internal_test.go
@@ -168,7 +168,7 @@ func TestTailnetAPIConnector_ResumeToken(t *testing.T) {
 	require.NoError(t, err)
 	mgr := jwtutils.StaticKeyManager{
 		ID:  uuid.New().String(),
-		Key: resumeTokenSigningKey,
+		Key: resumeTokenSigningKey[:],
 	}
 	resumeTokenProvider := tailnet.NewResumeTokenKeyProvider(mgr, clock, time.Hour)
 	svc, err := tailnet.NewClientService(tailnet.ClientServiceOptions{
@@ -287,7 +287,7 @@ func TestTailnetAPIConnector_ResumeTokenFailure(t *testing.T) {
 	require.NoError(t, err)
 	mgr := jwtutils.StaticKeyManager{
 		ID:  uuid.New().String(),
-		Key: resumeTokenSigningKey,
+		Key: resumeTokenSigningKey[:],
 	}
 	resumeTokenProvider := tailnet.NewResumeTokenKeyProvider(mgr, clock, time.Hour)
 	svc, err := tailnet.NewClientService(tailnet.ClientServiceOptions{

diff --git a/enterprise/wsproxy/wsproxy.go b/enterprise/wsproxy/wsproxy.go
@@ -31,6 +31,7 @@ import (
 	"github.com/coder/coder/v2/buildinfo"
 	"github.com/coder/coder/v2/cli/cliutil"
 	"github.com/coder/coder/v2/coderd"
+	"github.com/coder/coder/v2/coderd/cryptokeys"
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/coderd/tracing"
@@ -92,7 +93,9 @@ type Options struct {
 	// from the dashboardURL. This should only be used in development.
 	AllowAllCors bool
 
-	StatsCollectorOptions workspaceapps.StatsCollectorOptions
+	StatsCollectorOptions           workspaceapps.StatsCollectorOptions
+	WorkspaceAppsEncryptionKeycache cryptokeys.EncryptionKeycache
+	WorkspaceAppsSigningKeycache    cryptokeys.SigningKeycache
 }
 
 func (o *Options) Validate() error {
@@ -240,11 +243,6 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		return nil, xerrors.Errorf("handle register: %w", err)
 	}
 
-	secKey, err := workspaceapps.KeyFromString(regResp.AppSecurityKey)
-	if err != nil {
-		return nil, xerrors.Errorf("parse app security key: %w", err)
-	}
-
 	agentProvider, err := coderd.NewServerTailnet(ctx,
 		s.Logger,
 		nil,
@@ -277,14 +275,14 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		HostnameRegex: opts.AppHostnameRegex,
 		RealIPConfig:  opts.RealIPConfig,
 		SignedTokenProvider: &TokenProvider{
-			DashboardURL: opts.DashboardURL,
-			AccessURL:    opts.AccessURL,
-			AppHostname:  opts.AppHostname,
-			Client:       client,
-			SecurityKey:  secKey,
-			Logger:       s.Logger.Named("proxy_token_provider"),
+			DashboardURL:  opts.DashboardURL,
+			AccessURL:     opts.AccessURL,
+			AppHostname:   opts.AppHostname,
+			Client:        client,
+			SigningKey:    opts.WorkspaceAppsSigningKeycache,
+			EncryptingKey: opts.WorkspaceAppsEncryptionKeycache,
+			Logger:        s.Logger.Named("proxy_token_provider"),
 		},
-		AppSecurityKey: secKey,
 
 		DisablePathApps:  opts.DisablePathApps,
 		SecureAuthCookie: opts.SecureAuthCookie,

diff --git a/tailnet/resume_test.go b/tailnet/resume_test.go
@@ -188,6 +188,6 @@ func TestResumeTokenKeyProvider(t *testing.T) {
 func newKeySigner(key tailnet.ResumeTokenSigningKey) jwtutils.SigningKeyManager {
 	return jwtutils.StaticKeyManager{
 		ID:  uuid.New().String(),
-		Key: key,
+		Key: key[:],
 	}
 }