Refactor keycache to improve error handling

sreya · sreya · commit a7c63df36aba · 2024-10-01T23:25:05.000Z
- Use `cryptokeys.ErrKeyNotFound` for missing keys
- Use `cryptokeys.ErrKeyInvalid` for invalid keys
- Replace `xerrors` for streamlined error codes
diff --git a/enterprise/wsproxy/keycache.go b/enterprise/wsproxy/keycache.go
@@ -64,7 +64,7 @@ func (k *CryptoKeyCache) Signing(ctx context.Context) (codersdk.CryptoKey, error
 	latest := k.latest
 	k.keysMu.RUnlock()
 
-	now := k.Clock.Now().UTC()
+	now := k.Clock.Now()
 	if latest.CanSign(now) {
 		return latest, nil
 	}
@@ -87,14 +87,14 @@ func (k *CryptoKeyCache) Signing(ctx context.Context) (codersdk.CryptoKey, error
 	}
 
 	if !k.latest.CanSign(now) {
-		return codersdk.CryptoKey{}, xerrors.Errorf("no active keys found")
+		return codersdk.CryptoKey{}, cryptokeys.ErrKeyNotFound
 	}
 
 	return k.latest, nil
 }
 
 func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersdk.CryptoKey, error) {
-	now := k.Clock.Now().UTC()
+	now := k.Clock.Now()
 	k.keysMu.RLock()
 	if k.closed {
 		k.keysMu.RUnlock()
@@ -127,7 +127,7 @@ func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersd
 
 	key, ok = k.keys[sequence]
 	if !ok {
-		return codersdk.CryptoKey{}, xerrors.Errorf("key %d not found", sequence)
+		return codersdk.CryptoKey{}, cryptokeys.ErrKeyNotFound
 	}
 
 	return validKey(key, now)
@@ -155,7 +155,7 @@ func (k *CryptoKeyCache) fetch(ctx context.Context) (map[int32]codersdk.CryptoKe
 		return nil, codersdk.CryptoKey{}, xerrors.Errorf("get security keys: %w", err)
 	}
 
-	kmap, latest := toKeyMap(keys.CryptoKeys, k.Clock.Now().UTC())
+	kmap, latest := toKeyMap(keys.CryptoKeys, k.Clock.Now())
 	return kmap, latest, nil
 }
 
@@ -172,8 +172,8 @@ func toKeyMap(keys []codersdk.CryptoKey, now time.Time) (map[int32]codersdk.Cryp
 }
 
 func validKey(key codersdk.CryptoKey, now time.Time) (codersdk.CryptoKey, error) {
-	if !key.CanSign(now) {
-		return codersdk.CryptoKey{}, xerrors.Errorf("key %d is invalid", key.Sequence)
+	if !key.CanVerify(now) {
+		return codersdk.CryptoKey{}, cryptokeys.ErrKeyInvalid
 	}
 
 	return key, nil
diff --git a/enterprise/wsproxy/keycache_test.go b/enterprise/wsproxy/keycache_test.go
@@ -23,7 +23,7 @@ import (
 func TestCryptoKeyCache(t *testing.T) {
 	t.Parallel()
 
-	t.Run("Latest", func(t *testing.T) {
+	t.Run("Signing", func(t *testing.T) {
 		t.Parallel()
 
 		t.Run("HitsCache", func(t *testing.T) {
@@ -138,9 +138,27 @@ func TestCryptoKeyCache(t *testing.T) {
 			require.Equal(t, expected, got)
 			require.Equal(t, 1, fc.called)
 		})
+
+		t.Run("KeyNotFound", func(t *testing.T) {
+			t.Parallel()
+
+			var (
+				ctx    = testutil.Context(t, testutil.WaitShort)
+				logger = slogtest.Make(t, nil)
+				clock  = quartz.NewMock(t)
+			)
+
+			fc := newFakeCoderd(t, []codersdk.CryptoKey{})
+
+			cache, err := wsproxy.NewCryptoKeyCache(ctx, logger, wsproxysdk.New(fc.url), withClock(clock))
+			require.NoError(t, err)
+
+			_, err = cache.Verifying(ctx, 1)
+			require.ErrorIs(t, err, cryptokeys.ErrKeyNotFound)
+		})
 	})
 
-	t.Run("Version", func(t *testing.T) {
+	t.Run("Verifying", func(t *testing.T) {
 		t.Parallel()
 
 		t.Run("HitsCache", func(t *testing.T) {
@@ -241,7 +259,7 @@ func TestCryptoKeyCache(t *testing.T) {
 			require.Equal(t, 1, fc.called)
 		})
 
-		t.Run("NoInvalid", func(t *testing.T) {
+		t.Run("KeyInvalid", func(t *testing.T) {
 			t.Parallel()
 
 			var (
@@ -267,9 +285,27 @@ func TestCryptoKeyCache(t *testing.T) {
 			require.NoError(t, err)
 
 			_, err = cache.Verifying(ctx, expected.Sequence)
-			require.Error(t, err)
+			require.ErrorIs(t, err, cryptokeys.ErrKeyInvalid)
 			require.Equal(t, 1, fc.called)
 		})
+
+		t.Run("KeyNotFound", func(t *testing.T) {
+			t.Parallel()
+
+			var (
+				ctx    = testutil.Context(t, testutil.WaitShort)
+				logger = slogtest.Make(t, nil)
+				clock  = quartz.NewMock(t)
+			)
+
+			fc := newFakeCoderd(t, []codersdk.CryptoKey{})
+
+			cache, err := wsproxy.NewCryptoKeyCache(ctx, logger, wsproxysdk.New(fc.url), withClock(clock))
+			require.NoError(t, err)
+
+			_, err = cache.Verifying(ctx, 1)
+			require.ErrorIs(t, err, cryptokeys.ErrKeyNotFound)
+		})
 	})
 
 	t.Run("CacheRefreshes", func(t *testing.T) {
@@ -342,11 +378,10 @@ func TestCryptoKeyCache(t *testing.T) {
 
 		now := clock.Now()
 		expected := codersdk.CryptoKey{
-			Feature:   codersdk.CryptoKeyFeatureWorkspaceApp,
-			Secret:    "key1",
-			Sequence:  12,
-			StartsAt:  now,
-			DeletesAt: now.Add(time.Minute * 10),
+			Feature:  codersdk.CryptoKeyFeatureWorkspaceApp,
+			Secret:   "key1",
+			Sequence: 12,
+			StartsAt: now,
 		}
 		fc := newFakeCoderd(t, []codersdk.CryptoKey{
 			expected,

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ func (k *CryptoKeyCache) Signing(ctx context.Context) (codersdk.CryptoKey, error`
`64`	`64`	`latest := k.latest`
`65`	`65`	`k.keysMu.RUnlock()`
`66`	`66`
`67`		`- now := k.Clock.Now().UTC()`
	`67`	`+ now := k.Clock.Now()`
`68`	`68`	`if latest.CanSign(now) {`
`69`	`69`	`return latest, nil`
`70`	`70`	`}`
`@@ -87,14 +87,14 @@ func (k *CryptoKeyCache) Signing(ctx context.Context) (codersdk.CryptoKey, error`
`87`	`87`	`}`
`88`	`88`
`89`	`89`	`if !k.latest.CanSign(now) {`
`90`		`- return codersdk.CryptoKey{}, xerrors.Errorf("no active keys found")`
	`90`	`+ return codersdk.CryptoKey{}, cryptokeys.ErrKeyNotFound`
`91`	`91`	`}`
`92`	`92`
`93`	`93`	`return k.latest, nil`
`94`	`94`	`}`
`95`	`95`
`96`	`96`	`func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersdk.CryptoKey, error) {`
`97`		`- now := k.Clock.Now().UTC()`
	`97`	`+ now := k.Clock.Now()`
`98`	`98`	`k.keysMu.RLock()`
`99`	`99`	`if k.closed {`
`100`	`100`	`k.keysMu.RUnlock()`
`@@ -127,7 +127,7 @@ func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersd`
`127`	`127`
`128`	`128`	`key, ok = k.keys[sequence]`
`129`	`129`	`if !ok {`
`130`		`- return codersdk.CryptoKey{}, xerrors.Errorf("key %d not found", sequence)`
	`130`	`+ return codersdk.CryptoKey{}, cryptokeys.ErrKeyNotFound`
`131`	`131`	`}`
`132`	`132`
`133`	`133`	`return validKey(key, now)`
`@@ -155,7 +155,7 @@ func (k *CryptoKeyCache) fetch(ctx context.Context) (map[int32]codersdk.CryptoKe`
`155`	`155`	`return nil, codersdk.CryptoKey{}, xerrors.Errorf("get security keys: %w", err)`
`156`	`156`	`}`
`157`	`157`
`158`		`- kmap, latest := toKeyMap(keys.CryptoKeys, k.Clock.Now().UTC())`
	`158`	`+ kmap, latest := toKeyMap(keys.CryptoKeys, k.Clock.Now())`
`159`	`159`	`return kmap, latest, nil`
`160`	`160`	`}`
`161`	`161`
`@@ -172,8 +172,8 @@ func toKeyMap(keys []codersdk.CryptoKey, now time.Time) (map[int32]codersdk.Cryp`
`172`	`172`	`}`
`173`	`173`
`174`	`174`	`func validKey(key codersdk.CryptoKey, now time.Time) (codersdk.CryptoKey, error) {`
`175`		`- if !key.CanSign(now) {`
`176`		`- return codersdk.CryptoKey{}, xerrors.Errorf("key %d is invalid", key.Sequence)`
	`175`	`+ if !key.CanVerify(now) {`
	`176`	`+ return codersdk.CryptoKey{}, cryptokeys.ErrKeyInvalid`
`177`	`177`	`}`
`178`	`178`
`179`	`179`	`return key, nil`