fix: harden auth service and reduce KMS load

Author: codex

internal/dolthubauth/kms.go

@@ -5,9 +5,12 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/rand"
+	"crypto/sha256"
 	"encoding/json"
 	"fmt"
 	"strings"
+	"sync"
+	"time"
 
 	kms "cloud.google.com/go/kms/apiv1"
 	"cloud.google.com/go/kms/apiv1/kmspb"
@@ -42,8 +45,9 @@ func (c realKMSEnvelopeClient) GetCryptoKey(ctx context.Context, req *kmspb.GetC
 // GCPKMSEnvelopeCipher encrypts credentials with a local DEK and wraps the DEK
 // with Google Cloud KMS.
 type GCPKMSEnvelopeCipher struct {
-	client  kmsEnvelopeClient
-	keyName string
+	client   kmsEnvelopeClient
+	keyName  string
+	dekCache *kmsDEKCache
 }
 
 type kmsEnvelopePayload struct {
@@ -80,8 +84,9 @@ func newGCPKMSEnvelopeCipherWithClient(client kmsEnvelopeClient, keyName string)
 		return nil, fmt.Errorf("kms client is required")
 	}
 	return &GCPKMSEnvelopeCipher{
-		client:  client,
-		keyName: keyName,
+		client:   client,
+		keyName:  keyName,
+		dekCache: newKMSDEKCache(5*time.Minute, 1024),
 	}, nil
 }
 
@@ -160,14 +165,19 @@ func (c *GCPKMSEnvelopeCipher) Decrypt(ctx context.Context, ciphertext []byte, k
 		keyVersion = c.keyName
 	}
 
