Refactor key caching and logging behavior

- Improve code clarity by refactoring key caching logic. - Simplify logger initialization for signing and encryption caches. - Ensure consistent closing of caches in the API and server.
coder · sreya · Oct 25, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
commit 08570b76f9ca9f1660b817b503545286c021f52d
diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go
@@ -125,7 +125,7 @@ func TestWorkspaceActivityBump(t *testing.T) {
 			}
 
 			// maxTimeDrift is how long we are willing wait for a deadline to
-			// be increased. Since it could have been bumped at the intial
+			// be increased. Since it could have been bumped at the initial
 			maxTimeDrift := testutil.WaitMedium
 
 			updatedAfter := dbtime.Now()

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -461,7 +461,7 @@ func New(options *Options) *API {
 
 	if options.OIDCConvertKeyCache == nil {
 		options.OIDCConvertKeyCache, err = cryptokeys.NewSigningCache(ctx,
-			options.Logger.Named("oidc_convert_keycache"),
+			options.Logger,
 			fetcher,
 			codersdk.CryptoKeyFeatureOIDCConvert,
 		)
@@ -470,7 +470,7 @@ func New(options *Options) *API {
 
 	if options.AppSigningKeyCache == nil {
 		options.AppSigningKeyCache, err = cryptokeys.NewSigningCache(ctx,
-			options.Logger.Named("app_signing_keycache"),
+			options.Logger,
 			fetcher,
 			codersdk.CryptoKeyFeatureWorkspaceAppsToken,
 		)
@@ -479,7 +479,7 @@ func New(options *Options) *API {
 
 	if options.AppEncryptionKeyCache == nil {
 		options.AppEncryptionKeyCache, err = cryptokeys.NewEncryptionCache(ctx,
-			options.Logger.Named("app_encryption_keycache"),
+			options.Logger,
 			fetcher,
 			codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey,
 		)
@@ -522,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())
@@ -1474,6 +1474,9 @@ func (api *API) Close() error {
 	_ = api.agentProvider.Close()
 	_ = api.statsReporter.Close()
 	_ = api.NetworkTelemetryBatcher.Close()
+	_ = api.OIDCConvertKeyCache.Close()
+	_ = api.AppSigningKeyCache.Close()
+	_ = api.AppEncryptionKeyCache.Close()
 	return nil
 }
 

diff --git a/coderd/cryptokeys/cache.go b/coderd/cryptokeys/cache.go
@@ -3,6 +3,7 @@ package cryptokeys
 import (
 	"context"
 	"encoding/hex"
+	"fmt"
 	"io"
 	"strconv"
 	"sync"
@@ -108,6 +109,7 @@ func NewSigningCache(ctx context.Context, logger slog.Logger, fetcher Fetcher,
 	if !isSigningKeyFeature(feature) {
 		return nil, xerrors.Errorf("invalid feature: %s", feature)
 	}
+	logger = logger.Named(fmt.Sprintf("%s_signing_keycache", feature))
 	return newCache(ctx, logger, fetcher, feature, opts...)
 }
 
@@ -117,6 +119,7 @@ func NewEncryptionCache(ctx context.Context, logger slog.Logger, fetcher Fetcher
 	if !isEncryptionKeyFeature(feature) {
 		return nil, xerrors.Errorf("invalid feature: %s", feature)
 	}
+	logger = logger.Named(fmt.Sprintf("%s_encryption_keycache", feature))
 	return newCache(ctx, logger, fetcher, feature, opts...)
 }
 

diff --git a/enterprise/coderd/coderdenttest/proxytest.go b/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,

diff --git a/enterprise/wsproxy/keyfetcher.go b/enterprise/wsproxy/keyfetcher.go
@@ -3,10 +3,11 @@ package wsproxy
 import (
 	"context"
 
+	"golang.org/x/xerrors"
+
 	"github.com/coder/coder/v2/coderd/cryptokeys"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/enterprise/wsproxy/wsproxysdk"
-	"golang.org/x/xerrors"
 )
 
 var _ cryptokeys.Fetcher = &ProxyFetcher{}

diff --git a/enterprise/wsproxy/wsproxy.go b/enterprise/wsproxy/wsproxy.go
@@ -93,9 +93,7 @@ type Options struct {
 	// from the dashboardURL. This should only be used in development.
 	AllowAllCors bool
 
-	StatsCollectorOptions           workspaceapps.StatsCollectorOptions
-	WorkspaceAppsEncryptionKeycache cryptokeys.EncryptionKeycache
-	WorkspaceAppsSigningKeycache    cryptokeys.SigningKeycache
+	StatsCollectorOptions workspaceapps.StatsCollectorOptions
 }
 
 func (o *Options) Validate() error {
@@ -133,6 +131,9 @@ type Server struct {
 	// the moon's token.
 	SDKClient *wsproxysdk.Client
 
+	WorkspaceAppsEncryptionKeycache cryptokeys.EncryptionKeycache
+	WorkspaceAppsSigningKeycache    cryptokeys.SigningKeycache
+
 	// DERP
 	derpMesh                *derpmesh.Mesh
 	derpMeshTLSConfig       *tls.Config
@@ -199,8 +200,28 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 
+	encryptionCache, err := cryptokeys.NewEncryptionCache(ctx,
+		opts.Logger,
+		&ProxyFetcher{Client: client},
+		codersdk.CryptoKeyFeatureWorkspaceAppsAPIKey,
+	)
+	if err != nil {
+		return nil, xerrors.Errorf("create api key encryption cache: %w", err)
+	}
+	signingCache, err := cryptokeys.NewSigningCache(ctx,
+		opts.Logger,
+		&ProxyFetcher{Client: client},
+		codersdk.CryptoKeyFeatureWorkspaceAppsToken,
+	)
+	if err != nil {
+		return nil, xerrors.Errorf("create api token signing cache: %w", err)
+	}
+
 	r := chi.NewRouter()
 	s := &Server{
+		ctx:    ctx,
+		cancel: cancel,
+
 		Options:            opts,
 		Handler:            r,
 		DashboardURL:       opts.DashboardURL,
@@ -210,8 +231,6 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		SDKClient:          client,
 		derpMesh:           derpmesh.New(opts.Logger.Named("net.derpmesh"), derpServer, meshTLSConfig),
 		derpMeshTLSConfig:  meshTLSConfig,
-		ctx:                ctx,
-		cancel:             cancel,
 	}
 
 	// Register the workspace proxy with the primary coderd instance and start a
@@ -280,8 +299,8 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 			AccessURL:     opts.AccessURL,
 			AppHostname:   opts.AppHostname,
 			Client:        client,
-			SigningKey:    opts.WorkspaceAppsSigningKeycache,
-			EncryptingKey: opts.WorkspaceAppsEncryptionKeycache,
+			SigningKey:    signingCache,
+			EncryptingKey: encryptionCache,
 			Logger:        s.Logger.Named("proxy_token_provider"),
 		},
 
@@ -432,6 +451,8 @@ func (s *Server) Close() error {
 		err = multierror.Append(err, agentProviderErr)
 	}
 	s.SDKClient.SDKClient.HTTPClient.CloseIdleConnections()
+	_ = s.WorkspaceAppsSigningKeycache.Close()
+	_ = s.WorkspaceAppsEncryptionKeycache.Close()
 	return err
 }