Weekend fun with further optimisations.

2026-10-19 04:47:48 +00:00 · 2025-09-13 23:42:02 +01:00
parent a31fd64112
commit 7457a901a9
48 changed files with 8726 additions and 3268 deletions
@@ -7,6 +7,7 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/base64"
+	"encoding/binary"
 	"encoding/json"
 	"encoding/pem"
 	"fmt"
@@ -17,101 +18,78 @@ import (
 	"time"
 )

-// JWK represents a JSON Web Key used for verifying
-// JWT signatures. Supports both RSA and ECDSA key types.
+// JWK represents a JSON Web Key as defined in RFC 7517.
+// It can represent different key types including RSA, EC, and symmetric keys.
 type JWK struct {
-	// Kty specifies the key type ("RSA" or "EC")
+	// Key type (e.g., "RSA", "EC", "oct")
 	Kty string `json:"kty"`
-	// Kid is the unique identifier for this key
-	Kid string `json:"kid"`
-	// Use indicates the intended use of the key ("sig" for signature)
-	Use string `json:"use"`
-	// N is the RSA public key modulus (base64url-encoded)
-	N string `json:"n"`
-	// E is the RSA public key exponent (base64url-encoded)
-	E string `json:"e"`
-	// Alg specifies the algorithm intended for use with this key
-	Alg string `json:"alg"`
-	// Crv specifies the elliptic curve for EC keys (P-256, P-384, P-521)
-	Crv string `json:"crv"`
-	// X is the x coordinate for EC public keys (base64url-encoded)
-	X string `json:"x"`
-	// Y is the y coordinate for EC public keys (base64url-encoded)
-	Y string `json:"y"`
+	// Key use (e.g., "sig" for signature, "enc" for encryption)
+	Use string `json:"use,omitempty"`
+	// Key operations allowed
+	KeyOps []string `json:"key_ops,omitempty"`
+	// Algorithm intended for use with this key
+	Alg string `json:"alg,omitempty"`
+	// Key ID
+	Kid string `json:"kid,omitempty"`
+
+	// RSA specific fields
+	N string `json:"n,omitempty"` // Modulus
+	E string `json:"e,omitempty"` // Exponent
+
+	// EC specific fields
+	Crv string `json:"crv,omitempty"` // Curve
+	X   string `json:"x,omitempty"`   // X coordinate
+	Y   string `json:"y,omitempty"`   // Y coordinate
 }

-// JWKSet represents a collection of JSON Web Keys from a JWKS endpoint.
-// OIDC providers publish multiple keys in a set to support key rotation.
+// JWKSet represents a set of JSON Web Keys.
+// Typically fetched from an OIDC provider's JWKS endpoint.
 type JWKSet struct {
 	// Keys contains the array of JWK objects
 	Keys []JWK `json:"keys"`
 }

-// JWKCache provides thread-safe caching of JWKS (JSON Web Key Sets)
-// with automatic expiration and the ability to fetch new keys and
-// refresh when keys expire.
+// JWKCache provides thread-safe caching of JWKS using UniversalCache
 type JWKCache struct {
-	// internalCache stores the cached JWKS data
-	internalCache CacheInterface
-	// CacheLifetime defines how long keys are cached before refresh
-	CacheLifetime time.Duration
-	// maxSize limits the number of cached JWKS entries
-	maxSize int
-	// mutex protects concurrent access to the cache
+	cache *UniversalCache
 	mutex sync.RWMutex
 }

 // JWKCacheInterface defines the contract for JWK caching implementations.
-// It provides methods for fetching JWKS, cache maintenance, and graceful shutdown.
 type JWKCacheInterface interface {
-	// GetJWKS fetches or retrieves cached JWKS from the given URL
 	GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error)
-	// Cleanup performs cache maintenance (removing expired entries)
 	Cleanup()
-	// Close shuts down the cache and releases resources
 	Close()
 }

