Add closed state to DBCache for safe shutdown

sreya · sreya · commit 4dfdd4fe92cc · 2024-10-01T15:15:33.000Z
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.
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 {
