homelab-frontend/internal/auth/jwt.go

package auth

import (
	"context"
	"fmt"
	gohttp "net/http"
	"strings"
	"sync"
	"time"

	"github.com/MicahParks/keyfunc/v2"
	"github.com/golang-jwt/jwt/v5"
	"github.com/zitadel/zitadel-go/v3/pkg/authorization/oauth"
)

// issuerMatches reports whether iss is the configured issuer or a sub-path of
// it. Zitadel emits exactly "https://zitadel.riotpiao.com"; providers that
// mount an application path append a suffix (Authentik used
// "/application/o/<app>/"). A trailing slash on either side is ignored.
//
// An empty configured issuer disables the domain check; the JWKS signature
// check still applies. Production configuration always sets an issuer.
func issuerMatches(iss, configured string) bool {
	if configured == "" {
		return iss != ""
	}
	iss = strings.TrimSuffix(iss, "/")
	configured = strings.TrimSuffix(configured, "/")
	return iss == configured || strings.HasPrefix(iss, configured+"/")
}

// audienceValues normalises the aud claim. Zitadel access tokens carry an
// array, older providers carry a single string.
func audienceValues(claims jwt.MapClaims) ([]string, bool) {
	switch v := claims["aud"].(type) {
	case string:
		return []string{v}, v != ""
	case []string:
		return v, len(v) > 0
	case []interface{}:
		out := make([]string, 0, len(v))
		for _, item := range v {
			if s, ok := item.(string); ok {
				out = append(out, s)
			}
		}
		return out, len(out) > 0
	case jwt.ClaimStrings:
		return []string(v), len(v) > 0
	}
	return nil, false
}

// Validator validates JWTs against the JWKS of an OIDC provider (Zitadel).
// The issuer is configuration-driven: tokens whose iss claim equals the
// configured issuer, or lives under it as a sub-path, are accepted.
//
// When a ZitadelVerifier is attached (see NewZitadelValidator), validation
// delegates to the official ZITADEL Go SDK and this type only adapts the
// result to jwt.MapClaims for the rest of the gateway.
type Validator struct {
	issuer   string // Expected issuer (iss claim); empty disables the domain check
	audience string // Expected audience; currently logged only (any non-empty aud is accepted)
	jwksURL  string
	jwks     *keyfunc.JWKS
	mu       sync.Mutex
	zitadel  *ZitadelVerifier // non-nil when backed by the ZITADEL SDK
}

// NewZitadelValidator builds a Validator backed by the ZITADEL Go SDK.
// issuer is the expected iss claim, used for logging and for the legacy
// claim checks if the SDK verifier is ever bypassed. clientIDs is the
// audience allow-list. Discovery runs here, so a wrong domain fails at
// startup.
func NewZitadelValidator(ctx context.Context, issuer string, clientIDs []string, cfg ZitadelConfig) (*Validator, error) {
	return NewZitadelValidatorWithClient(ctx, issuer, clientIDs, cfg, nil)
}

// NewZitadelValidatorWithClient is NewZitadelValidator with an injected HTTP
// client, for tests that serve discovery and JWKS from httptest.
func NewZitadelValidatorWithClient(ctx context.Context, issuer string, clientIDs []string, cfg ZitadelConfig, httpClient *gohttp.Client) (*Validator, error) {
	cfg.ClientIDs = clientIDs
	if cfg.HTTPClient == nil {
		cfg.HTTPClient = httpClient
	}
	sdk, err := NewZitadelVerifier(ctx, cfg)
	if err != nil {
		return nil, err
	}
	return &Validator{
		issuer:  issuer,
		jwksURL: "", // unused on the SDK path
		zitadel: sdk,
	}, nil
}

// claimsFromIntrospection converts the SDK's verified claim set into
// jwt.MapClaims so the rest of the gateway keeps one claim representation.
// The SDK has already checked the signature, exp, nbf and iss; this only
// copies values across.
func claimsFromIntrospection(authCtx *oauth.IntrospectionContext) jwt.MapClaims {
	claims := jwt.MapClaims{}
	if authCtx == nil {
		return claims
	}
	for k, v := range authCtx.Claims {
		claims[k] = v
	}
	if authCtx.Issuer != "" {
		claims["iss"] = authCtx.Issuer
	}
	if authCtx.Subject != "" {
		claims["sub"] = authCtx.Subject
	}
	if len(authCtx.Audience) > 0 {
		aud := make([]interface{}, 0, len(authCtx.Audience))
		for _, a := range authCtx.Audience {
			aud = append(aud, a)
		}
		claims["aud"] = aud
	}
	if authCtx.Expiration != 0 {
		claims["exp"] = float64(authCtx.Expiration)
	}
	if authCtx.NotBefore != 0 {
		claims["nbf"] = float64(authCtx.NotBefore)
	}
	if authCtx.IssuedAt != 0 {
		claims["iat"] = float64(authCtx.IssuedAt)
	}
	if authCtx.ClientID != "" {
		claims["client_id"] = authCtx.ClientID
	}
	return claims
}

