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

package auth

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/coreos/go-oidc/v3/oidc"
	"golang.org/x/oauth2"
)

var ErrOIDCInvalidGrant = errors.New("OIDC refresh token is invalid")

var errOIDCResponseTooLarge = errors.New("OIDC response exceeds size limit")

type OIDCPublicClientConfig struct {
	Issuer                string
	ClientID              string
	RedirectURI           string
	PostLogoutRedirectURI string
	Scopes                []string
	HTTPClient            *http.Client
}

type OIDCTokenResponse struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	IDToken      string `json:"id_token"`
	TokenType    string `json:"token_type"`
	ExpiresIn    int    `json:"expires_in"`
}

type OIDCPublicClient struct {
	config          OIDCPublicClientConfig
	client          *http.Client
	mu              sync.Mutex
	oauth2Config    *oauth2.Config
	metadata        oidcClientDiscovery
	idTokenVerifier *oidc.IDTokenVerifier
}

type oidcClientDiscovery struct {
	Issuer                string `json:"issuer"`
	AuthorizationEndpoint string `json:"authorization_endpoint"`
	TokenEndpoint         string `json:"token_endpoint"`
	JWKSURI               string `json:"jwks_uri"`
	RevocationEndpoint    string `json:"revocation_endpoint"`
	EndSessionEndpoint    string `json:"end_session_endpoint"`
}

func NewOIDCPublicClient(config OIDCPublicClientConfig) (*OIDCPublicClient, error) {
	config.Issuer = strings.TrimRight(strings.TrimSpace(config.Issuer), "/")
	config.ClientID = strings.TrimSpace(config.ClientID)
	config.RedirectURI = strings.TrimSpace(config.RedirectURI)
	config.PostLogoutRedirectURI = strings.TrimSpace(config.PostLogoutRedirectURI)
	config.Scopes = normalizedScopes(config.Scopes)
	for _, scope := range config.Scopes {
		if strings.EqualFold(scope, oidc.ScopeOfflineAccess) {
			return nil, errors.New("offline_access is not allowed for Gateway browser sessions")
		}
	}
	if validatePublicURL(config.Issuer) != nil || config.ClientID == "" || validatePublicURL(config.RedirectURI) != nil || validatePublicURL(config.PostLogoutRedirectURI) != nil {
		return nil, errors.New("issuer, public client id and exact redirect URLs are required")
	}
	return &OIDCPublicClient{config: config, client: newOIDCHTTPClient(config.HTTPClient)}, nil
}

// ValidateConfiguration eagerly checks public-client Discovery metadata.
func (c *OIDCPublicClient) ValidateConfiguration(ctx context.Context) error {
	_, _, err := c.configuration(ctx)
	return err
}

func (c *OIDCPublicClient) AuthorizationURL(ctx context.Context, state, nonce, pkceVerifier string) (string, error) {
	if strings.TrimSpace(state) == "" || strings.TrimSpace(nonce) == "" || !validPKCEVerifier(pkceVerifier) {
		return "", errors.New("state, nonce and PKCE verifier are required")
	}
	config, _, err := c.configuration(ctx)
	if err != nil {
		return "", err
	}
	return config.AuthCodeURL(state, oidc.Nonce(nonce), oauth2.S256ChallengeOption(pkceVerifier)), nil
}

func (c *OIDCPublicClient) ExchangeCode(ctx context.Context, code, verifier string) (OIDCTokenResponse, error) {
	if strings.TrimSpace(code) == "" || !validPKCEVerifier(verifier) {
		return OIDCTokenResponse{}, errors.New("authorization code and PKCE verifier are required")
	}
	config, _, err := c.configuration(ctx)
	if err != nil {
		return OIDCTokenResponse{}, err
	}
	token, err := config.Exchange(c.requestContext(ctx), code, oauth2.VerifierOption(verifier))
	if err != nil {
		return OIDCTokenResponse{}, oidcTokenError(err)
	}
	return oidcTokenResponse(token)
}

func (c *OIDCPublicClient) Refresh(ctx context.Context, refreshToken string) (OIDCTokenResponse, error) {
	if strings.TrimSpace(refreshToken) == "" {
		return OIDCTokenResponse{}, ErrOIDCInvalidGrant
	}
	config, _, err := c.configuration(ctx)
	if err != nil {
		return OIDCTokenResponse{}, err
	}
	token, err := config.TokenSource(c.requestContext(ctx), &oauth2.Token{RefreshToken: refreshToken}).Token()
	if err != nil {
		return OIDCTokenResponse{}, oidcTokenError(err)
	}
	return oidcTokenResponse(token)
}

