homelab-frontend/internal/auth/jwt.go

package auth

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

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

// isValidIssuer checks if issuer is from Authentik (any provider/app).
// Accepts: https://authentik.riotpiao.com/application/o/{provider}/
func isValidIssuer(iss string) bool {
	return strings.Contains(iss, "authentik.riotpiao.com/application/o/") &&
		strings.HasSuffix(iss, "/")
}

// 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) - accept any Authentik provider issuer
	// (portfolio-agent, memory-agent, api-gw, etc.)
	// All use same signing key so JWKS validation is sufficient
	if iss, ok := claims["iss"].(string); !ok {
		return nil, fmt.Errorf("missing issuer claim")
	} else if !isValidIssuer(iss) {
		return nil, fmt.Errorf("invalid issuer: %s", 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).
// Checks both "permissions" claim (for users) and "roles" claim (for service accounts).
// Returns true if any required permission is found or wildcard "*" exists.
func (v *Validator) CheckPermissions(claims jwt.MapClaims, required ...string) bool {
	// Try permissions claim first (for user tokens)
	if permsIface, ok := claims["permissions"]; ok {
		if perms, ok := permsIface.([]interface{}); ok {
			if v.checkPermList(perms, required...) {
				return true
			}
		}
	}

	// Fall back to roles claim (for service account tokens)
	if rolesIface, ok := claims["roles"]; ok {
		if roles, ok := rolesIface.([]interface{}); ok {
			if v.checkPermList(roles, required...) {
				return true
			}
		}
	}

	return false
}

// checkPermList is a helper that checks a permission/role list.
func (v *Validator) checkPermList(perms []interface{}, required ...string) bool {
	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
}