// NewValidator creates a new JWT validator for a service.
// issuer is the expected iss claim (exact value or prefix). audience is kept
// for compatibility; only non-emptiness is enforced. JWKS fetching is lazy
// (deferred until first validation).
func NewValidator(issuer, audience, jwksURL string) *Validator {
	return &Validator{
		issuer:   issuer,   // deprecated param, kept for compat
		audience: audience, // deprecated param, kept for compat
		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 v.jwks != nil || v.zitadel == nil {
		if err := v.ensureJWKS(); err != nil {
			return nil, err
		}
	}

	// SDK path: ZITADEL Go SDK performs signature, expiry and issuer checks.
	if v.zitadel != nil {
		// ZITADEL issues two access-token shapes. A plain JWS has three
		// dot-separated segments and can be verified locally. A JWE has five,
		// is encrypted with the instance key, and only ZITADEL can read it.
		// Report that case explicitly so the operator changes the client to
		// OIDC_TOKEN_TYPE_BEARER instead of debugging a signature failure.
		if strings.Count(tokenString, ".") == 4 {
			return nil, fmt.Errorf("token is an encrypted JWE (5 segments); " +
				"set accessTokenType=OIDC_TOKEN_TYPE_BEARER on the ZITADEL client so the gateway can verify it locally")
		}

		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		authCtx, err := v.zitadel.CheckAuthorization(ctx, authHeader)
		if err != nil {
			return nil, fmt.Errorf("token validation failed: %v", err)
		}
		return claimsFromIntrospection(authCtx), nil
	}

	// 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 against the configured issuer (exact match or sub-path).
	if iss, ok := claims["iss"].(string); !ok {
		return nil, fmt.Errorf("missing issuer claim")
	} else if !issuerMatches(iss, v.issuer) {
		return nil, fmt.Errorf("invalid issuer: %s (want %s)", iss, v.issuer)
	}

	// Check aud - Zitadel sends an array; accept a string or array.
	if auds, ok := audienceValues(claims); !ok {
		return nil, fmt.Errorf("missing audience claim")
	} else if len(auds) == 0 {
		return nil, fmt.Errorf("empty audience claim")
	}
	// The signature check plus the issuer check establish the trust boundary;
	// the audience value itself is not hardcoded.

	return claims, nil
}

// CheckPermissions checks if claims contain required permission(s).
// Understands three token shapes:
//   - "permissions": []string            (generic / legacy tokens)
//   - "roles": []string or map[string]…  (legacy service-account tokens)
//   - "urn:zitadel:iam:org:project:roles": map[string]…  (Zitadel access tokens,
//     where the map keys are role keys)
//
// Returns true if any required permission is found or wildcard "*" exists.
func (v *Validator) CheckPermissions(claims jwt.MapClaims, required ...string) bool {
	// 1. permissions claim (array of strings)
	if permsIface, ok := claims["permissions"]; ok {
		if perms, ok := permsIface.([]interface{}); ok {
			if v.checkPermList(perms, required...) {
				return true
			}
		}
	}

	// 2. roles claim: array (legacy) or map keyed by role name
	if rolesIface, ok := claims["roles"]; ok {
		switch roles := rolesIface.(type) {
		case []interface{}:
			if v.checkPermList(roles, required...) {
				return true
			}
		case map[string]interface{}:
			if v.checkPermMap(roles, required...) {
				return true
			}
		}
	}

	// 3. Zitadel access-token role assertion: map keyed by role key
	if rolesIface, ok := claims["urn:zitadel:iam:org:project:roles"]; ok {
		if roles, ok := rolesIface.(map[string]interface{}); ok {
			if v.checkPermMap(roles, required...) {
				return true
			}
		}
	}

	return false
}

// checkPermMap reports whether any key in a role map matches a required
// permission, or the map carries a wildcard key.
func (v *Validator) checkPermMap(roles map[string]interface{}, required ...string) bool {
	for roleKey := range roles {
		if roleKey == "*" {
			return true
		}
		for _, req := range required {
			if roleKey == req {
				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
}