at-container-registry/pkg/auth/cache.go

// Package auth provides service token caching and management for AppView.
// Service tokens are JWTs issued by a user's PDS to authorize AppView to
// act on their behalf when communicating with hold services. Tokens are
// cached with automatic expiry parsing and 10-second safety margins.
package auth

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"
)

// serviceTokenEntry represents a cached service token.
type serviceTokenEntry struct {
	token     string
	expiresAt time.Time
}

// Cache stores per-(DID, hold DID) service tokens with automatic expiry.
// The zero value is not usable; construct via NewCache. A default package
// instance backs the GetServiceToken/SetServiceToken/... wrappers; tests
// that need isolation can construct their own.
type Cache struct {
	mu     sync.RWMutex
	tokens map[string]*serviceTokenEntry
}

// NewCache returns an empty Cache.
func NewCache() *Cache {
	return &Cache{tokens: make(map[string]*serviceTokenEntry)}
}

// defaultCache backs the package-level wrappers (GetServiceToken, etc.).
// Existing callers reach it transparently via those wrappers; tests that
// need isolation should construct a fresh Cache instead.
var defaultCache = NewCache()

// Get returns the cached token for (did, holdDID) and its expiry. If the
// entry is expired it is removed and (zero values, zero time) is returned.
func (c *Cache) Get(did, holdDID string) (string, time.Time) {
	cacheKey := did + ":" + holdDID

	c.mu.RLock()
	entry, exists := c.tokens[cacheKey]
	c.mu.RUnlock()

	if !exists {
		return "", time.Time{}
	}

	if time.Now().After(entry.expiresAt) {
		c.mu.Lock()
		delete(c.tokens, cacheKey)
		c.mu.Unlock()
		return "", time.Time{}
	}

	return entry.token, entry.expiresAt
}

// Set stores token for (did, holdDID), parsing its JWT exp claim and
// applying a 10s safety margin so the cache expires before the real
// token does. Falls back to a 50s TTL if the JWT can't be parsed.
func (c *Cache) Set(did, holdDID, token string) error {
	cacheKey := did + ":" + holdDID

	expiry, err := parseJWTExpiry(token)
	if err != nil {
		slog.Warn("Failed to parse JWT expiry, using default 50s", "error", err, "cacheKey", cacheKey)
		expiry = time.Now().Add(50 * time.Second)
	} else {
		expiry = expiry.Add(-10 * time.Second)
	}

	c.mu.Lock()
	c.tokens[cacheKey] = &serviceTokenEntry{
		token:     token,
		expiresAt: expiry,
	}
	c.mu.Unlock()

	slog.Debug("Cached service token",
		"cacheKey", cacheKey,
		"expiresIn", time.Until(expiry).Round(time.Second))

	return nil
}

// Invalidate removes the cached entry for (did, holdDID). No-op if absent.
func (c *Cache) Invalidate(did, holdDID string) {
	cacheKey := did + ":" + holdDID

	c.mu.Lock()
	delete(c.tokens, cacheKey)
	c.mu.Unlock()

	slog.Debug("Invalidated service token", "cacheKey", cacheKey)
}

// Stats returns total/valid/expired counts for debugging.
func (c *Cache) Stats() map[string]any {
	c.mu.RLock()
	defer c.mu.RUnlock()

	validCount := 0
	expiredCount := 0
	now := time.Now()

	for _, entry := range c.tokens {
		if now.Before(entry.expiresAt) {
			validCount++
		} else {
			expiredCount++
		}
	}

	return map[string]any{
		"total_entries":  len(c.tokens),
		"valid_tokens":   validCount,
		"expired_tokens": expiredCount,
	}
}

// CleanExpired removes all expired entries.
func (c *Cache) CleanExpired() {
	c.mu.Lock()
	defer c.mu.Unlock()

	now := time.Now()
	removed := 0

	for key, entry := range c.tokens {
		if now.After(entry.expiresAt) {
			delete(c.tokens, key)
			removed++
		}
	}

	if removed > 0 {
		slog.Debug("Cleaned expired service tokens", "count", removed)
	}
}

// Clear removes every entry. Intended for tests that need isolation
// between subtests, not for production code.
func (c *Cache) Clear() {
	c.mu.Lock()
	for k := range c.tokens {
		delete(c.tokens, k)
	}
	c.mu.Unlock()
}

// GetServiceToken returns the cached service token for (did, holdDID).
// Returns ("", zero time) if absent or expired. Delegates to the default
// package cache; tests that need isolation should construct a Cache.
func GetServiceToken(did, holdDID string) (token string, expiresAt time.Time) {
	return defaultCache.Get(did, holdDID)
}

// SetServiceToken stores token under (did, holdDID) in the default cache,
// applying the standard 10s safety margin against the JWT's exp claim.
func SetServiceToken(did, holdDID, token string) error {
	return defaultCache.Set(did, holdDID, token)
}

// InvalidateServiceToken removes (did, holdDID) from the default cache.
func InvalidateServiceToken(did, holdDID string) {
	defaultCache.Invalidate(did, holdDID)
}

// GetCacheStats returns default-cache statistics for debugging.
func GetCacheStats() map[string]any {
	return defaultCache.Stats()
}

// CleanExpiredTokens prunes expired entries from the default cache.
func CleanExpiredTokens() {
	defaultCache.CleanExpired()
}

// DefaultCache returns the package-level cache that backs the
// GetServiceToken/SetServiceToken/... wrappers. Callers (notably
// ServiceAuthFetcher) use this when they want to read back a value
// that GetOrFetchServiceToken* wrote.
func DefaultCache() *Cache {
	return defaultCache
}

// parseJWTExpiry extracts the exp claim from a JWT without verifying its
// signature. We trust tokens from the user's PDS, so signature
// verification isn't needed here.
func parseJWTExpiry(tokenString string) (time.Time, error) {
	parts := strings.Split(tokenString, ".")
	if len(parts) != 3 {
		return time.Time{}, fmt.Errorf("invalid JWT format: expected 3 parts, got %d", len(parts))
	}

	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return time.Time{}, fmt.Errorf("failed to decode JWT payload: %w", err)
	}

	var claims struct {
		Exp int64 `json:"exp"`
	}
	if err := json.Unmarshal(payload, &claims); err != nil {
		return time.Time{}, fmt.Errorf("failed to parse JWT claims: %w", err)
	}

	if claims.Exp == 0 {
		return time.Time{}, fmt.Errorf("JWT missing exp claim")
	}

	return time.Unix(claims.Exp, 0), nil
}