func (c *OIDCPublicClient) VerifyIDToken(ctx context.Context, raw, expectedNonce string) (string, error) {
	expectedNonce = strings.TrimSpace(expectedNonce)
	if strings.TrimSpace(raw) == "" || expectedNonce == "" {
		return "", oidcUnauthorized("ID token validation context is invalid", nil)
	}
	if _, _, err := c.configuration(ctx); err != nil {
		return "", oidcUnauthorized("ID token provider discovery failed", err)
	}
	c.mu.Lock()
	verifier := c.idTokenVerifier
	c.mu.Unlock()
	if verifier == nil {
		return "", oidcUnauthorized("ID token verifier is unavailable", nil)
	}
	token, err := verifier.Verify(c.requestContext(ctx), raw)
	if err != nil {
		return "", oidcUnauthorized("ID token signature or registered claims are invalid", nil)
	}
	if token.Subject == "" || token.Nonce != expectedNonce {
		return "", oidcUnauthorized("ID token subject or nonce is invalid", nil)
	}
	var claims map[string]any
	if err := token.Claims(&claims); err != nil {
		return "", oidcUnauthorized("ID token claims are invalid", nil)
	}
	if _, ok := numericDateClaim(claims["nbf"]); !ok {
		return "", oidcUnauthorized("ID token nbf is missing", nil)
	}
	return token.Subject, nil
}

func (c *OIDCPublicClient) RevokeRefreshToken(ctx context.Context, refreshToken string) error {
	if strings.TrimSpace(refreshToken) == "" {
		return nil
	}
	_, metadata, err := c.configuration(ctx)
	if err != nil {
		return err
	}
	if metadata.RevocationEndpoint == "" {
		return errors.New("OIDC revocation endpoint is unavailable")
	}
	response, err := c.postForm(ctx, metadata.RevocationEndpoint, url.Values{
		"client_id": {c.config.ClientID}, "token": {refreshToken}, "token_type_hint": {"refresh_token"},
	})
	if err != nil {
		return err
	}
	defer response.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, maxOIDCResponseBytes))
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return fmt.Errorf("OIDC revocation returned HTTP %d", response.StatusCode)
	}
	return nil
}

func (c *OIDCPublicClient) EndSessionURL(ctx context.Context, idTokenHint string) (string, error) {
	_, metadata, err := c.configuration(ctx)
	if err != nil {
		return "", err
	}
	if metadata.EndSessionEndpoint == "" {
		return c.config.PostLogoutRedirectURI, nil
	}
	parsed, err := url.Parse(metadata.EndSessionEndpoint)
	if err != nil {
		return "", errors.New("OIDC end session endpoint is invalid")
	}
	query := parsed.Query()
	query.Set("client_id", c.config.ClientID)
	query.Set("post_logout_redirect_uri", c.config.PostLogoutRedirectURI)
	if strings.TrimSpace(idTokenHint) != "" {
		query.Set("id_token_hint", idTokenHint)
	}
	parsed.RawQuery = query.Encode()
	return parsed.String(), nil
}

func (c *OIDCPublicClient) postForm(ctx context.Context, endpoint string, form url.Values) (*http.Response, error) {
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.Header.Set("Accept", "application/json")
	return c.client.Do(request)
}

