easyai-ai-gateway/apps/api/internal/auth/oidc.go

package auth

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rsa"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

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

const maxOIDCResponseBytes = 1 << 20
const defaultOIDCHTTPTimeout = 10 * time.Second

type OIDCConfig struct {
	AppEnv                          string
	Issuer                          string
	Audience                        string
	TenantID                        string
	RolePrefix                      string
	RequiredScopes                  []string
	JWKSCacheTTL                    time.Duration
	IntrospectionEnabled            bool
	IntrospectionClientID           string
	IntrospectionClientSecret       string
	IntrospectionCredentialProvider func(context.Context) (string, []byte, error)
	HTTPClient                      *http.Client
	SecurityEventEvaluator          func(context.Context, OIDCSecurityEventIdentity) (OIDCSecurityEventEvaluation, error)
	IntrospectionObserver           func(string)
	JWKSRefreshFailureObserver      func()
}

type OIDCSecurityEventIdentity struct {
	Issuer   string
	TenantID string
	Subject  string
	IssuedAt time.Time
}

type OIDCSecurityEventEvaluation struct {
	Enabled              bool
	Revoked              bool
	RequireIntrospection bool
}

type OIDCVerifier struct {
	config                OIDCConfig
	client                *http.Client
	mu                    sync.Mutex
	keys                  map[string]any
	expiresAt             time.Time
	introspectionEndpoint string
}

type discoveryDocument struct {
	Issuer                string `json:"issuer"`
	JWKSURI               string `json:"jwks_uri"`
	IntrospectionEndpoint string `json:"introspection_endpoint"`
}

type jsonWebKeySet struct {
	Keys []jsonWebKey `json:"keys"`
}

type jsonWebKey struct {
	KID string `json:"kid"`
	KTY string `json:"kty"`
	Use string `json:"use"`
	Alg string `json:"alg"`
	N   string `json:"n"`
	E   string `json:"e"`
	Crv string `json:"crv"`
	X   string `json:"x"`
	Y   string `json:"y"`
}

func NewOIDCVerifier(config OIDCConfig) (*OIDCVerifier, error) {
	config.Issuer = strings.TrimRight(strings.TrimSpace(config.Issuer), "/")
	config.Audience = strings.TrimSpace(config.Audience)
	config.TenantID = strings.TrimSpace(config.TenantID)
	config.RolePrefix = strings.TrimSpace(config.RolePrefix)
	if err := validatePublicURL(config.Issuer, config.AppEnv); err != nil || config.Audience == "" || config.TenantID == "" || config.RolePrefix == "" {
		return nil, errors.New("issuer, audience, tenant and role prefix are required")
	}
	if config.IntrospectionEnabled && config.IntrospectionCredentialProvider == nil &&
		(strings.TrimSpace(config.IntrospectionClientID) == "" || config.IntrospectionClientSecret == "") {
		return nil, errors.New("introspection client credentials are required when introspection is enabled")
	}
	if config.JWKSCacheTTL <= 0 {
		config.JWKSCacheTTL = 5 * time.Minute
	}
	client := config.HTTPClient
	if client == nil {
		transport := http.DefaultTransport.(*http.Transport).Clone()
		transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12}
		transport.DisableCompression = true
		client = &http.Client{Timeout: defaultOIDCHTTPTimeout, CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
			return http.ErrUseLastResponse
		}, Transport: transport}
	}
	return &OIDCVerifier{config: config, client: client, keys: map[string]any{}}, nil
}

// ValidateConfiguration eagerly validates Discovery, JWKS and, when enabled,
// the RFC 7662 endpoint and machine credential before a runtime is activated.
func (v *OIDCVerifier) ValidateConfiguration(ctx context.Context) error {
	if err := v.refresh(ctx, true); err != nil {
		return err
	}
	if v.config.IntrospectionEnabled || v.config.SecurityEventEvaluator != nil {
		if _, err := v.introspect(ctx, "identity-configuration-probe"); err != nil {
			return err
		}
	}
	return nil
}

