fixup! Further pursue of perfection.

jwt.go

@@ -16,19 +16,28 @@ import (
 	"time"
 )
 
+// Replay attack protection cache and synchronization primitives.
+// This cache tracks JWT IDs (jti claims) to prevent token reuse attacks.
 var (
-	replayCacheMu        sync.RWMutex // Use RWMutex for better read performance
-	replayCache          *Cache       // Replace unbounded map with bounded Cache
-	replayCacheOnce      sync.Once
+	// replayCacheMu protects access to the replay cache instance
+	replayCacheMu sync.RWMutex
+	// replayCache stores JWT IDs with expiration to prevent replay attacks
+	replayCache *Cache
+	// replayCacheOnce ensures the replay cache is initialized only once
+	replayCacheOnce sync.Once
+	// replayCacheCleanupWG waits for cleanup goroutine to finish
 	replayCacheCleanupWG sync.WaitGroup
-	replayCacheCtx       context.Context
-	replayCacheCancel    context.CancelFunc
+	// replayCacheCtx provides context for cleanup operations
+	replayCacheCtx context.Context
+	// replayCacheCancel cancels the cleanup context
+	replayCacheCancel context.CancelFunc
+	// replayCacheCleanupMu protects cleanup operations
 	replayCacheCleanupMu sync.Mutex
 )
 
-// initReplayCache initializes the global replay cache for JWT ID tracking.
-// It uses sync.Once to ensure thread-safe single initialization.
+// initReplayCache initializes the JWT replay protection cache with bounded size.
 // The cache is bounded to 10,000 entries to prevent unbounded memory growth.
