homelab-frontend/internal/auth/jwt.go

package auth

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/MicahParks/keyfunc/v2"
	"github.com/golang-jwt/jwt/v5"
)

// Validator validates JWTs against Authentik JWKS.
type Validator struct {
	issuer   string
	audience string
	jwksURL  string
	jwks     *keyfunc.JWKS
	mu       sync.Mutex
}

// NewValidator creates a new JWT validator for a service.
// JWKS fetching is lazy (deferred until first validation).
func NewValidator(issuer, audience, jwksURL string) *Validator {
	return &Validator{
		issuer:   issuer,
		audience: audience,
		jwksURL:  jwksURL,
		jwks:     nil, // Lazy-loaded on first use
	}
}

// ensureJWKS fetches JWKS on first use (lazy initialization, thread-safe).
func (v *Validator) ensureJWKS() error {
	v.mu.Lock()
	defer v.mu.Unlock()

	if v.jwks != nil {
		return nil
	}

	options := keyfunc.Options{
		Ctx:              context.Background(),
		RefreshInterval:  15 * time.Minute,
		RefreshRateLimit: 5 * time.Minute,
		RefreshTimeout:   10 * time.Second,
		RefreshErrorHandler: func(err error) {
			fmt.Printf("JWKS refresh error for %s: %v\n", v.issuer, err)
		},
	}

	jwks, err := keyfunc.Get(v.jwksURL, options)
	if err != nil {
		return fmt.Errorf("failed to fetch JWKS from %s: %v", v.jwksURL, err)
	}

	v.jwks = jwks
	return nil
}

// ValidateBearerToken extracts and validates the Bearer token from Authorization header.
// Returns claims on success, error message on failure.
func (v *Validator) ValidateBearerToken(authHeader string) (jwt.MapClaims, error) {
	if authHeader == "" {
		return nil, fmt.Errorf("missing Authorization header")
	}

	// Extract token from "Bearer <token>"
	tokenString := ""
	if len(authHeader) > 7 && authHeader[:7] == "Bearer " {
		tokenString = authHeader[7:]
	} else {
		return nil, fmt.Errorf("invalid Authorization header format")
	}

	// Ensure JWKS is loaded (lazy)
	if err := v.ensureJWKS(); err != nil {
		return nil, err
	}

	// Parse and validate
	claims := jwt.MapClaims{}
	token, err := jwt.ParseWithClaims(tokenString, claims, v.jwks.Keyfunc)
	if err != nil {
		return nil, fmt.Errorf("token validation failed: %v", err)
	}

	if !token.Valid {
		return nil, fmt.Errorf("token is invalid")
	}

	// Verify required claims
	now := time.Now()
	const skew = 60 * time.Second

	// Check exp
	if exp, ok := claims["exp"].(float64); ok {
		if time.Now().After(time.Unix(int64(exp), 0).Add(skew)) {
			return nil, fmt.Errorf("token expired")
		}
	}

	// Check nbf (not before)
	if nbf, ok := claims["nbf"].(float64); ok {
		if now.Before(time.Unix(int64(nbf), 0).Add(-skew)) {
			return nil, fmt.Errorf("token not yet valid")
		}
	}

	// Check iss (issuer)
	if iss, ok := claims["iss"].(string); !ok || iss != v.issuer {
		return nil, fmt.Errorf("invalid issuer: expected %s, got %s", v.issuer, iss)
	}

	// Check aud (audience)
	if aud, ok := claims["aud"].(string); !ok || aud != v.audience {
		return nil, fmt.Errorf("invalid audience: expected %s, got %s", v.audience, aud)
	}

	return claims, nil
}

// CheckPermissions checks if claims contain required permission(s).
// Returns true if any required permission is found or wildcard "*" exists.
func (v *Validator) CheckPermissions(claims jwt.MapClaims, required ...string) bool {
	permsIface, ok := claims["permissions"]
	if !ok {
		return false
	}

	perms, ok := permsIface.([]interface{})
	if !ok {
		return false
	}

	for _, perm := range perms {
		permStr, ok := perm.(string)
		if !ok {
			continue
		}
		if permStr == "*" {
			return true
		}
		for _, req := range required {
			if permStr == req {
				return true
			}
		}
	}

	return false
}

// DecodeToken decodes JWT payload without verification (for debugging/testing).
func DecodeToken(tokenString string) (jwt.MapClaims, error) {
	claims := jwt.MapClaims{}
	_, _, err := new(jwt.Parser).ParseUnverified(tokenString, claims)
	if err != nil {
		return nil, err
	}
	return claims, nil
}