func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
	parser := jwt.NewParser(jwt.WithValidMethods([]string{"RS256", "ES256"}))
	unverified, _, err := parser.ParseUnverified(raw, jwt.MapClaims{})
	if err != nil || unverified == nil {
		return nil, oidcUnauthorized("TOKEN_ENVELOPE_INVALID", "token envelope is invalid", err)
	}
	unverifiedClaims, ok := unverified.Claims.(jwt.MapClaims)
	if !ok {
		return nil, oidcUnauthorized("TOKEN_ENVELOPE_INVALID", "token envelope is invalid", nil)
	}
	if stringClaim(unverifiedClaims, "iss") == "" {
		return nil, oidcUnauthorized("ISSUER_MISSING", "issuer is missing", nil)
	}
	if len(stringSliceClaim(unverifiedClaims, "aud")) == 0 {
		return nil, oidcUnauthorized("AUDIENCE_MISSING", "audience is missing", nil)
	}
	if _, exists := unverifiedClaims["exp"]; !exists {
		return nil, oidcUnauthorized("EXP_MISSING", "exp is missing", nil)
	}
	kid, _ := unverified.Header["kid"].(string)
	if kid == "" {
		return nil, oidcUnauthorized("KID_MISSING", "kid is missing", nil)
	}
	key, err := v.key(ctx, kid)
	if err != nil {
		return nil, oidcUnauthorized("SIGNING_KEY_LOOKUP_FAILED", "signing key lookup failed", err)
	}
	token, err := jwt.Parse(raw, func(token *jwt.Token) (any, error) {
		if token.Header["kid"] != kid {
			return nil, ErrUnauthorized
		}
		return key, nil
	}, jwt.WithValidMethods([]string{"RS256", "ES256"}), jwt.WithIssuer(v.config.Issuer),
		jwt.WithAudience(v.config.Audience), jwt.WithExpirationRequired(), jwt.WithLeeway(30*time.Second))
	if err != nil || !token.Valid {
		return nil, oidcUnauthorized(registeredClaimsValidationCategory(err), "signature or registered claims are invalid", err)
	}
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok || stringClaim(claims, "sub") == "" || stringClaim(claims, "tid") != v.config.TenantID {
		return nil, oidcUnauthorized("STABLE_IDENTITY_CLAIMS_INVALID", "stable identity claims are invalid", nil)
	}
	expiresAt, ok := numericDateClaim(claims["exp"])
	if !ok {
		return nil, oidcUnauthorized("EXP_INVALID", "exp is invalid", nil)
	}
	if _, ok := numericDateClaim(claims["nbf"]); !ok {
		return nil, oidcUnauthorized("NBF_MISSING", "nbf is missing", nil)
	}
	issuedAt, hasIssuedAt := numericDateClaim(claims["iat"])
	scopes := scopeClaims(claims)
	if !containsAll(scopes, v.config.RequiredScopes) {
		return nil, oidcUnauthorized("REQUIRED_SCOPE_MISSING", "required scope is missing", nil)
	}
	roles := gatewayRoles(stringSliceClaim(claims, "roles"), v.config.RolePrefix)
	if len(roles) == 0 {
		roles = gatewayRoles(stringSliceClaim(claims, "role"), v.config.RolePrefix)
	}
	if len(roles) == 0 {
		return nil, oidcUnauthorized("MAPPED_ROLE_MISSING", "mapped role is missing", nil)
	}
	if v.config.SecurityEventEvaluator != nil {
		evaluation, evaluateErr := v.config.SecurityEventEvaluator(ctx, OIDCSecurityEventIdentity{
			Issuer: v.config.Issuer, TenantID: v.config.TenantID, Subject: stringClaim(claims, "sub"), IssuedAt: issuedAt,
		})
		if evaluateErr != nil {
			return nil, NewRequestAuthError(http.StatusServiceUnavailable, "OIDC_SECURITY_EVENT_STATE_UNAVAILABLE", "认证撤销状态暂时不可用")
		}
		if evaluation.Enabled && !hasIssuedAt {
			return nil, oidcUnauthorized("SECURITY_EVENT_IAT_MISSING", "iat is required when security events are enabled", nil)
		}
		if evaluation.Revoked {
			return nil, oidcUnauthorized("TOKEN_REVOKED", "token was issued before the revocation watermark", nil)
		}
		if evaluation.RequireIntrospection {
			active, introspectionErr := v.introspect(ctx, raw)
			if introspectionErr != nil {
				return nil, NewRequestAuthError(http.StatusServiceUnavailable, "OIDC_INTROSPECTION_UNAVAILABLE", "认证中心内省暂时不可用")
			}
			if !active {
				return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", nil)
			}
		} else if !evaluation.Enabled && v.config.IntrospectionEnabled {
			active, introspectionErr := v.introspect(ctx, raw)
			if introspectionErr != nil || !active {
				return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", introspectionErr)
			}
		}
	} else if v.config.IntrospectionEnabled {
		active, err := v.introspect(ctx, raw)
		if err != nil || !active {
			return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", err)
		}
	}
	username := stringClaim(claims, "preferred_username")
	if username == "" {
		username = stringClaim(claims, "username")
	}
	return &User{
		ID: stringClaim(claims, "sub"), Username: username, Roles: roles,
		TenantID: v.config.TenantID, Source: "oidc", TokenExpiresAt: expiresAt, TokenIssuedAt: issuedAt, Issuer: v.config.Issuer,
	}, nil
}

