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

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

// Global cache for service tokens (DID:HoldDID -> token)
// Service tokens are JWTs issued by a user's PDS to authorize AppView to act on their behalf
// when communicating with hold services. These tokens are scoped to specific holds and have
// limited lifetime (typically 60s, can request up to 5min).
var (
	globalServiceTokens   = make(map[string]*serviceTokenEntry)
	globalServiceTokensMu sync.RWMutex
)

// GetServiceToken retrieves a cached service token for the given DID and hold DID
// Returns empty string if no valid cached token exists
func GetServiceToken(did, holdDID string) (token string, expiresAt time.Time) {
	cacheKey := did + ":" + holdDID

	globalServiceTokensMu.RLock()
	entry, exists := globalServiceTokens[cacheKey]
	globalServiceTokensMu.RUnlock()

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

	// Check if token is still valid
	if time.Now().After(entry.expiresAt) {
		// Token expired, remove from cache
		globalServiceTokensMu.Lock()
		delete(globalServiceTokens, cacheKey)
		globalServiceTokensMu.Unlock()
		return "", time.Time{}
	}

	return entry.token, entry.expiresAt
}

// SetServiceToken stores a service token in the cache
// Automatically parses the JWT to extract the expiry time
// Applies a 10-second safety margin (cache expires 10s before actual JWT expiry)
func SetServiceToken(did, holdDID, token string) error {
	cacheKey := did + ":" + holdDID

	// Parse JWT to extract expiry (don't verify signature - we trust the PDS)
	expiry, err := parseJWTExpiry(token)
	if err != nil {
		// If parsing fails, use default 50s TTL (conservative fallback)
		slog.Warn("Failed to parse JWT expiry, using default 50s", "error", err, "cacheKey", cacheKey)
		expiry = time.Now().Add(50 * time.Second)
	} else {
		// Apply 10s safety margin to avoid using nearly-expired tokens
		expiry = expiry.Add(-10 * time.Second)
	}

	globalServiceTokensMu.Lock()
	globalServiceTokens[cacheKey] = &serviceTokenEntry{
		token:     token,
		expiresAt: expiry,
	}
	globalServiceTokensMu.Unlock()

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

	return nil
}

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

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

	// Parse the JSON payload
	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
}

// InvalidateServiceToken removes a service token from the cache
// Used when we detect that a token is invalid or the user's session has expired
func InvalidateServiceToken(did, holdDID string) {
	cacheKey := did + ":" + holdDID

	globalServiceTokensMu.Lock()
	delete(globalServiceTokens, cacheKey)
	globalServiceTokensMu.Unlock()

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

// GetCacheStats returns statistics about the service token cache for debugging
func GetCacheStats() map[string]any {
	globalServiceTokensMu.RLock()
	defer globalServiceTokensMu.RUnlock()

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

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

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

// CleanExpiredTokens removes expired tokens from the cache
// Can be called periodically to prevent unbounded growth (though expired tokens
// are also removed lazily on access)
func CleanExpiredTokens() {
	globalServiceTokensMu.Lock()
	defer globalServiceTokensMu.Unlock()

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

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

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