-	unwrapped, err := c.client.Decrypt(ctx, &kmspb.DecryptRequest{
-		Name:       kmsEnvelopeCryptoKeyName(keyVersion, c.keyName),
-		Ciphertext: payload.WrappedDEK,
-	})
-	if err != nil {
-		return nil, fmt.Errorf("unwrap dek: %w", err)
+	cacheKey := kmsEnvelopeDEKCacheKey(keyVersion, payload.WrappedDEK)
+	dek, ok := c.dekCache.Get(cacheKey)
+	if !ok {
+		unwrapped, err := c.client.Decrypt(ctx, &kmspb.DecryptRequest{
+			Name:       kmsEnvelopeCryptoKeyName(keyVersion, c.keyName),
+			Ciphertext: payload.WrappedDEK,
+		})
+		if err != nil {
+			return nil, fmt.Errorf("unwrap dek: %w", err)
+		}
+		dek = append([]byte(nil), unwrapped.GetPlaintext()...)
+		c.dekCache.Set(cacheKey, dek)
 	}
-	dek := append([]byte(nil), unwrapped.GetPlaintext()...)
 	defer clearBytes(dek)
 
 	block, err := aes.NewCipher(dek)
@@ -202,6 +212,96 @@ func kmsEnvelopeAAD(keyVersion string) []byte {
 	return []byte("wasteland:dolthub-auth:gcp-kms-envelope:" + strings.TrimSpace(keyVersion))
 }
 
+func kmsEnvelopeDEKCacheKey(keyVersion string, wrappedDEK []byte) string {
+	sum := sha256.Sum256(wrappedDEK)
+	return strings.TrimSpace(keyVersion) + ":" + fmt.Sprintf("%x", sum[:])
+}
+
+type kmsDEKCache struct {
+	mu         sync.Mutex
+	ttl        time.Duration
+	maxEntries int
+	entries    map[string]kmsDEKCacheEntry
+}
+
+type kmsDEKCacheEntry struct {
+	dek       []byte
+	expiresAt time.Time
+	usedAt    time.Time
+}
+
+func newKMSDEKCache(ttl time.Duration, maxEntries int) *kmsDEKCache {
+	if ttl <= 0 || maxEntries <= 0 {
+		return nil
+	}
+	return &kmsDEKCache{
+		ttl:        ttl,
+		maxEntries: maxEntries,
+		entries:    make(map[string]kmsDEKCacheEntry),
+	}
+}
+
+func (c *kmsDEKCache) Get(key string) ([]byte, bool) {
+	if c == nil {
+		return nil, false
+	}
+	now := time.Now()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	entry, ok := c.entries[key]
+	if !ok {
+		return nil, false
+	}
+	if now.After(entry.expiresAt) {
+		clearBytes(entry.dek)
+		delete(c.entries, key)
+		return nil, false
+	}
+	entry.usedAt = now
+	c.entries[key] = entry
+	return append([]byte(nil), entry.dek...), true
+}
+
+func (c *kmsDEKCache) Set(key string, dek []byte) {
+	if c == nil || key == "" || len(dek) == 0 {
+		return
+	}
+	now := time.Now()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if old, ok := c.entries[key]; ok {
+		clearBytes(old.dek)
+	}
+	c.entries[key] = kmsDEKCacheEntry{
+		dek:       append([]byte(nil), dek...),
+		expiresAt: now.Add(c.ttl),
+		usedAt:    now,
+	}
+	c.evictLocked(now)
+}
+
+func (c *kmsDEKCache) evictLocked(now time.Time) {
+	for key, entry := range c.entries {
+		if now.After(entry.expiresAt) {
+			clearBytes(entry.dek)
+			delete(c.entries, key)
+		}
+	}
+	for len(c.entries) > c.maxEntries {
+		var oldestKey string
+		var oldestAt time.Time
+		for key, entry := range c.entries {
+			if oldestKey == "" || entry.usedAt.Before(oldestAt) {
+				oldestKey = key
+				oldestAt = entry.usedAt
+			}
+		}
+		entry := c.entries[oldestKey]
+		clearBytes(entry.dek)
+		delete(c.entries, oldestKey)
+	}
+}
+
 func clearBytes(buf []byte) {
 	for i := range buf {
 		buf[i] = 0

internal/dolthubauth/kms_test.go

@@ -14,6 +14,7 @@ type fakeKMSEnvelopeClient struct {
 	lastDecryptName string
 	lastGetName     string
 	versionName     string
+	decryptCalls    int
 	checkErr        error
 	encryptErr      error
 	decryptErr      error
@@ -34,6 +35,7 @@ func (f *fakeKMSEnvelopeClient) Decrypt(_ context.Context, req *kmspb.DecryptReq
 	if f.decryptErr != nil {
 		return nil, f.decryptErr
 	}
+	f.decryptCalls++
 	f.lastDecryptName = req.GetName()
 	if len(req.GetCiphertext()) < len("wrapped:") || string(req.GetCiphertext()[:len("wrapped:")]) != "wrapped:" {
 		return nil, errors.New("unexpected wrapped dek")
@@ -98,6 +100,36 @@ func TestGCPKMSEnvelopeCipher_RoundTrip(t *testing.T) {
 	}
 }
 
+func TestGCPKMSEnvelopeCipher_DecryptCachesUnwrappedDEK(t *testing.T) {
+	t.Parallel()
+
+	const keyName = "projects/example-project/locations/us-central1/keyRings/example-ring/cryptoKeys/dolthub-auth-staging"
+	client := &fakeKMSEnvelopeClient{
+		versionName: keyName + "/cryptoKeyVersions/7",
+	}
+	cipher, err := newGCPKMSEnvelopeCipherWithClient(client, keyName)
+	if err != nil {
+		t.Fatalf("newGCPKMSEnvelopeCipherWithClient() error = %v", err)
+	}
+
+	encoded, keyVersion, backend, err := cipher.Encrypt(context.Background(), []byte("super-secret-token"))
+	if err != nil {
+		t.Fatalf("Encrypt() error = %v", err)
+	}
+	for range 2 {
+		plaintext, err := cipher.Decrypt(context.Background(), encoded, keyVersion, backend)
+		if err != nil {
+			t.Fatalf("Decrypt() error = %v", err)
+		}
+		if string(plaintext) != "super-secret-token" {
+			t.Fatalf("plaintext = %q", plaintext)
+		}
+	}
+	if client.decryptCalls != 1 {
+		t.Fatalf("decryptCalls = %d, want 1", client.decryptCalls)
+	}
+}
+
 func TestNewCredentialCipher_KMSPrimaryWithLegacyLocalFallback(t *testing.T) {
 	t.Parallel()
 

internal/dolthubauth/server_test.go

@@ -167,6 +167,55 @@ func TestServerRedeemInvalidJSON(t *testing.T) {
 	}
 }
 
+func TestServerRedeemExpiredTokenReturnsCORSJSON(t *testing.T) {
+	srv, err := NewServer(validConfig(), Dependencies{
+		Store:      expiredRedeemStore{},
+		KeyManager: fakeCipher{},
+	})
+	if err != nil {
+		t.Fatalf("NewServer() error = %v", err)
+	}
+
+	reqBody, _ := json.Marshal(RedeemConnectTokenRequest{
+		ConnectToken: "connect-token",
+		RedeemSecret: "redeem-secret",
+		APIKey:       "secret-token",
+		Metadata: UserMetadata{
+			RigHandle: "alice",
+			Wastelands: []WastelandConfig{{
+				Upstream: "hop/wl-commons",
+				ForkOrg:  "alice-org",
+				ForkDB:   "wl-commons",
+				Mode:     "pr",
+				Signing:  true,
+			}},
+		},
+	})
+	req := httptest.NewRequest(http.MethodPost, "/v1/connect-tokens/redeem", bytes.NewReader(reqBody))
+	req.Header.Set("Origin", "https://app.example")
+	rec := httptest.NewRecorder()
+	srv.Handler().ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusUnauthorized {
+		t.Fatalf("status = %d body = %s", rec.Code, rec.Body.String())
+	}
+	if got := rec.Header().Get("Access-Control-Allow-Origin"); got != "https://app.example" {
+		t.Fatalf("Access-Control-Allow-Origin = %q", got)
+	}
+	if got := rec.Header().Get("X-Wasteland-Auth-Error-Code"); got != "expired_connect_token" {
+		t.Fatalf("X-Wasteland-Auth-Error-Code = %q", got)
+	}
+	if !strings.Contains(rec.Body.String(), `"error_code":"expired_connect_token"`) {
+		t.Fatalf("body = %s", rec.Body.String())
+	}
+}
+
+type expiredRedeemStore struct{ fakeStore }
+
+func (f expiredRedeemStore) RedeemConnectToken(context.Context, RedeemInput) (*Connection, error) {
+	return nil, ErrExpiredConnectToken
+}
+
 func TestServerCORSRejectsUnknownOrigin(t *testing.T) {
 	srv, err := NewServer(validConfig(), Dependencies{
 		Store:      fakeStore{},

internal/dolthubauth/store.go

@@ -107,7 +107,7 @@ func (s *PostgresStore) RedeemConnectToken(ctx context.Context, input RedeemInpu
 	if usedAt != nil {
 		return nil, ErrInvalidConnectToken
 	}
-	if input.Now.After(expiresAt) {
+	if connectTokenExpired(input.Now, expiresAt) {
 		return nil, ErrExpiredConnectToken
 	}
 	if tenantID != input.TenantID || environment != input.Environment {
@@ -558,6 +558,10 @@ func isUniqueViolation(err error) bool {
 	return errors.As(err, &pgErr) && pgErr.Code == "23505"
 }
 
+func connectTokenExpired(now, expiresAt time.Time) bool {
+	return !now.Before(expiresAt)
+}
+
 func (s *PostgresStore) reapExpiredConnectTokens(ctx context.Context, now time.Time) error {
 	_, err := s.pool.Exec(ctx, `
 		DELETE FROM connect_tokens

internal/dolthubauth/store_test.go

@@ -3,6 +3,7 @@ package dolthubauth
 import (
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/jackc/pgx/v5/pgconn"
 )
@@ -36,3 +37,25 @@ func TestIsUniqueViolation(t *testing.T) {
 		}
 	})
 }
+
+func TestConnectTokenExpired(t *testing.T) {
+	expiresAt := time.Date(2026, 4, 25, 5, 30, 0, 0, time.UTC)
+
+	tests := []struct {
+		name string
+		now  time.Time
+		want bool
+	}{
+		{name: "before expiry", now: expiresAt.Add(-time.Nanosecond), want: false},
+		{name: "at expiry boundary", now: expiresAt, want: true},
+		{name: "after expiry", now: expiresAt.Add(time.Nanosecond), want: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := connectTokenExpired(tt.now, expiresAt); got != tt.want {
+				t.Fatalf("connectTokenExpired(%s, %s) = %v, want %v", tt.now, expiresAt, got, tt.want)
+			}
+		})
+	}
+}

internal/hosted/authservice_resolver_test.go

@@ -2,11 +2,15 @@ package hosted
 
 import (
 	"context"
+	"io"
 	"net/http"
+	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/gastownhall/wasteland/internal/dolthubauth"
+	"github.com/gastownhall/wasteland/internal/remote"
 )
 
 type roundTripFunc func(*http.Request) (*http.Response, error)
@@ -68,3 +72,33 @@ func TestAuthServiceWorkspaceResolver_BuildClientWarmsPendingCache(t *testing.T)
 	}
 	cache.Stop()
 }
+
+func TestPendingUpstreamCache_DoesNotRefreshWithoutReaders(t *testing.T) {
+	var calls atomic.Int32
+	provider := remote.NewDoltHubProviderWithClient(&http.Client{
+		Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+			calls.Add(1)
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Header:     make(http.Header),
+				Body:       io.NopCloser(strings.NewReader(`{"pulls":[]}`)),
+			}, nil
+		}),
+	})
+	cache := newPendingUpstreamCache(provider, "hop", "wl-commons", 10*time.Millisecond)
+	defer cache.Stop()
+
+	deadline := time.Now().Add(time.Second)
+	for calls.Load() == 0 && time.Now().Before(deadline) {
+		time.Sleep(time.Millisecond)
+	}
+	first := calls.Load()
+	if first == 0 {
+		t.Fatal("initial pending refresh did not run")
+	}
+
+	time.Sleep(50 * time.Millisecond)
+	if got := calls.Load(); got != first {
+		t.Fatalf("pending cache refreshed without readers: calls = %d, want %d", got, first)
+	}
+}

internal/hosted/shared_cache.go

@@ -58,21 +58,6 @@ func newPendingUpstreamCache(provider *remote.DoltHubProvider, upOrg, upDB strin
 
 	c.scheduleRefresh(context.Background())
 
-	c.wg.Add(1)
-	go func() {
-		defer c.wg.Done()
-		ticker := time.NewTicker(interval)
-		defer ticker.Stop()
-		for {
-			select {
-			case <-ticker.C:
-				c.scheduleRefresh(context.Background())
-			case <-c.stop:
-				return
-			}
-		}
-	}()
-
 	return c
 }