func (c *OIDCPublicClient) configuration(ctx context.Context) (*oauth2.Config, oidcClientDiscovery, error) {
	c.mu.Lock()
	if c.oauth2Config != nil {
		config, metadata := c.oauth2Config, c.metadata
		c.mu.Unlock()
		return config, metadata, nil
	}
	c.mu.Unlock()

	provider, err := oidc.NewProvider(c.requestContext(ctx), c.config.Issuer)
	if err != nil {
		return nil, oidcClientDiscovery{}, errors.New("OIDC discovery failed")
	}
	var metadata oidcClientDiscovery
	if err := provider.Claims(&metadata); err != nil || metadata.Issuer != c.config.Issuer ||
		validatePublicURL(metadata.AuthorizationEndpoint) != nil || validatePublicURL(metadata.TokenEndpoint) != nil || validatePublicURL(metadata.JWKSURI) != nil ||
		metadata.RevocationEndpoint != "" && validatePublicURL(metadata.RevocationEndpoint) != nil ||
		metadata.EndSessionEndpoint != "" && validatePublicURL(metadata.EndSessionEndpoint) != nil {
		return nil, oidcClientDiscovery{}, errors.New("OIDC discovery metadata is invalid")
	}
	endpoint := provider.Endpoint()
	// Gateway is a public client. Force client_id into the form and never probe
	// HTTP Basic authentication with an empty client secret.
	endpoint.AuthStyle = oauth2.AuthStyleInParams
	config := &oauth2.Config{
		ClientID: c.config.ClientID, RedirectURL: c.config.RedirectURI,
		Endpoint: endpoint, Scopes: append([]string(nil), c.config.Scopes...),
	}
	verifierContext := oidc.ClientContext(context.Background(), c.client)
	idTokenVerifier := provider.VerifierContext(verifierContext, &oidc.Config{
		ClientID: c.config.ClientID, SupportedSigningAlgs: []string{oidc.RS256, oidc.ES256},
	})
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.oauth2Config == nil {
		c.oauth2Config = config
		c.metadata = metadata
		c.idTokenVerifier = idTokenVerifier
	}
	return c.oauth2Config, c.metadata, nil
}

func (c *OIDCPublicClient) requestContext(ctx context.Context) context.Context {
	return oidc.ClientContext(ctx, c.client)
}

func oidcTokenError(err error) error {
	var retrieveError *oauth2.RetrieveError
	if errors.As(err, &retrieveError) {
		if retrieveError.ErrorCode == "invalid_grant" {
			return ErrOIDCInvalidGrant
		}
		if retrieveError.Response != nil {
			return fmt.Errorf("OIDC token endpoint returned HTTP %d", retrieveError.Response.StatusCode)
		}
	}
	return errors.New("OIDC token endpoint request failed")
}

func oidcTokenResponse(token *oauth2.Token) (OIDCTokenResponse, error) {
	if token == nil || strings.TrimSpace(token.AccessToken) == "" {
		return OIDCTokenResponse{}, errors.New("OIDC token response is invalid")
	}
	idToken, _ := token.Extra("id_token").(string)
	return OIDCTokenResponse{
		AccessToken: token.AccessToken, RefreshToken: token.RefreshToken, IDToken: idToken,
		TokenType: token.TokenType, ExpiresIn: int(token.ExpiresIn),
	}, nil
}

func validPKCEVerifier(value string) bool {
	if len(value) < 43 || len(value) > 128 {
		return false
	}
	for _, char := range value {
		if char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || char >= '0' && char <= '9' || strings.ContainsRune("-._~", char) {
			continue
		}
		return false
	}
	return true
}

func newOIDCHTTPClient(base *http.Client) *http.Client {
	if base == nil {
		base = &http.Client{Timeout: defaultOIDCHTTPTimeout}
	}
	client := *base
	transport := client.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}
	client.Transport = oidcResponseLimitTransport{base: transport}
	client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse }
	return &client
}

type oidcResponseLimitTransport struct {
	base http.RoundTripper
}

func (t oidcResponseLimitTransport) RoundTrip(request *http.Request) (*http.Response, error) {
	response, err := t.base.RoundTrip(request)
	if err != nil {
		return nil, err
	}
	response.Body = &oidcLimitedReadCloser{body: response.Body, remaining: maxOIDCResponseBytes}
	return response, nil
}

type oidcLimitedReadCloser struct {
	body      io.ReadCloser
	remaining int64
}

func (r *oidcLimitedReadCloser) Read(buffer []byte) (int, error) {
	if r.remaining == 0 {
		var extra [1]byte
		if count, err := r.body.Read(extra[:]); count > 0 {
			return 0, errOIDCResponseTooLarge
		} else {
			return 0, err
		}
	}
	if int64(len(buffer)) > r.remaining {
		buffer = buffer[:r.remaining]
	}
	count, err := r.body.Read(buffer)
	r.remaining -= int64(count)
	return count, err
}

func (r *oidcLimitedReadCloser) Close() error {
	return r.body.Close()
}

func normalizedScopes(scopes []string) []string {
	result := make([]string, 0, len(scopes)+1)
	seen := map[string]struct{}{}
	for _, scope := range append([]string{oidc.ScopeOpenID}, scopes...) {
		scope = strings.TrimSpace(scope)
		if scope == "" {
			continue
		}
		if _, ok := seen[scope]; ok {
			continue
		}
		seen[scope] = struct{}{}
		result = append(result, scope)
	}
	return result
}