type oidcValidationError struct {
	category string
	reason   string
	cause    error
}

func (err *oidcValidationError) Error() string {
	if err.cause != nil {
		return fmt.Sprintf("%v: %s: %v", ErrUnauthorized, err.reason, err.cause)
	}
	return fmt.Sprintf("%v: %s", ErrUnauthorized, err.reason)
}

func (err *oidcValidationError) Unwrap() []error {
	if err.cause == nil {
		return []error{ErrUnauthorized}
	}
	return []error{ErrUnauthorized, err.cause}
}

func OIDCValidationCategory(err error) string {
	var validationError *oidcValidationError
	if errors.As(err, &validationError) {
		return validationError.category
	}
	return "UNCLASSIFIED"
}

func registeredClaimsValidationCategory(err error) string {
	switch {
	case errors.Is(err, jwt.ErrTokenRequiredClaimMissing):
		return "REQUIRED_CLAIM_MISSING"
	case errors.Is(err, jwt.ErrTokenInvalidAudience):
		return "AUDIENCE_INVALID"
	case errors.Is(err, jwt.ErrTokenInvalidIssuer):
		return "ISSUER_INVALID"
	case errors.Is(err, jwt.ErrTokenExpired):
		return "TOKEN_EXPIRED"
	case errors.Is(err, jwt.ErrTokenNotValidYet):
		return "TOKEN_NOT_VALID_YET"
	case errors.Is(err, jwt.ErrTokenUsedBeforeIssued):
		return "TOKEN_USED_BEFORE_ISSUED"
	case errors.Is(err, jwt.ErrTokenSignatureInvalid):
		return "SIGNATURE_INVALID"
	case errors.Is(err, jwt.ErrTokenMalformed):
		return "TOKEN_MALFORMED"
	case errors.Is(err, jwt.ErrTokenUnverifiable):
		return "TOKEN_UNVERIFIABLE"
	case errors.Is(err, jwt.ErrTokenInvalidClaims):
		return "CLAIMS_INVALID"
	default:
		return "REGISTERED_CLAIMS_INVALID"
	}
}

func oidcUnauthorized(category, reason string, cause error) error {
	return &oidcValidationError{
		category: category,
		reason:   reason,
		cause:    cause,
	}
}

func (v *OIDCVerifier) key(ctx context.Context, kid string) (any, error) {
	if err := v.refresh(ctx, false); err != nil {
		if v.config.JWKSRefreshFailureObserver != nil {
			v.config.JWKSRefreshFailureObserver()
		}
		return nil, err
	}
	v.mu.Lock()
	key, ok := v.keys[kid]
	v.mu.Unlock()
	if ok {
		return key, nil
	}
	if err := v.refresh(ctx, true); err != nil {
		if v.config.JWKSRefreshFailureObserver != nil {
			v.config.JWKSRefreshFailureObserver()
		}
		return nil, err
	}
	v.mu.Lock()
	defer v.mu.Unlock()
	key, ok = v.keys[kid]
	if !ok {
		return nil, errors.New("signing key not found")
	}
	return key, nil
}

