fixup! Further pursue of perfection.

jwk.go

@@ -17,106 +17,107 @@ import (
 	"time"
 )
 
-// JWK represents a JSON Web Key as defined in RFC 7517.
-// It contains the cryptographic key parameters used for verifying
+// JWK represents a JSON Web Key used for verifying
 // JWT signatures. Supports both RSA and ECDSA key types.
 type JWK struct {
-	Kty string `json:"kty"` // Key type (RSA, EC)
-	Kid string `json:"kid"` // Key ID
-	Use string `json:"use"` // Key use (sig, enc)
-	N   string `json:"n"`   // RSA modulus
-	E   string `json:"e"`   // RSA public exponent
-	Alg string `json:"alg"` // Algorithm
-	Crv string `json:"crv"` // ECDSA curve
-	X   string `json:"x"`   // ECDSA x coordinate
-	Y   string `json:"y"`   // ECDSA y coordinate
+	// Kty specifies the key type ("RSA" or "EC")
+	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"`
 }
 
-// JWKSet represents a set of JSON Web Keys as returned by
-// an OIDC provider's JWKS endpoint. It contains multiple keys
-// to support key rotation.
+// JWKSet represents a collection of JSON Web Keys from a JWKS endpoint.
+// OIDC providers publish multiple keys in a set to support key rotation.
 type JWKSet struct {
+	// Keys contains the array of JWK objects
 	Keys []JWK `json:"keys"`
 }
 
-// JWKCache provides thread-safe caching of JSON Web Key Sets.
-// It fetches JWKS from OIDC providers and caches them to reduce
-// network requests. The cache supports expiration and automatic
+// 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.
 type JWKCache struct {
+	// internalCache stores the cached JWKS data
 	internalCache *Cache
+	// CacheLifetime defines how long keys are cached before refresh
 	CacheLifetime time.Duration
-	maxSize       int
-	mutex         sync.RWMutex
+	// maxSize limits the number of cached JWKS entries
+	maxSize int
+	// mutex protects concurrent access to the cache
+	mutex sync.RWMutex
 }
 
-// JWKCacheInterface defines the contract for JWK cache implementations.
-// It provides methods for retrieving JWKS, performing cleanup, and
-// graceful shutdown.
+// 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()
 }
 