-// NewJWKCache creates a new JWK cache with default configuration.
-// It initializes a cache with a 1-hour lifetime and maximum size of 100 entries.
+// NewJWKCache creates a new JWK cache using the global cache manager
 func NewJWKCache() *JWKCache {
-	config := DefaultUnifiedCacheConfig()
-	config.MaxSize = 100
-	unifiedCache := NewUnifiedCache(config)
-	cache := &JWKCache{
-		CacheLifetime: 1 * time.Hour,
-		maxSize:       100,
-		internalCache: NewCacheAdapter(unifiedCache),
+	manager := GetUniversalCacheManager(nil)
+	return &JWKCache{
+		cache: manager.GetJWKCache(),
 	}
-	return cache
 }

 // GetJWKS retrieves JWKS from cache or fetches from the remote URL if not cached.
-// It uses double-checked locking to prevent concurrent fetches of the same JWKS.
-// Parameters:
-//   - ctx: Context for request cancellation and deadlines
-//   - jwksURL: The JWKS endpoint URL to fetch from
-//   - httpClient: HTTP client to use for the request
-//
-// Returns:
-//   - *JWKSet: The retrieved or cached JWKS
-//   - error: Any error encountered during fetch or cache operations
 func (c *JWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
-	if c.internalCache != nil {
-		if cachedJwks, found := c.internalCache.Get(jwksURL); found {
-			return cachedJwks.(*JWKSet), nil
+	// Check cache first
+	if cachedValue, found := c.cache.Get(jwksURL); found {
+		if jwks, ok := cachedValue.(*JWKSet); ok {
+			return jwks, nil
 		}
 	}

 	c.mutex.Lock()
 	defer c.mutex.Unlock()

-	if c.internalCache != nil {
-		if cachedJwks, found := c.internalCache.Get(jwksURL); found {
-			return cachedJwks.(*JWKSet), nil
+	// Double-check after acquiring lock
+	if cachedValue, found := c.cache.Get(jwksURL); found {
+		if jwks, ok := cachedValue.(*JWKSet); ok {
+			return jwks, nil
 		}
 	}

+	// Fetch from URL
 	jwks, err := fetchJWKS(ctx, jwksURL, httpClient)
 	if err != nil {
 		return nil, err
@@ -121,162 +99,92 @@ func (c *JWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http
 		return nil, fmt.Errorf("JWKS response contains no keys")
 	}

-	lifetime := c.CacheLifetime
-	if lifetime == 0 {
-		lifetime = 1 * time.Hour
-	}
-
-	if c.internalCache != nil {
-		c.internalCache.Set(jwksURL, jwks, lifetime)
-	}
+	// Cache for 1 hour
+	c.cache.Set(jwksURL, jwks, 1*time.Hour)

 	return jwks, nil
 }

-// Cleanup removes expired entries from the JWK cache to free memory.
-// It delegates to the internal cache's cleanup method.
+// Cleanup is a no-op as cleanup is handled by UniversalCache
 func (c *JWKCache) Cleanup() {
-	if c != nil && c.internalCache != nil {
-		c.internalCache.Cleanup()
-	}
+	// Handled internally by UniversalCache
 }

-// Close shuts down the JWK cache and releases all resources.
-// It stops auto-cleanup routines and closes the internal cache.
+// Close is a no-op as the cache is managed globally
 func (c *JWKCache) Close() {
-	if c.internalCache != nil {
-		c.internalCache.Close()
-	}
+	// Managed by global cache manager
 }

-// SetMaxSize configures the maximum number of JWKS entries to cache.
-// This helps prevent unbounded memory growth in long-running applications.
-func (c *JWKCache) SetMaxSize(size int) {
-	c.maxSize = size
-	if c.internalCache != nil {
-		c.internalCache.SetMaxSize(size)
-	}
-}
-
-// fetchJWKS retrieves JWKS from a remote URL using the provided HTTP client.
-// It handles request creation, response validation, and JSON decoding.
-// Parameters:
-//   - ctx: Context for request cancellation and timeouts
-//   - jwksURL: The JWKS endpoint URL to fetch from
-//   - httpClient: HTTP client configured with appropriate timeouts
-//
-// Returns:
-//   - *JWKSet: The decoded JWKS response
-//   - An error if the request fails, the status code is not OK, or the response body cannot be decoded
+// fetchJWKS fetches JWKS from a remote URL
 func fetchJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
 	req, err := http.NewRequestWithContext(ctx, "GET", jwksURL, nil)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create JWKS request: %w", err)
+		return nil, fmt.Errorf("error creating JWKS request: %w", err)
 	}

 	resp, err := httpClient.Do(req)
 	if err != nil {
-		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
+		return nil, fmt.Errorf("error fetching JWKS: %w", err)
 	}
-	defer func() {
-		io.Copy(io.Discard, resp.Body)
-		resp.Body.Close()
-	}()
+	defer resp.Body.Close()

 	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("failed to fetch JWKS: unexpected status code %d", resp.StatusCode)
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("JWKS fetch failed with status %d: %s", resp.StatusCode, body)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("error reading JWKS response: %w", err)
 	}

 	var jwks JWKSet
-	if err := json.NewDecoder(resp.Body).Decode(&jwks); err != nil {
-		return nil, fmt.Errorf("failed to decode JWKS: %w", err)
+	if err := json.Unmarshal(body, &jwks); err != nil {
+		return nil, fmt.Errorf("error parsing JWKS: %w", err)
 	}

 	return &jwks, nil
 }

-// jwkToPEM converts a JSON Web Key to PEM-encoded public key format.
-// Supports RSA and ECDSA key types for JWT signature verification.
-// Parameters:
-//   - jwk: The JWK to convert
-//
-// Returns:
-//   - []byte: PEM-encoded public key
-//   - An error if the key type is unsupported or conversion fails
-func jwkToPEM(jwk *JWK) ([]byte, error) {
-	converter, ok := jwkConverters[jwk.Kty]
-	if !ok {
-		return nil, fmt.Errorf("unsupported key type: %s", jwk.Kty)
+// ToRSAPublicKey converts a JWK to an RSA public key.
+// Returns an error if the JWK is not an RSA key or if the key data is invalid.
+func (jwk *JWK) ToRSAPublicKey() (*rsa.PublicKey, error) {
+	if jwk.Kty != "RSA" {
+		return nil, fmt.Errorf("not an RSA key")
 	}
-	return converter(jwk)
-}

-// jwkToPEMConverter defines the function signature for JWK to PEM conversion functions.
-type jwkToPEMConverter func(*JWK) ([]byte, error)
-
-// jwkConverters maps JWK key types to their respective PEM conversion functions.
-var jwkConverters = map[string]jwkToPEMConverter{
-	"RSA": rsaJWKToPEM,
-	"EC":  ecJWKToPEM,
-}
-
-// rsaJWKToPEM converts an RSA JSON Web Key to PEM format.
-// It decodes the base64url-encoded modulus (n) and exponent (e) parameters
-// and constructs a standard RSA public key structure.
-// Parameters:
-//   - jwk: The RSA JWK to convert
-//
-// Returns:
-//   - []byte: PEM-encoded RSA public key
-//   - An error if decoding parameters fails or key marshaling fails
-func rsaJWKToPEM(jwk *JWK) ([]byte, error) {
 	nBytes, err := base64.RawURLEncoding.DecodeString(jwk.N)
 	if err != nil {
-		return nil, fmt.Errorf("failed to decode JWK 'n' parameter: %w", err)
+		return nil, fmt.Errorf("error decoding modulus: %w", err)
 	}
+
 	eBytes, err := base64.RawURLEncoding.DecodeString(jwk.E)
 	if err != nil {
-		return nil, fmt.Errorf("failed to decode JWK 'e' parameter: %w", err)
+		return nil, fmt.Errorf("error decoding exponent: %w", err)
 	}

-	n := new(big.Int).SetBytes(nBytes)
-	e := new(big.Int).SetBytes(eBytes)
-
-	pubKey := &rsa.PublicKey{
-		N: n,
-		E: int(e.Int64()),
+	// Convert exponent bytes to int
+	var e int
+	if len(eBytes) <= 8 {
+		// Pad to 8 bytes for uint64
+		paddedE := make([]byte, 8)
+		copy(paddedE[8-len(eBytes):], eBytes)
+		e = int(binary.BigEndian.Uint64(paddedE))
+	} else {
+		return nil, fmt.Errorf("exponent too large")
 	}

-	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
-	if err != nil {
-		return nil, fmt.Errorf("failed to marshal RSA public key: %w", err)
-	}
-
-	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
-		Type:  "PUBLIC KEY",
-		Bytes: pubKeyBytes,
-	})
-
-	return pubKeyPEM, nil
+	return &rsa.PublicKey{
+		N: new(big.Int).SetBytes(nBytes),
+		E: e,
+	}, nil
 }

-// ecJWKToPEM converts an Elliptic Curve JSON Web Key to PEM format.
-// It decodes the base64url-encoded x and y coordinates and supports
-// standard NIST curves (P-256, P-384, P-521).
-// Parameters:
-//   - jwk: The EC JWK to convert
-//
-// Returns:
-//   - []byte: PEM-encoded EC public key
-//   - An error if decoding parameters fails, the curve is unsupported, or key marshaling fails
-func ecJWKToPEM(jwk *JWK) ([]byte, error) {
-	xBytes, err := base64.RawURLEncoding.DecodeString(jwk.X)
-	if err != nil {
-		return nil, fmt.Errorf("failed to decode JWK 'x' parameter: %w", err)
-	}
-	yBytes, err := base64.RawURLEncoding.DecodeString(jwk.Y)
-	if err != nil {
-		return nil, fmt.Errorf("failed to decode JWK 'y' parameter: %w", err)
+// ToECDSAPublicKey converts a JWK to an ECDSA public key.
+// Returns an error if the JWK is not an EC key or if the key data is invalid.
+func (jwk *JWK) ToECDSAPublicKey() (*ecdsa.PublicKey, error) {
+	if jwk.Kty != "EC" {
+		return nil, fmt.Errorf("not an EC key")
 	}

 	var curve elliptic.Curve
@@ -288,24 +196,68 @@ func ecJWKToPEM(jwk *JWK) ([]byte, error) {
 	case "P-521":
 		curve = elliptic.P521()
 	default:
-		return nil, fmt.Errorf("unsupported elliptic curve: %s", jwk.Crv)
+		return nil, fmt.Errorf("unsupported curve: %s", jwk.Crv)
 	}

-	pubKey := &ecdsa.PublicKey{
+	xBytes, err := base64.RawURLEncoding.DecodeString(jwk.X)
+	if err != nil {
+		return nil, fmt.Errorf("error decoding X coordinate: %w", err)
+	}
+
+	yBytes, err := base64.RawURLEncoding.DecodeString(jwk.Y)
+	if err != nil {
+		return nil, fmt.Errorf("error decoding Y coordinate: %w", err)
+	}
+
+	return &ecdsa.PublicKey{
 		Curve: curve,
 		X:     new(big.Int).SetBytes(xBytes),
 		Y:     new(big.Int).SetBytes(yBytes),
+	}, nil
+}
+
+// GetKey finds a key by its ID (kid) in the JWKSet.
+// Returns nil if no key with the given ID is found.
+func (jwks *JWKSet) GetKey(kid string) *JWK {
+	for _, key := range jwks.Keys {
+		if key.Kid == kid {
+			return &key
+		}
+	}
+	return nil
+}
+
+// jwkToPEM converts a JWK to PEM format for signature verification
+func jwkToPEM(jwk *JWK) ([]byte, error) {
+	var publicKey interface{}
+	var err error
+
+	switch jwk.Kty {
+	case "RSA":
+		publicKey, err = jwk.ToRSAPublicKey()
+		if err != nil {
+			return nil, fmt.Errorf("failed to convert RSA JWK: %w", err)
+		}
+	case "EC":
+		publicKey, err = jwk.ToECDSAPublicKey()
+		if err != nil {
+			return nil, fmt.Errorf("failed to convert EC JWK: %w", err)
+		}
+	default:
+		return nil, fmt.Errorf("unsupported key type: %s", jwk.Kty)
 	}

-	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
+	// Marshal the public key to DER format
+	pubKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
 	if err != nil {
-		return nil, fmt.Errorf("failed to marshal EC public key: %w", err)
+		return nil, fmt.Errorf("failed to marshal public key: %w", err)
 	}

-	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
+	// Encode to PEM format
+	pemBlock := &pem.Block{
 		Type:  "PUBLIC KEY",
 		Bytes: pubKeyBytes,
-	})
+	}

-	return pubKeyPEM, nil
+	return pem.EncodeToMemory(pemBlock), nil
 }