func (v *OIDCVerifier) refresh(ctx context.Context, force bool) error {
	v.mu.Lock()
	defer v.mu.Unlock()
	if !force && len(v.keys) > 0 && time.Now().Before(v.expiresAt) {
		return nil
	}
	discoveryURL := v.config.Issuer + "/.well-known/openid-configuration"
	var discovery discoveryDocument
	if err := v.fetchJSON(ctx, discoveryURL, &discovery); err != nil {
		return err
	}
	if discovery.Issuer != v.config.Issuer || validatePublicURL(discovery.JWKSURI, v.config.AppEnv) != nil {
		return errors.New("OIDC discovery metadata is invalid")
	}
	if (v.config.IntrospectionEnabled || v.config.SecurityEventEvaluator != nil) && validatePublicURL(discovery.IntrospectionEndpoint, v.config.AppEnv) != nil {
		return errors.New("OIDC introspection metadata is invalid")
	}
	var set jsonWebKeySet
	if err := v.fetchJSON(ctx, discovery.JWKSURI, &set); err != nil {
		return err
	}
	keys := make(map[string]any, len(set.Keys))
	for _, item := range set.Keys {
		if item.KID == "" || item.Use != "" && item.Use != "sig" {
			continue
		}
		key, err := item.publicKey()
		if err == nil {
			keys[item.KID] = key
		}
	}
	if len(keys) == 0 {
		return errors.New("OIDC JWKS contains no supported signing keys")
	}
	v.keys = keys
	v.expiresAt = time.Now().Add(v.config.JWKSCacheTTL)
	v.introspectionEndpoint = discovery.IntrospectionEndpoint
	return nil
}

func (v *OIDCVerifier) introspect(ctx context.Context, raw string) (bool, error) {
	outcome := "failed"
	defer func() {
		if v.config.IntrospectionObserver != nil {
			v.config.IntrospectionObserver(outcome)
		}
	}()
	v.mu.Lock()
	endpoint := v.introspectionEndpoint
	v.mu.Unlock()
	if endpoint == "" {
		return false, errors.New("OIDC introspection endpoint is unavailable")
	}
	form := url.Values{"token": {raw}, "token_type_hint": {"access_token"}}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode()))
	if err != nil {
		return false, err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.Header.Set("Accept", "application/json")
	clientID := v.config.IntrospectionClientID
	secret := []byte(v.config.IntrospectionClientSecret)
	if v.config.IntrospectionCredentialProvider != nil {
		clientID, secret, err = v.config.IntrospectionCredentialProvider(ctx)
		if err != nil {
			return false, err
		}
	}
	defer clear(secret)
	if strings.TrimSpace(clientID) == "" || len(secret) < 16 {
		return false, errors.New("OIDC introspection credential is unavailable")
	}
	request.SetBasicAuth(clientID, string(secret))
	response, err := v.client.Do(request)
	if err != nil {
		return false, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, maxOIDCResponseBytes))
		return false, fmt.Errorf("OIDC introspection returned HTTP %d", response.StatusCode)
	}
	var result struct {
		Active bool `json:"active"`
	}
	decoder := json.NewDecoder(io.LimitReader(response.Body, maxOIDCResponseBytes))
	if err := decoder.Decode(&result); err != nil {
		return false, errors.New("OIDC introspection response is invalid")
	}
	if result.Active {
		outcome = "active"
	} else {
		outcome = "inactive"
	}
	return result.Active, nil
}

