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

package auth

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

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

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
	metadata oidcClientDiscovery
}

type oidcClientDiscovery struct {
	Issuer                string `json:"issuer"`
	AuthorizationEndpoint string `json:"authorization_endpoint"`
	TokenEndpoint         string `json:"token_endpoint"`
	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, "offline_access") {
			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")
	}
	client := config.HTTPClient
	if client == nil {
		client = &http.Client{Timeout: defaultOIDCHTTPTimeout, CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
			return http.ErrUseLastResponse
		}}
	}
	return &OIDCPublicClient{config: config, client: client}, nil
}

func (c *OIDCPublicClient) AuthorizationURL(ctx context.Context, state, nonce, codeChallenge string) (string, error) {
	if strings.TrimSpace(state) == "" || strings.TrimSpace(nonce) == "" || strings.TrimSpace(codeChallenge) == "" {
		return "", errors.New("state, nonce and PKCE challenge are required")
	}
	metadata, err := c.discovery(ctx)
	if err != nil {
		return "", err
	}
	parsed, err := url.Parse(metadata.AuthorizationEndpoint)
	if err != nil {
		return "", errors.New("OIDC authorization endpoint is invalid")
	}
	query := parsed.Query()
	query.Set("response_type", "code")
	query.Set("client_id", c.config.ClientID)
	query.Set("redirect_uri", c.config.RedirectURI)
	query.Set("scope", strings.Join(c.config.Scopes, " "))
	query.Set("state", state)
	query.Set("nonce", nonce)
	query.Set("code_challenge", codeChallenge)
	query.Set("code_challenge_method", "S256")
	parsed.RawQuery = query.Encode()
	return parsed.String(), nil
}

func (c *OIDCPublicClient) ExchangeCode(ctx context.Context, code, verifier string) (OIDCTokenResponse, error) {
	if strings.TrimSpace(code) == "" || strings.TrimSpace(verifier) == "" {
		return OIDCTokenResponse{}, errors.New("authorization code and PKCE verifier are required")
	}
	return c.token(ctx, url.Values{
		"grant_type": {"authorization_code"}, "client_id": {c.config.ClientID},
		"redirect_uri": {c.config.RedirectURI}, "code": {code}, "code_verifier": {verifier},
	})
}

func (c *OIDCPublicClient) Refresh(ctx context.Context, refreshToken string) (OIDCTokenResponse, error) {
	if strings.TrimSpace(refreshToken) == "" {
		return OIDCTokenResponse{}, ErrOIDCInvalidGrant
	}
	return c.token(ctx, url.Values{
		"grant_type": {"refresh_token"}, "client_id": {c.config.ClientID}, "refresh_token": {refreshToken},
	})
}

func (c *OIDCPublicClient) RevokeRefreshToken(ctx context.Context, refreshToken string) error {
	if strings.TrimSpace(refreshToken) == "" {
		return nil
	}
	metadata, err := c.discovery(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.discovery(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) token(ctx context.Context, form url.Values) (OIDCTokenResponse, error) {
	metadata, err := c.discovery(ctx)
	if err != nil {
		return OIDCTokenResponse{}, err
	}
	response, err := c.postForm(ctx, metadata.TokenEndpoint, form)
	if err != nil {
		return OIDCTokenResponse{}, err
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		var oauthError struct {
			Error string `json:"error"`
		}
		_ = json.NewDecoder(io.LimitReader(response.Body, maxOIDCResponseBytes)).Decode(&oauthError)
		if oauthError.Error == "invalid_grant" {
			return OIDCTokenResponse{}, ErrOIDCInvalidGrant
		}
		return OIDCTokenResponse{}, fmt.Errorf("OIDC token endpoint returned HTTP %d", response.StatusCode)
	}
	var result OIDCTokenResponse
	if err := json.NewDecoder(io.LimitReader(response.Body, maxOIDCResponseBytes)).Decode(&result); err != nil || strings.TrimSpace(result.AccessToken) == "" {
		return OIDCTokenResponse{}, errors.New("OIDC token response is invalid")
	}
	return result, 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) discovery(ctx context.Context) (oidcClientDiscovery, error) {
	c.mu.Lock()
	if c.metadata.Issuer != "" {
		metadata := c.metadata
		c.mu.Unlock()
		return metadata, nil
	}
	c.mu.Unlock()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, c.config.Issuer+"/.well-known/openid-configuration", nil)
	if err != nil {
		return oidcClientDiscovery{}, err
	}
	request.Header.Set("Accept", "application/json")
	response, err := c.client.Do(request)
	if err != nil {
		return oidcClientDiscovery{}, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return oidcClientDiscovery{}, fmt.Errorf("OIDC discovery returned HTTP %d", response.StatusCode)
	}
	var metadata oidcClientDiscovery
	if err := json.NewDecoder(io.LimitReader(response.Body, maxOIDCResponseBytes)).Decode(&metadata); err != nil ||
		metadata.Issuer != c.config.Issuer || validatePublicURL(metadata.AuthorizationEndpoint) != nil || validatePublicURL(metadata.TokenEndpoint) != nil ||
		metadata.RevocationEndpoint != "" && validatePublicURL(metadata.RevocationEndpoint) != nil ||
		metadata.EndSessionEndpoint != "" && validatePublicURL(metadata.EndSessionEndpoint) != nil {
		return oidcClientDiscovery{}, errors.New("OIDC discovery metadata is invalid")
	}
	c.mu.Lock()
	c.metadata = metadata
	c.mu.Unlock()
	return metadata, nil
}

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