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 a safety margin
// (ServiceTokenSafetyMargin) so the cache stops serving a token before the
// hold would reject it.
package auth

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

// ServiceTokenSafetyMargin is how far ahead of a service token's real exp the
// AppView stops treating it as usable. Every cache that holds a service token
// (this one, and the registry middleware's per-process validation cache)
// subtracts it, and the registry JWT's exp is stamped from the value this
// package returns, so the JWT never outlives the credential behind it.
//
// It is 60s because that is exactly distribution's token.Leeway: the registry
// auth package accepts a registry JWT for 60s past its exp. With a 60s margin
// the JWT is stamped at (service token exp - 60s), so the last moment
// distribution will accept it is the service token's real exp. Shrinking this
// below distribution's leeway reopens the window this constant closes: a client
// would hold an accepted JWT while the service token behind it is already dead,
// the hold would answer 403 "token has expired", and the pull would fail
// instead of re-authenticating.
const ServiceTokenSafetyMargin = 60 * time.Second

// unparsableTokenTTL is how long a token whose exp claim could not be read is
// cached. PDS-granted service tokens are requested with a 5 minute expiry (see
// servicetoken.go), so a fixed 50s is comfortably inside any plausible real
// lifetime and the next request re-mints.
const unparsableTokenTTL = 50 * time.Second

// 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 ServiceTokenSafetyMargin so the cache expires before the real
// token does. Falls back to unparsableTokenTTL if the JWT can't be parsed.
//
// A PDS is free to grant less than the margin (ATCR asks for 5 minutes;
// reference PDSes grant up to an hour, others may grant less). Subtracting a
// 60s margin from a 30s token would store an entry that is already expired,
// which Get would evict on sight, so every single request would re-mint: a
// refetch storm against the user's PDS. In that case the entry is kept for half
// of whatever life the token actually has instead, which is always positive
// while the token is alive and still leaves headroom proportional to it. A
// token that arrives already expired gets a past expiry, which is correct: it
// is unusable and the next call must mint a new one.
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 fallback TTL",
			"error", err, "cacheKey", cacheKey, "ttl", unparsableTokenTTL)
		expiry = time.Now().Add(unparsableTokenTTL)
	} else if remaining := time.Until(expiry); remaining <= ServiceTokenSafetyMargin {
		slog.Warn("PDS granted a service token shorter than the safety margin",
			"cacheKey", cacheKey,
			"grantedLife", remaining.Round(time.Second),
			"margin", ServiceTokenSafetyMargin)
		expiry = time.Now().Add(remaining / 2)
	} else {
		expiry = expiry.Add(-ServiceTokenSafetyMargin)
	}

	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 ServiceTokenSafetyMargin 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
}

// ServiceTokenExpiry reports the exp claim of a PDS-issued service token,
// without verifying its signature (we trust tokens minted by the user's PDS,
// and the hold verifies them anyway).
//
// Exported for callers that hold a service token outside this cache and must
// not keep it past its real life. The registry middleware's validation cache is
// the one that matters: it used to pin every token for a flat 45s, so a token
// with 12s left was handed to the hold for another 33s after it died.
func ServiceTokenExpiry(token string) (time.Time, error) {
	return parseJWTExpiry(token)
}

// 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
}