func (v *OIDCVerifier) fetchJSON(ctx context.Context, endpoint string, output any) error {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return err
	}
	request.Header.Set("Accept", "application/json")
	response, err := v.client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("OIDC endpoint returned HTTP %d", response.StatusCode)
	}
	payload, err := io.ReadAll(io.LimitReader(response.Body, maxOIDCResponseBytes+1))
	if err != nil || len(payload) > maxOIDCResponseBytes {
		return errors.New("OIDC response is invalid")
	}
	decoder := json.NewDecoder(strings.NewReader(string(payload)))
	return decoder.Decode(output)
}

func (j jsonWebKey) publicKey() (any, error) {
	switch {
	case j.KTY == "RSA" && (j.Alg == "" || j.Alg == "RS256"):
		n, err := decodeBigInt(j.N)
		if err != nil {
			return nil, err
		}
		e, err := decodeBigInt(j.E)
		if err != nil || !e.IsInt64() || e.Int64() <= 0 {
			return nil, errors.New("invalid RSA exponent")
		}
		return &rsa.PublicKey{N: n, E: int(e.Int64())}, nil
	case j.KTY == "EC" && j.Crv == "P-256" && (j.Alg == "" || j.Alg == "ES256"):
		x, err := decodeBigInt(j.X)
		if err != nil {
			return nil, err
		}
		y, err := decodeBigInt(j.Y)
		if err != nil || !elliptic.P256().IsOnCurve(x, y) {
			return nil, errors.New("invalid EC point")
		}
		return &ecdsa.PublicKey{Curve: elliptic.P256(), X: x, Y: y}, nil
	default:
		return nil, errors.New("unsupported JWK")
	}
}

func decodeBigInt(value string) (*big.Int, error) {
	payload, err := base64.RawURLEncoding.DecodeString(value)
	if err != nil || len(payload) == 0 {
		return nil, errors.New("invalid JWK integer")
	}
	return new(big.Int).SetBytes(payload), nil
}

func validatePublicURL(raw, appEnv string) error {
	parsed, err := url.Parse(raw)
	if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
		return errors.New("invalid URL")
	}
	if parsed.Scheme == "https" {
		return nil
	}
	if parsed.Scheme == "http" && isLocalOIDCEnvironment(appEnv) && isLoopbackOIDCHost(parsed.Hostname()) {
		return nil
	}
	return errors.New("URL must use HTTPS")
}

func isLocalOIDCEnvironment(value string) bool {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "development", "dev", "local", "test":
		return true
	default:
		return false
	}
}

func isLoopbackOIDCHost(value string) bool {
	value = strings.ToLower(strings.TrimSpace(value))
	if value == "localhost" {
		return true
	}
	ip := net.ParseIP(value)
	return ip != nil && ip.IsLoopback()
}

func numericDateClaim(value any) (time.Time, bool) {
	switch typed := value.(type) {
	case float64:
		return time.Unix(int64(typed), 0), true
	case json.Number:
		seconds, err := typed.Int64()
		return time.Unix(seconds, 0), err == nil
	default:
		return time.Time{}, false
	}
}

func scopeClaims(claims jwt.MapClaims) []string {
	result := make([]string, 0)
	seen := map[string]bool{}
	for _, scope := range strings.Fields(stringClaim(claims, "scope")) {
		if !seen[scope] {
			result = append(result, scope)
			seen[scope] = true
		}
	}
	for _, scope := range stringSliceClaim(claims, "scp") {
		if !seen[scope] {
			result = append(result, scope)
			seen[scope] = true
		}
	}
	return result
}

func containsAll(actual, required []string) bool {
	set := make(map[string]struct{}, len(actual))
	for _, item := range actual {
		set[item] = struct{}{}
	}
	for _, item := range required {
		if _, ok := set[item]; !ok {
			return false
		}
	}
	return true
}

func gatewayRoles(external []string, prefix string) []string {
	allowed := map[string]bool{"user": true, "creator": true, "operator": true, "manager": true, "admin": true}
	roles := make([]string, 0, len(external))
	seen := map[string]bool{}
	for _, role := range external {
		if !strings.HasPrefix(role, prefix) {
			continue
		}
		local := strings.TrimPrefix(role, prefix)
		if allowed[local] && !seen[local] {
			roles = append(roles, local)
			seen[local] = true
		}
	}
	return roles
}