Add closed state to DBCache for safe shutdown

Introduce `closed` state to the DBCache to ensure operations such as Verifying and Signing are safely terminated when the cache is closed. This update prevents resource access and ensures proper shutdown sequences, enhancing system reliability.
coder · sreya · Oct 1, 2024 · Sep 26, 2024 · Sep 26, 2024 · Sep 26, 2024
commit 4dfdd4fe92cc262d16b68149b4609086e29cba9b
diff --git a/coderd/cryptokeys/dbkeycache.go b/coderd/cryptokeys/dbkeycache.go
@@ -31,6 +31,7 @@ type DBCache struct {
 	timer     *quartz.Timer
 	// invalidateAt is the time at which the keys cache should be invalidated.
 	invalidateAt time.Time
+	closed       bool
 }
 
 type DBCacheOption func(*DBCache)
@@ -64,8 +65,13 @@ func NewDBCache(logger slog.Logger, db database.Store, feature database.CryptoKe
 // it is neither deleted nor has breached its deletion date. It should only be
 // used for verifying or decrypting payloads. To sign/encrypt call Signing.
 func (d *DBCache) Verifying(ctx context.Context, sequence int32) (codersdk.CryptoKey, error) {
-	now := d.clock.Now()
 	d.keysMu.RLock()
+	if d.closed {
+		d.keysMu.RUnlock()
+		return codersdk.CryptoKey{}, ErrClosed
+	}
+
+	now := d.clock.Now()
 	key, ok := d.keys[sequence]
 	d.keysMu.RUnlock()
 	if ok {
@@ -97,6 +103,12 @@ func (d *DBCache) Verifying(ctx context.Context, sequence int32) (codersdk.Crypt
 // both past its start time and before its deletion time.
 func (d *DBCache) Signing(ctx context.Context) (codersdk.CryptoKey, error) {
 	d.keysMu.RLock()
+
+	if d.closed {
+		d.keysMu.RUnlock()
+		return codersdk.CryptoKey{}, ErrClosed
+	}
+
 	latest := d.latestKey
 	d.keysMu.RUnlock()
 
@@ -185,5 +197,11 @@ func (d *DBCache) newTimer() *quartz.Timer {
 func (d *DBCache) Close() {
 	d.keysMu.Lock()
 	defer d.keysMu.Unlock()
+
+	if d.closed {
+		return
+	}
+
 	d.timer.Stop()
+	d.closed = true
 }
diff --git a/coderd/cryptokeys/dbkeycache_test.go b/coderd/cryptokeys/dbkeycache_test.go
@@ -104,4 +104,40 @@ func TestDBKeyCache(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, db2sdk.CryptoKey(expectedKey), got)
 	})
+
+	t.Run("Closed", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			db, _  = dbtestutil.NewDB(t)
+			clock  = quartz.NewMock(t)
+			ctx    = testutil.Context(t, testutil.WaitShort)
+			logger = slogtest.Make(t, nil)
+		)
+
+		expectedKey := dbgen.CryptoKey(t, db, database.CryptoKey{
+			Feature:  database.CryptoKeyFeatureWorkspaceApps,
+			Sequence: 10,
+			StartsAt: clock.Now(),
+		})
+
+		k := cryptokeys.NewDBCache(logger, db, database.CryptoKeyFeatureWorkspaceApps, cryptokeys.WithDBCacheClock(clock))
+		defer k.Close()
+
+		got, err := k.Signing(ctx)
+		require.NoError(t, err)
+		require.Equal(t, db2sdk.CryptoKey(expectedKey), got)
+
+		got, err = k.Verifying(ctx, expectedKey.Sequence)
+		require.NoError(t, err)
+		require.Equal(t, db2sdk.CryptoKey(expectedKey), got)
+
+		k.Close()
+
+		_, err = k.Signing(ctx)
+		require.ErrorIs(t, err, cryptokeys.ErrClosed)
+
+		_, err = k.Verifying(ctx, expectedKey.Sequence)
+		require.ErrorIs(t, err, cryptokeys.ErrClosed)
+	})
 }
diff --git a/coderd/cryptokeys/keycache.go b/coderd/cryptokeys/keycache.go
@@ -8,9 +8,11 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 )
 
-var ErrKeyNotFound = xerrors.New("key not found")
-
-var ErrKeyInvalid = xerrors.New("key is invalid for use")
+var (
+	ErrKeyNotFound = xerrors.New("key not found")
+	ErrKeyInvalid  = xerrors.New("key is invalid for use")
+	ErrClosed      = xerrors.New("closed")
+)
 
 // Keycache provides an abstraction for fetching signing keys.
 type Keycache interface {