-// GetJWKS retrieves the JSON Web Key Set (JWKS) from the cache or fetches it from the provider.
-// It first checks if a valid, non-expired JWKS is present in the cache. If so, it returns the cached version.
-// Otherwise, it attempts to fetch the JWKS from the specified jwksURL using the provided httpClient.
-// If the fetch is successful, the JWKS is stored in the cache with an expiration time based on CacheLifetime
-// (defaulting to 1 hour if not set) and returned.
-// This method uses double-checked locking to minimize contention when the cache needs refreshing.
-//
-// Parameters:
-//   - ctx: Context for the HTTP request if fetching is required.
-//   - jwksURL: The URL of the OIDC provider's JWKS endpoint.
-//   - httpClient: The HTTP client to use for fetching the JWKS.
-//
-// Returns:
-//   - A pointer to the JWKSet containing the keys.
-//   - An error if fetching fails or the response cannot be decoded.
-
 // NewJWKCache creates a new JWK cache with default configuration.
 // It initializes a cache with a 1-hour lifetime and maximum size of 100 entries.
 func NewJWKCache() *JWKCache {
 	cache := &JWKCache{
 		CacheLifetime: 1 * time.Hour,
-		maxSize:       100, // Default maximum size
+		maxSize:       100,
 		internalCache: NewCache(),
 	}
 	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) {
-	// Use only the internalCache for storage to avoid double storage
 	if c.internalCache != nil {
 		if cachedJwks, found := c.internalCache.Get(jwksURL); found {
 			return cachedJwks.(*JWKSet), nil
 		}
 	}
 
-	// Acquire write lock for potential update
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
-	// Double-check after acquiring write lock
 	if c.internalCache != nil {
 		if cachedJwks, found := c.internalCache.Get(jwksURL); found {
 			return cachedJwks.(*JWKSet), nil
 		}
 	}
 
-	// Fetch new JWKS
 	jwks, err := fetchJWKS(ctx, jwksURL, httpClient)
 	if err != nil {
 		return nil, err
 	}
 
-	// STABILITY FIX: Validate JWKS contains keys before caching
 	if len(jwks.Keys) == 0 {
 		return nil, fmt.Errorf("JWKS response contains no keys")
 	}
 
-	// Store in the internalCache only (avoid double storage)
 	lifetime := c.CacheLifetime
 	if lifetime == 0 {
 		lifetime = 1 * time.Hour
@@ -129,9 +130,7 @@ func (c *JWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http
 	return jwks, nil
 }
 
-// Cleanup removes the cached JWKS if it has expired.
-// This is intended to be called periodically to ensure stale JWKS data is cleared.
-// Cleanup removes expired entries from the cache.
+// Cleanup removes expired entries from the JWK cache to free memory.
 // It delegates to the internal cache's cleanup method.
 func (c *JWKCache) Cleanup() {
 	if c.internalCache != nil {
@@ -139,15 +138,16 @@ func (c *JWKCache) Cleanup() {
 	}
 }
 
-// Close shuts down the cache's auto-cleanup routine.
+// Close shuts down the JWK cache and releases all resources.
+// It stops auto-cleanup routines and closes the internal cache.
 func (c *JWKCache) Close() {
-	// Delegate to internal cache's Close method
 	if c.internalCache != nil {
 		c.internalCache.Close()
 	}
 }
 
-// SetMaxSize sets the maximum number of items in the cache
+// 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 {
@@ -155,19 +155,17 @@ func (c *JWKCache) SetMaxSize(size int) {
 	}
 }
 
-// fetchJWKS retrieves the JSON Web Key Set (JWKS) from the specified URL.
-// It uses the provided context and HTTP client to make the request.
-//
+// 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 the HTTP request.
-//   - jwksURL: The URL of the OIDC provider's JWKS endpoint.
-//   - httpClient: The HTTP client to use for the request.
+//   - ctx: Context for request cancellation and timeouts
+//   - jwksURL: The JWKS endpoint URL to fetch from
+//   - httpClient: HTTP client configured with appropriate timeouts
 //
 // Returns:
-//   - A pointer to the fetched JWKSet.
-//   - An error if the request fails, the status code is not OK, or the response body cannot be decoded.
+//   - *JWKSet: The decoded JWKS response
+//   - An error if the request fails, the status code is not OK, or the response body cannot be decoded
 func fetchJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
-	// Create a request with context to enforce timeout
 	req, err := http.NewRequestWithContext(ctx, "GET", jwksURL, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create JWKS request: %w", err)
@@ -178,7 +176,6 @@ func fetchJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*J
 		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
 	}
 	defer func() {
-		// Always drain the body before closing to ensure connection can be reused
 		io.Copy(io.Discard, resp.Body)
 		resp.Body.Close()
 	}()
@@ -195,16 +192,14 @@ func fetchJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*J
 	return &jwks, nil
 }
 
-// jwkToPEM converts a JWK (JSON Web Key) object into PEM (Privacy-Enhanced Mail) format.
-// It selects the appropriate conversion function based on the JWK's key type ("kty").
-// Currently supports "RSA" and "EC" key types.
-//
+// jwkToPEM converts a JSON Web Key to PEM-encoded public key format.
+// Supports RSA and ECDSA key types for JWT signature verification.
 // Parameters:
-//   - jwk: A pointer to the JWK object to convert.
+//   - jwk: The JWK to convert
 //
 // Returns:
-//   - A byte slice containing the public key in PEM format.
-//   - An error if the key type is unsupported or conversion fails.
+//   - []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 {
@@ -213,24 +208,24 @@ func jwkToPEM(jwk *JWK) ([]byte, error) {
 	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 JWK into PEM format.
-// It decodes the modulus (n) and exponent (e) from base64 URL encoding,
-// constructs an rsa.PublicKey, marshals it into PKIX format, and then
-// encodes it as a PEM block.
-//
+// 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: A pointer to the RSA JWK object (must have "kty": "RSA").
+//   - jwk: The RSA JWK to convert
 //
 // Returns:
-//   - A byte slice containing the RSA public key in PEM format.
-//   - An error if decoding parameters fails or key marshaling fails.
+//   - []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 {
@@ -262,18 +257,15 @@ func rsaJWKToPEM(jwk *JWK) ([]byte, error) {
 	return pubKeyPEM, nil
 }
 
-// ecJWKToPEM converts an EC (Elliptic Curve) JWK into PEM format.
-// It decodes the X and Y coordinates from base64 URL encoding, determines the
-// elliptic curve based on the "crv" parameter (P-256, P-384, P-521),
-// constructs an ecdsa.PublicKey, marshals it into PKIX format, and then
-// encodes it as a PEM block.
-//
+// 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: A pointer to the EC JWK object (must have "kty": "EC").
+//   - jwk: The EC JWK to convert
 //
 // Returns:
-//   - A byte slice containing the EC public key in PEM format.
-//   - An error if decoding parameters fails, the curve is unsupported, or key marshaling fails.
+//   - []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 {