+// This function uses sync.Once to ensure thread-safe single initialization.
 func initReplayCache() {
 	replayCacheOnce.Do(func() {
 		replayCache = NewCache()
@@ -36,11 +45,10 @@ func initReplayCache() {
 	})
 }
 
-// cleanupReplayCache gracefully shuts down the replay cache.
-// It acquires a write lock, closes the cache, and sets it to nil
-// to ensure proper cleanup during shutdown.
+// cleanupReplayCache performs graceful shutdown of the replay cache system.
+// It cancels the cleanup context, waits for background goroutines to finish,
+// and properly closes the cache to ensure proper cleanup during shutdown.
 func cleanupReplayCache() {
-	// Cancel cleanup goroutine first
 	replayCacheCleanupMu.Lock()
 	if replayCacheCancel != nil {
 		replayCacheCancel()
@@ -49,7 +57,6 @@ func cleanupReplayCache() {
 	}
 	replayCacheCleanupMu.Unlock()
 
-	// Wait for cleanup goroutine to finish
 	replayCacheCleanupWG.Wait()
 
 	replayCacheMu.Lock()
@@ -58,18 +65,14 @@ func cleanupReplayCache() {
 	if replayCache != nil {
 		replayCache.Close()
 		replayCache = nil
-		// Reset the once to allow re-initialization
 		replayCacheOnce = sync.Once{}
 	}
 }
 
-// getReplayCacheStats returns current statistics about the replay cache.
-// Due to sync.Pool limitations, it returns 0 for current size and the
-// configured maximum size of 10,000.
-//
+// getReplayCacheStats returns statistics about the replay cache state.
 // Returns:
-//   - size: Current number of entries (always 0 due to implementation).
-//   - maxSize: Maximum allowed entries (10,000).
+//   - size: Current number of entries in the cache (currently always 0 due to interface limitations)
+//   - maxSize: Maximum allowed entries (10,000)
 func getReplayCacheStats() (size int, maxSize int) {
 	replayCacheMu.RLock()
 	defer replayCacheMu.RUnlock()
@@ -81,24 +84,20 @@ func getReplayCacheStats() (size int, maxSize int) {
 	return 0, 10000
 }
 
-// startReplayCacheCleanup initiates a background goroutine that periodically
-// cleans up expired entries from the replay cache. It runs every 5 minutes
-// and logs cache statistics if a logger is provided.
-//
+// startReplayCacheCleanup starts a background goroutine for periodic cache maintenance.
+// The goroutine runs every 5 minutes to clean expired entries and log cache statistics.
 // Parameters:
-//   - ctx: Context for cancellation.
-//   - logger: Logger for debug output (can be nil).
+//   - ctx: Parent context for cancellation
+//   - logger: Logger for debug output (can be nil)
 func startReplayCacheCleanup(ctx context.Context, logger *Logger) {
 	replayCacheCleanupMu.Lock()
 	defer replayCacheCleanupMu.Unlock()
 
-	// Cancel any existing cleanup goroutine
 	if replayCacheCancel != nil {
 		replayCacheCancel()
 		replayCacheCleanupWG.Wait()
 	}
 
-	// Create new context for this cleanup goroutine
 	replayCacheCtx, replayCacheCancel = context.WithCancel(ctx)
 
 	replayCacheCleanupWG.Add(1)
@@ -131,42 +130,47 @@ func startReplayCacheCleanup(ctx context.Context, logger *Logger) {
 	}()
 }
 
+// ClockSkewToleranceFuture defines the maximum allowable clock skew for future time validation.
+// Tokens are considered valid for an additional 2 minutes past their expiration time.
 var ClockSkewToleranceFuture = 2 * time.Minute
 
+// ClockSkewTolerancePast defines the maximum allowable clock skew for past time validation.
+// Tokens are considered valid if issued up to 10 seconds in the future.
 var ClockSkewTolerancePast = 10 * time.Second
 
+// ClockSkewTolerance is an alias for ClockSkewToleranceFuture for backward compatibility.
 var ClockSkewTolerance = ClockSkewToleranceFuture
 
-// JWT represents a JSON Web Token as defined in RFC 7519.
-// JWT represents a parsed JSON Web Token with its three components.
-// It provides structured access to the header, claims, and signature
+// JWT represents a parsed JSON Web Token with its constituent parts.
+// It provides a structured representation of JWT components
 // for validation and processing within the OIDC middleware.
 type JWT struct {
-	Header    map[string]interface{}
-	Claims    map[string]interface{}
-	Token     string
+	// Header contains the JWT header claims (alg, typ, kid, etc.)
+	Header map[string]interface{}
+	// Claims contains the JWT payload claims (iss, sub, aud, exp, etc.)
+	Claims map[string]interface{}
+	// Token is the original JWT token string
+	Token string
+	// Signature contains the decoded JWT signature bytes
 	Signature []byte
 }
 
-// parseJWT decodes a raw JWT string into its constituent parts: header, claims, and signature.
-// It splits the token string by '.', decodes each part using base64 URL decoding,
-// and unmarshals the header and claims JSON into maps. The raw signature bytes are stored.
-// It performs basic format validation (expecting 3 parts).
-// Note: This function does *not* validate the signature or the claims.
-//
+// parseJWT parses a JWT token string into its constituent parts.
+// It decodes the base64url-encoded header, claims, and signature components
+// and unmarshals the JSON data into structured maps. Uses memory pools
+// for efficient memory allocation during parsing.
 // Parameters:
-//   - tokenString: The raw JWT string.
+//   - tokenString: The JWT token string to parse
 //
 // Returns:
-//   - A pointer to a JWT struct containing the decoded parts.
-//   - An error if the token format is invalid or decoding/unmarshaling fails.
+//   - *JWT: Parsed JWT structure with header, claims, and signature
+//   - An error if the token format is invalid or decoding/unmarshaling fails
 func parseJWT(tokenString string) (*JWT, error) {
 	parts := strings.Split(tokenString, ".")
 	if len(parts) != 3 {
 		return nil, fmt.Errorf("invalid JWT format: expected 3 parts, got %d", len(parts))
 	}
 
-	// Use memory pool for efficient buffer management
 	pools := GetGlobalMemoryPools()
 	jwtBuf := pools.GetJWTParsingBuffer()
 	defer pools.PutJWTParsingBuffer(jwtBuf)
@@ -175,7 +179,6 @@ func parseJWT(tokenString string) (*JWT, error) {
 		Token: tokenString,
 	}
 
-	// Decode header using pooled buffer
 	headerLen := base64.RawURLEncoding.DecodedLen(len(parts[0]))
 	if headerLen > cap(jwtBuf.HeaderBuf) {
 		jwtBuf.HeaderBuf = make([]byte, headerLen)
@@ -197,7 +200,6 @@ func parseJWT(tokenString string) (*JWT, error) {
 		return nil, fmt.Errorf("invalid JWT format: header is nil after unmarshaling")
 	}
 
-	// Decode claims using pooled buffer
 	claimsLen := base64.RawURLEncoding.DecodedLen(len(parts[1]))
 	if claimsLen > cap(jwtBuf.PayloadBuf) {
 		jwtBuf.PayloadBuf = make([]byte, claimsLen)
@@ -219,7 +221,6 @@ func parseJWT(tokenString string) (*JWT, error) {
 		return nil, fmt.Errorf("invalid JWT format: claims is nil after unmarshaling")
 	}
 
-	// Decode signature using pooled buffer
 	sigLen := base64.RawURLEncoding.DecodedLen(len(parts[2]))
 	if sigLen > cap(jwtBuf.SignatureBuf) {
 		jwtBuf.SignatureBuf = make([]byte, sigLen)
@@ -232,34 +233,23 @@ func parseJWT(tokenString string) (*JWT, error) {
 		return nil, fmt.Errorf("invalid JWT format: failed to decode signature: %v", err)
 	}
 
-	// Copy signature to JWT struct (create new slice to avoid pool retention)
 	jwt.Signature = make([]byte, n)
 	copy(jwt.Signature, jwtBuf.SignatureBuf[:n])
 
 	return jwt, nil
 }
 
-// Verify performs standard claim validation on the JWT according to RFC 7519.
-// It checks the following:
-// - Algorithm ('alg') is supported.
-// - Issuer ('iss') matches the expected issuerURL.
-// - Audience ('aud') contains the expected clientID.
-// - Expiration time ('exp') is in the future (within tolerance).
-// - Issued at time ('iat') is in the past (within tolerance).
-// - Not before time ('nbf'), if present, is in the past (within tolerance).
-// - Subject ('sub') claim exists and is not empty.
-// - JWT ID ('jti'), if present, is checked against a replay cache to prevent token reuse.
-//
+// Verify performs comprehensive JWT token validation according to OIDC specifications.
+// It validates the token signature algorithm, issuer, audience, expiration, issued-at time,
+// not-before time (if present), and prevents replay attacks using JTI claims.
 // Parameters:
-//   - issuerURL: The expected issuer URL (e.g., "https://accounts.google.com").
-//   - clientID: The expected audience value (the client ID of this application).
-//   - skipReplayCheck: If true, skips JTI replay detection (used for revalidation of cached tokens).
+//   - issuerURL: Expected issuer URL to validate against
+//   - clientID: Expected audience (client ID) to validate against
+//   - skipReplayCheck: Optional parameter to skip replay attack protection
 //
 // Returns:
-//   - nil if all standard claims are valid.
-//   - An error describing the first validation failure encountered.
+//   - An error describing the first validation failure encountered
 func (j *JWT) Verify(issuerURL, clientID string, skipReplayCheck ...bool) error {
-	// Validate algorithm to prevent algorithm switching attacks
 	alg, ok := j.Header["alg"].(string)
 	if !ok {
 		return fmt.Errorf("missing 'alg' header")
@@ -354,16 +344,14 @@ func (j *JWT) Verify(issuerURL, clientID string, skipReplayCheck ...bool) error
 	return nil
 }
 
-// verifyAudience checks if the expected audience is present in the token's 'aud' claim.
-// The 'aud' claim can be a single string or an array of strings.
-//
+// verifyAudience validates the JWT audience claim against the expected client ID.
+// The audience claim can be either a single string or an array of strings.
 // Parameters:
-//   - tokenAudience: The 'aud' claim value extracted from the token (can be string or []interface{}).
-//   - expectedAudience: The audience value expected for this application (client ID).
+//   - tokenAudience: The audience claim from the JWT (string or []interface{})
+//   - expectedAudience: The expected audience value (typically the OAuth client ID)
 //
 // Returns:
-//   - nil if the expected audience is found.
-//   - An error if the claim type is invalid or the expected audience is not present.
+//   - An error if the claim type is invalid or the expected audience is not present
 func verifyAudience(tokenAudience interface{}, expectedAudience string) error {
 	switch aud := tokenAudience.(type) {
 	case string:
@@ -387,15 +375,13 @@ func verifyAudience(tokenAudience interface{}, expectedAudience string) error {
 	return nil
 }
 
-// verifyIssuer checks if the token's 'iss' claim matches the expected issuer URL.
-//
+// verifyIssuer validates the JWT issuer claim against the expected issuer URL.
 // Parameters:
-//   - tokenIssuer: The 'iss' claim value from the token.
-//   - expectedIssuer: The expected issuer URL configured for the OIDC provider.
+//   - tokenIssuer: The issuer claim from the JWT
+//   - expectedIssuer: The expected issuer URL from OIDC configuration
 //
 // Returns:
-//   - nil if the issuers match.
-//   - An error if the issuers do not match.
+//   - An error if the issuers do not match
 func verifyIssuer(tokenIssuer, expectedIssuer string) error {
 	if tokenIssuer != expectedIssuer {
 		return fmt.Errorf("invalid issuer (token: %s, expected: %s)", tokenIssuer, expectedIssuer)
@@ -403,17 +389,15 @@ func verifyIssuer(tokenIssuer, expectedIssuer string) error {
 	return nil
 }
 
-// verifyTimeConstraint checks time-based claims ('exp', 'iat', 'nbf') against the current time,
-// allowing for configurable clock skew. It uses different tolerances for past and future checks.
-//
+// verifyTimeConstraint validates time-based JWT claims with clock skew tolerance.
+// It handles both future constraints (exp) and past constraints (iat, nbf).
 // Parameters:
-//   - unixTime: The timestamp value from the claim (as a float64 Unix time).
-//   - claimName: The name of the claim being verified ("exp", "iat", "nbf").
-//   - future: A boolean indicating the direction of the check (true for 'exp', false for 'iat'/'nbf').
+//   - unixTime: The Unix timestamp from the JWT claim
+//   - claimName: Name of the claim being validated (for error messages)
+//   - future: If true, validates against future tolerance; if false, against past tolerance
 //
 // Returns:
-//   - nil if the time constraint is met within the allowed tolerance.
-//   - An error describing the failure (e.g., "token has expired", "token used before issued").
+//   - An error describing the failure (e.g., "token has expired", "token used before issued")
 func verifyTimeConstraint(unixTime float64, claimName string, future bool) error {
 	claimTime := time.Unix(int64(unixTime), 0)
 	now := time.Now()
@@ -438,39 +422,34 @@ func verifyTimeConstraint(unixTime float64, claimName string, future bool) error
 	return err
 }
 
-// verifyExpiration checks the 'exp' (Expiration Time) claim.
+// verifyExpiration validates the JWT expiration time (exp claim) with clock skew tolerance.
 // It calls verifyTimeConstraint with future=true.
 func verifyExpiration(expiration float64) error {
 	return verifyTimeConstraint(expiration, "exp", true)
 }
 
-// verifyIssuedAt checks the 'iat' (Issued At) claim.
+// verifyIssuedAt validates the JWT issued-at time (iat claim) with clock skew tolerance.
 // It calls verifyTimeConstraint with future=false.
 func verifyIssuedAt(issuedAt float64) error {
 	return verifyTimeConstraint(issuedAt, "iat", false)
 }
 
-// verifyNotBefore checks the 'nbf' (Not Before) claim.
+// verifyNotBefore validates the JWT not-before time (nbf claim) with clock skew tolerance.
 // It calls verifyTimeConstraint with future=false.
 func verifyNotBefore(notBefore float64) error {
 	return verifyTimeConstraint(notBefore, "nbf", false)
 }
 
-// verifySignature validates the JWT's signature using the provided public key.
-// It parses the public key from PEM format, selects the appropriate hashing algorithm
-// based on the 'alg' parameter (SHA256/384/512), hashes the token's signing input
-// (header + "." + payload), and then verifies the signature against the hash using
-// the corresponding RSA (PKCS1v15 or PSS) or ECDSA verification method.
-//
+// verifySignature verifies the JWT signature using the provided public key.
+// Supports RSA (RS256/384/512, PS256/384/512) and ECDSA (ES256/384/512) algorithms.
 // Parameters:
-//   - tokenString: The raw, complete JWT string.
-//   - publicKeyPEM: The public key corresponding to the private key used for signing, in PEM format.
-//   - alg: The algorithm specified in the JWT header (e.g., "RS256", "ES384").
+//   - tokenString: The complete JWT token string
+//   - publicKeyPEM: The public key in PEM format
+//   - alg: The signing algorithm specified in the JWT header
 //
 // Returns:
-//   - nil if the signature is valid.
-//   - An error if the token format is invalid, decoding fails, key parsing fails,
-//     the algorithm is unsupported, or the signature verification fails.
+//   - An error if the key parsing fails, the algorithm is unsupported,
+//     or the signature verification fails
 func verifySignature(tokenString string, publicKeyPEM []byte, alg string) error {
 	parts := strings.Split(tokenString, ".")
 	if len(parts) != 3 {