feat-dev-e2e3/pkg/auth/jwt.go

package auth

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

var (
	// ErrInvalidToken is returned when the token is malformed or signature is invalid.
	ErrInvalidToken = errors.New("invalid token")
	// ErrExpiredToken is returned when the token has expired.
	ErrExpiredToken = errors.New("token expired")
	// ErrInvalidClaims is returned when the token claims are invalid.
	ErrInvalidClaims = errors.New("invalid claims")
)

// JWTConfig configures the JWT validator.
type JWTConfig struct {
	// Secret is the HMAC secret key for HS256/HS384/HS512
	Secret []byte
	// PublicKey is the RSA/ECDSA public key for RS*/ES* algorithms (optional)
	PublicKey any
	// Issuer is the expected issuer claim (optional)
	Issuer string
	// Audience is the expected audience claim (optional)
	Audience string
}

// JWTClaims extends jwt.RegisteredClaims with custom fields.
type JWTClaims struct {
	jwt.RegisteredClaims
	// UserID is the user identifier
	UserID string `json:"uid,omitempty"`
	// Email is the user's email
	Email string `json:"email,omitempty"`
	// Roles are the user's roles
	Roles []string `json:"roles,omitempty"`
	// Scopes are the permitted scopes
	Scopes []string `json:"scopes,omitempty"`
}

// JWTValidator validates JWT tokens.
type JWTValidator struct {
	secret    []byte
	publicKey any
	issuer    string
	audience  string
}

// NewJWTValidator creates a new JWT validator.
func NewJWTValidator(cfg JWTConfig) *JWTValidator {
	return &JWTValidator{
		secret:    cfg.Secret,
		publicKey: cfg.PublicKey,
		issuer:    cfg.Issuer,
		audience:  cfg.Audience,
	}
}

// Validate validates a JWT token and returns the user.
func (v *JWTValidator) Validate(ctx context.Context, tokenString string) (*User, error) {
	// Parse and validate the token
	token, err := jwt.ParseWithClaims(tokenString, &JWTClaims{}, func(token *jwt.Token) (any, error) {
		// Check signing method
		switch token.Method.(type) {
		case *jwt.SigningMethodHMAC:
			if v.secret == nil {
				return nil, fmt.Errorf("HMAC secret not configured")
			}
			return v.secret, nil
		case *jwt.SigningMethodRSA, *jwt.SigningMethodECDSA:
			if v.publicKey == nil {
				return nil, fmt.Errorf("public key not configured")
			}
			return v.publicKey, nil
		default:
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
	})

	if err != nil {
		if errors.Is(err, jwt.ErrTokenExpired) {
			return nil, ErrExpiredToken
		}
		return nil, fmt.Errorf("%w: %v", ErrInvalidToken, err)
	}

	claims, ok := token.Claims.(*JWTClaims)
	if !ok || !token.Valid {
		return nil, ErrInvalidClaims
	}

	// Validate issuer if configured
	if v.issuer != "" && claims.Issuer != v.issuer {
		return nil, fmt.Errorf("%w: invalid issuer", ErrInvalidClaims)
	}

	// Validate audience if configured
	if v.audience != "" {
		found := false
		for _, aud := range claims.Audience {
			if aud == v.audience {
				found = true
				break
			}
		}
		if !found {
			return nil, fmt.Errorf("%w: invalid audience", ErrInvalidClaims)
		}
	}

	// Build user from claims
	user := &User{
		ID:     claims.UserID,
		Email:  claims.Email,
		Roles:  claims.Roles,
		Scopes: claims.Scopes,
	}

	// Fallback to subject if no user ID
	if user.ID == "" {
		user.ID = claims.Subject
	}

	return user, nil
}

// -----------------------------------------------------------------------------
// Token Generation (for testing and admin tools)
// -----------------------------------------------------------------------------

// GenerateToken creates a new JWT token for the given user.
// expiresIn specifies the token lifetime.
func GenerateToken(secret []byte, user *User, expiresIn time.Duration) (string, error) {
	now := time.Now()

	claims := JWTClaims{
		RegisteredClaims: jwt.RegisteredClaims{
			Subject:   user.ID,
			ExpiresAt: jwt.NewNumericDate(now.Add(expiresIn)),
			IssuedAt:  jwt.NewNumericDate(now),
			NotBefore: jwt.NewNumericDate(now),
		},
		UserID: user.ID,
		Email:  user.Email,
		Roles:  user.Roles,
		Scopes: user.Scopes,
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(secret)
}

// GenerateTokenWithIssuer creates a new JWT token with issuer and audience claims.
func GenerateTokenWithIssuer(secret []byte, user *User, expiresIn time.Duration, issuer, audience string) (string, error) {
	now := time.Now()

	claims := JWTClaims{
		RegisteredClaims: jwt.RegisteredClaims{
			Subject:   user.ID,
			Issuer:    issuer,
			Audience:  []string{audience},
			ExpiresAt: jwt.NewNumericDate(now.Add(expiresIn)),
			IssuedAt:  jwt.NewNumericDate(now),
			NotBefore: jwt.NewNumericDate(now),
		},
		UserID: user.ID,
		Email:  user.Email,
		Roles:  user.Roles,
		Scopes: user.Scopes,
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(secret)
}