probo/e2e/internal/testutil/oauth2.go

// Copyright (c) 2025-2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package testutil

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"math/rand/v2"
	"net/http"
	"net/url"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

type (
	OAuth2TokenResponse struct {
		AccessToken  string `json:"access_token"`
		TokenType    string `json:"token_type"`
		ExpiresIn    int64  `json:"expires_in"`
		RefreshToken string `json:"refresh_token,omitempty"`
		IDToken      string `json:"id_token,omitempty"`
		Scope        string `json:"scope,omitempty"`
	}

	OAuth2ErrorResponse struct {
		Code        string `json:"error"`
		Description string `json:"error_description,omitempty"`
	}

	OAuth2RegisterResponse struct {
		ClientID                string   `json:"client_id"`
		ClientSecret            string   `json:"client_secret,omitempty"`
		ClientName              string   `json:"client_name"`
		Visibility              string   `json:"visibility"`
		RedirectURIs            []string `json:"redirect_uris"`
		GrantTypes              []string `json:"grant_types"`
		ResponseTypes           []string `json:"response_types"`
		TokenEndpointAuthMethod string   `json:"token_endpoint_auth_method"`
		Scopes                  string   `json:"scopes"`
	}

	OAuth2IntrospectResponse struct {
		Active    bool   `json:"active"`
		Scope     string `json:"scope,omitempty"`
		ClientID  string `json:"client_id,omitempty"`
		Sub       string `json:"sub,omitempty"`
		Exp       int64  `json:"exp,omitempty"`
		Iat       int64  `json:"iat,omitempty"`
		TokenType string `json:"token_type,omitempty"`
	}

	OAuth2DeviceAuthResponse struct {
		DeviceCode              string `json:"device_code"`
		UserCode                string `json:"user_code"`
		VerificationURI         string `json:"verification_uri"`
		VerificationURIComplete string `json:"verification_uri_complete"`
		ExpiresIn               int    `json:"expires_in"`
		Interval                int    `json:"interval"`
	}

	OAuth2DiscoveryResponse struct {
		Issuer                                    string   `json:"issuer"`
		AuthorizationEndpoint                     string   `json:"authorization_endpoint"`
		TokenEndpoint                             string   `json:"token_endpoint"`
		UserinfoEndpoint                          string   `json:"userinfo_endpoint"`
		JwksURI                                   string   `json:"jwks_uri"`
		RegistrationEndpoint                      string   `json:"registration_endpoint"`
		IntrospectionEndpoint                     string   `json:"introspection_endpoint"`
		RevocationEndpoint                        string   `json:"revocation_endpoint"`
		DeviceAuthorizationEndpoint               string   `json:"device_authorization_endpoint"`
		ScopesSupported                           []string `json:"scopes_supported"`
		ResponseTypesSupported                    []string `json:"response_types_supported"`
		GrantTypesSupported                       []string `json:"grant_types_supported"`
		TokenEndpointAuthMethodsSupported         []string `json:"token_endpoint_auth_methods_supported"`
		RevocationEndpointAuthMethodsSupported    []string `json:"revocation_endpoint_auth_methods_supported"`
		IntrospectionEndpointAuthMethodsSupported []string `json:"introspection_endpoint_auth_methods_supported"`
		SubjectTypesSupported                     []string `json:"subject_types_supported"`
		IDTokenSigningAlgValuesSupported          []string `json:"id_token_signing_alg_values_supported"`
		CodeChallengeMethodsSupported             []string `json:"code_challenge_methods_supported"`
		ClaimsSupported                           []string `json:"claims_supported"`
		ProtectedResources                        []string `json:"protected_resources,omitempty"`
		ClientIDMetadataDocumentSupported         bool     `json:"client_id_metadata_document_supported"`
	}

	OAuth2ProtectedResourceMetadataResponse struct {
		Resource               string   `json:"resource"`
		AuthorizationServers   []string `json:"authorization_servers"`
		BearerMethodsSupported []string `json:"bearer_methods_supported"`
		ScopesSupported        []string `json:"scopes_supported"`
	}

	OAuth2JWKSResponse struct {
		Keys []map[string]any `json:"keys"`
	}

	OAuth2UserInfoResponse struct {
		Sub           string `json:"sub"`
		Email         string `json:"email,omitempty"`
		EmailVerified bool   `json:"email_verified,omitempty"`
		Name          string `json:"name,omitempty"`
	}

	OAuth2HTTPResponse struct {
		StatusCode int
		Header     http.Header
		Body       []byte
	}
)

func oauth2BaseURL(c *Client) string {
	return c.BaseURL() + "/api/connect/v1/oauth2"
}

func postForm(
	httpClient *http.Client,
	url string,
	values url.Values,
) (*OAuth2HTTPResponse, error) {
	resp, err := httpClient.PostForm(url, values)
	if err != nil {
		return nil, fmt.Errorf("cannot post form: %w", err)
	}

	defer func() { _ = resp.Body.Close() }()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("cannot read response body: %w", err)
	}

	return &OAuth2HTTPResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: body}, nil
}

func postJSON(
	httpClient *http.Client,
	url string,
	payload any,
) (*OAuth2HTTPResponse, error) {
	data, err := json.Marshal(payload)
	if err != nil {
		return nil, fmt.Errorf("cannot marshal payload: %w", err)
	}

	resp, err := httpClient.Post(url, "application/json", bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("cannot post json: %w", err)
	}

	defer func() { _ = resp.Body.Close() }()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("cannot read response body: %w", err)
	}

	return &OAuth2HTTPResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: body}, nil
}

func getJSON(
	httpClient *http.Client,
	url string,
	headers map[string]string,
) (*OAuth2HTTPResponse, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, fmt.Errorf("cannot create request: %w", err)
	}

	for k, v := range headers {
		req.Header.Set(k, v)
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("cannot execute request: %w", err)
	}

	defer func() { _ = resp.Body.Close() }()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("cannot read response body: %w", err)
	}

	return &OAuth2HTTPResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: body}, nil
}

func postFormWithBasicAuth(
	httpClient *http.Client,
	rawURL string,
	values url.Values,
	username, password string,
) (*OAuth2HTTPResponse, error) {
	req, err := http.NewRequest(
		"POST",
		rawURL,
		strings.NewReader(values.Encode()),
	)
	if err != nil {
		return nil, fmt.Errorf("cannot create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth(username, password)

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("cannot execute request: %w", err)
	}

	defer func() { _ = resp.Body.Close() }()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("cannot read response body: %w", err)
	}

	return &OAuth2HTTPResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: body}, nil
}

// OAuth2Discovery fetches the OpenID Connect discovery document.
func OAuth2Discovery(c *Client) (*OAuth2DiscoveryResponse, *OAuth2HTTPResponse, error) {
	raw, err := getJSON(c.HTTPClient(), c.BaseURL()+"/.well-known/openid-configuration", nil)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2DiscoveryResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode discovery response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2JWKS fetches the JSON Web Key Set.
func OAuth2JWKS(c *Client) (*OAuth2JWKSResponse, *OAuth2HTTPResponse, error) {
	raw, err := getJSON(c.HTTPClient(), oauth2BaseURL(c)+"/jwks", nil)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2JWKSResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode jwks response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2ProtectedResourceMetadata fetches the RFC 9728 protected resource
// metadata document.
func OAuth2ProtectedResourceMetadata(
	c *Client,
) (*OAuth2ProtectedResourceMetadataResponse, *OAuth2HTTPResponse, error) {
	raw, err := getJSON(c.HTTPClient(), c.BaseURL()+"/.well-known/oauth-protected-resource", nil)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2ProtectedResourceMetadataResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode protected resource metadata: %w", err)
	}

	return &result, raw, nil
}

// OAuth2RegisterClient registers a new OAuth2 client via dynamic registration.
func OAuth2RegisterClient(
	c *Client,
	input map[string]any,
) (*OAuth2RegisterResponse, *OAuth2HTTPResponse, error) {
	raw, err := postJSON(c.HTTPClient(), oauth2BaseURL(c)+"/register", input)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusCreated {
		return nil, raw, nil
	}

	var result OAuth2RegisterResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode register response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2Authorize performs a GET to the authorize endpoint and returns the
// raw HTTP response without following redirects.
func OAuth2Authorize(
	c *Client,
	params url.Values,
) (*OAuth2HTTPResponse, error) {
	noRedirectClient := &http.Client{
		Jar:     c.HTTPClient().Jar,
		Timeout: c.HTTPClient().Timeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	reqURL := oauth2BaseURL(c) + "/authorize?" + params.Encode()

	resp, err := noRedirectClient.Get(reqURL)
	if err != nil {
		return nil, fmt.Errorf("cannot get authorize: %w", err)
	}

	defer func() { _ = resp.Body.Close() }()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("cannot read response body: %w", err)
	}

	return &OAuth2HTTPResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: body}, nil
}

// OAuth2AuthorizeCodeFromRedirect extracts the authorization code from the
// Location header of a 302 response.
func OAuth2AuthorizeCodeFromRedirect(resp *OAuth2HTTPResponse) (string, error) {
	loc := resp.Header.Get("Location")
	if loc == "" {
		return "", fmt.Errorf("no Location header in redirect response (status=%d body=%s)", resp.StatusCode, string(resp.Body))
	}

	u, err := url.Parse(loc)
	if err != nil {
		return "", fmt.Errorf("cannot parse redirect url: %w", err)
	}

	code := u.Query().Get("code")
	if code == "" {
		return "", fmt.Errorf("no code in redirect url: %s", loc)
	}

	return code, nil
}

// OAuth2ConsentApprove approves an OAuth2 consent via the GraphQL mutation.
// It returns a simulated HTTP 302 response with the redirect URL in the
// Location header so existing callers can extract the authorization code.
func OAuth2ConsentApprove(c *Client, consentID string) (*OAuth2HTTPResponse, error) {
	return oauth2ConsentDecide(c, consentID, true)
}

// OAuth2ConsentDeny denies an OAuth2 consent via the GraphQL mutation.
// It returns a simulated HTTP 302 response with the redirect URL in the
// Location header so existing callers can inspect the error parameters.
func OAuth2ConsentDeny(c *Client, consentID string) (*OAuth2HTTPResponse, error) {
	return oauth2ConsentDecide(c, consentID, false)
}

func oauth2ConsentDecide(c *Client, consentID string, approved bool) (*OAuth2HTTPResponse, error) {
	const query = `
		mutation ApproveConsent($input: ApproveConsentInput!) {
			approveConsent(input: $input) {
				redirectURL
				deviceAuthorized
			}
		}
	`

	var result struct {
		ApproveConsent struct {
			RedirectURL      *string `json:"redirectURL"`
			DeviceAuthorized *bool   `json:"deviceAuthorized"`
		} `json:"approveConsent"`
	}

	err := c.ExecuteConnect(
		query,
		map[string]any{
			"input": map[string]any{
				"consentId": consentID,
				"approved":  approved,
			},
		},
		&result,
	)
	if err != nil {
		return &OAuth2HTTPResponse{
			StatusCode: http.StatusBadRequest,
			Header:     http.Header{},
			Body:       []byte(err.Error()),
		}, nil
	}

	resp := &OAuth2HTTPResponse{
		StatusCode: http.StatusFound,
		Header:     http.Header{},
	}

	if result.ApproveConsent.RedirectURL != nil {
		resp.Header.Set("Location", *result.ApproveConsent.RedirectURL)
	}

	return resp, nil
}

// OAuth2TokenWithCode exchanges an authorization code for tokens.
func OAuth2TokenWithCode(
	c *Client,
	clientID, clientSecret, code, redirectURI, codeVerifier string,
) (*OAuth2TokenResponse, *OAuth2HTTPResponse, error) {
	values := url.Values{
		"grant_type":   {"authorization_code"},
		"code":         {code},
		"redirect_uri": {redirectURI},
	}

	if codeVerifier != "" {
		values.Set("code_verifier", codeVerifier)
	}

	raw, err := postFormWithBasicAuth(
		c.HTTPClient(),
		oauth2BaseURL(c)+"/token",
		values,
		clientID,
		clientSecret,
	)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2TokenResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode token response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2TokenWithCodePostAuth exchanges an authorization code for tokens
// using client_secret_post authentication (credentials in POST body).
func OAuth2TokenWithCodePostAuth(
	c *Client,
	clientID, clientSecret, code, redirectURI, codeVerifier string,
) (*OAuth2TokenResponse, *OAuth2HTTPResponse, error) {
	values := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"redirect_uri":  {redirectURI},
		"client_id":     {clientID},
		"client_secret": {clientSecret},
	}

	if codeVerifier != "" {
		values.Set("code_verifier", codeVerifier)
	}

	raw, err := postForm(c.HTTPClient(), oauth2BaseURL(c)+"/token", values)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2TokenResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode token response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2TokenWithRefreshToken refreshes tokens using a refresh token.
func OAuth2TokenWithRefreshToken(
	c *Client,
	clientID, clientSecret, refreshToken string,
) (*OAuth2TokenResponse, *OAuth2HTTPResponse, error) {
	values := url.Values{
		"grant_type":    {"refresh_token"},
		"refresh_token": {refreshToken},
	}

	raw, err := postFormWithBasicAuth(
		c.HTTPClient(),
		oauth2BaseURL(c)+"/token",
		values,
		clientID,
		clientSecret,
	)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2TokenResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode token response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2TokenWithDeviceCode polls the token endpoint for device code grant.
func OAuth2TokenWithDeviceCode(
	c *Client,
	clientID, deviceCode string,
) (*OAuth2TokenResponse, *OAuth2ErrorResponse, *OAuth2HTTPResponse, error) {
	values := url.Values{
		"grant_type":  {"urn:ietf:params:oauth:grant-type:device_code"},
		"client_id":   {clientID},
		"device_code": {deviceCode},
	}

	raw, err := postForm(c.HTTPClient(), oauth2BaseURL(c)+"/token", values)
	if err != nil {
		return nil, nil, nil, err
	}

	if raw.StatusCode == http.StatusOK {
		var result OAuth2TokenResponse
		if err := json.Unmarshal(raw.Body, &result); err != nil {
			return nil, nil, raw, fmt.Errorf("cannot decode token response: %w", err)
		}

		return &result, nil, raw, nil
	}

	var errResp OAuth2ErrorResponse
	if err := json.Unmarshal(raw.Body, &errResp); err != nil {
		return nil, nil, raw, nil
	}

	return nil, &errResp, raw, nil
}

// OAuth2TokenRaw posts arbitrary form values to the token endpoint.
func OAuth2TokenRaw(
	c *Client,
	values url.Values,
) (*OAuth2HTTPResponse, error) {
	raw, err := postForm(c.HTTPClient(), oauth2BaseURL(c)+"/token", values)
	if err != nil {
		return nil, err
	}

	return raw, nil
}

// OAuth2TokenRawWithBasicAuth posts form values to the token endpoint with
// HTTP Basic authentication.
func OAuth2TokenRawWithBasicAuth(
	c *Client,
	values url.Values,
	username, password string,
) (*OAuth2HTTPResponse, error) {
	raw, err := postFormWithBasicAuth(
		c.HTTPClient(),
		oauth2BaseURL(c)+"/token",
		values,
		username,
		password,
	)
	if err != nil {
		return nil, err
	}

	return raw, nil
}

// OAuth2DeviceAuth starts the device authorization flow.
func OAuth2DeviceAuth(
	c *Client,
	clientID, scope string,
) (*OAuth2DeviceAuthResponse, *OAuth2HTTPResponse, error) {
	values := url.Values{
		"client_id": {clientID},
	}

	if scope != "" {
		values.Set("scope", scope)
	}

	raw, err := postForm(c.HTTPClient(), oauth2BaseURL(c)+"/device", values)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2DeviceAuthResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode device auth response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2DeviceVerify authorizes a device code via the GraphQL authorizeDevice
// mutation. It performs the full consent flow: submitting the user code, and if
// consent is required, approving it via approveOAuth2Consent.
func OAuth2DeviceVerify(c *Client, userCode string) (*OAuth2HTTPResponse, error) {
	const authorizeQuery = `
		mutation AuthorizeDevice($input: AuthorizeDeviceInput!) {
			authorizeDevice(input: $input) {
				success
				consentId
			}
		}
	`

	var authorizeResult struct {
		AuthorizeDevice struct {
			Success   bool    `json:"success"`
			ConsentID *string `json:"consentId"`
		} `json:"authorizeDevice"`
	}

	err := c.ExecuteConnect(
		authorizeQuery,
		map[string]any{
			"input": map[string]any{"userCode": userCode},
		},
		&authorizeResult,
	)
	if err != nil {
		body, _ := json.Marshal(map[string]string{"error": err.Error()})
		return &OAuth2HTTPResponse{StatusCode: http.StatusOK, Body: body}, nil
	}

	if authorizeResult.AuthorizeDevice.Success {
		return &OAuth2HTTPResponse{StatusCode: http.StatusOK, Body: []byte(`{"success":true}`)}, nil
	}

	consentID := authorizeResult.AuthorizeDevice.ConsentID
	if consentID == nil {
		body, _ := json.Marshal(map[string]string{"error": "unexpected response"})
		return &OAuth2HTTPResponse{StatusCode: http.StatusInternalServerError, Body: body}, nil
	}

	const approveQuery = `
		mutation ApproveConsent($input: ApproveConsentInput!) {
			approveConsent(input: $input) {
				deviceAuthorized
			}
		}
	`

	var approveResult struct {
		ApproveConsent struct {
			DeviceAuthorized *bool `json:"deviceAuthorized"`
		} `json:"approveConsent"`
	}

	err = c.ExecuteConnect(
		approveQuery,
		map[string]any{
			"input": map[string]any{"consentId": *consentID, "approved": true},
		},
		&approveResult,
	)
	if err != nil {
		body, _ := json.Marshal(map[string]string{"error": err.Error()})
		return &OAuth2HTTPResponse{StatusCode: http.StatusOK, Body: body}, nil
	}

	return &OAuth2HTTPResponse{StatusCode: http.StatusOK, Body: []byte(`{"success":true}`)}, nil
}

// OAuth2UserInfo fetches the UserInfo endpoint with a Bearer token.
func OAuth2UserInfo(
	c *Client,
	accessToken string,
) (*OAuth2UserInfoResponse, *OAuth2HTTPResponse, error) {
	headers := map[string]string{}
	if accessToken != "" {
		headers["Authorization"] = "Bearer " + accessToken
	}

	raw, err := getJSON(c.HTTPClient(), oauth2BaseURL(c)+"/userinfo", headers)
	if err != nil {
		return nil, nil, err
	}

	if raw.StatusCode != http.StatusOK {
		return nil, raw, nil
	}

	var result OAuth2UserInfoResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode userinfo response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2UserInfoRaw fetches the UserInfo endpoint with custom query params
// and no Authorization header (for testing that query/body tokens are rejected).
func OAuth2UserInfoRaw(
	c *Client,
	queryParams url.Values,
) (*OAuth2HTTPResponse, error) {
	reqURL := oauth2BaseURL(c) + "/userinfo"
	if len(queryParams) > 0 {
		reqURL += "?" + queryParams.Encode()
	}

	raw, err := getJSON(c.HTTPClient(), reqURL, nil)
	if err != nil {
		return nil, err
	}

	return raw, nil
}

// OAuth2Introspect introspects a token using client credentials.
func OAuth2Introspect(
	c *Client,
	clientID, clientSecret, token string,
) (*OAuth2IntrospectResponse, *OAuth2HTTPResponse, error) {
	return OAuth2IntrospectWithHint(c, clientID, clientSecret, token, "")
}

// OAuth2IntrospectWithHint introspects a token with an optional
// token_type_hint per RFC 7662.
func OAuth2IntrospectWithHint(
	c *Client,
	clientID, clientSecret, token, tokenTypeHint string,
) (*OAuth2IntrospectResponse, *OAuth2HTTPResponse, error) {
	values := url.Values{
		"token": {token},
	}

	if tokenTypeHint != "" {
		values.Set("token_type_hint", tokenTypeHint)
	}

	raw, err := postFormWithBasicAuth(
		c.HTTPClient(),
		oauth2BaseURL(c)+"/introspect",
		values,
		clientID,
		clientSecret,
	)
	if err != nil {
		return nil, nil, err
	}

	var result OAuth2IntrospectResponse
	if err := json.Unmarshal(raw.Body, &result); err != nil {
		return nil, raw, fmt.Errorf("cannot decode introspect response: %w", err)
	}

	return &result, raw, nil
}

// OAuth2Revoke revokes a token using client credentials.
func OAuth2Revoke(
	c *Client,
	clientID, clientSecret, token string,
) (*OAuth2HTTPResponse, error) {
	return OAuth2RevokeWithHint(c, clientID, clientSecret, token, "")
}

// OAuth2RevokeWithHint revokes a token with an optional token_type_hint.
func OAuth2RevokeWithHint(
	c *Client,
	clientID, clientSecret, token, tokenTypeHint string,
) (*OAuth2HTTPResponse, error) {
	values := url.Values{
		"token": {token},
	}

	if tokenTypeHint != "" {
		values.Set("token_type_hint", tokenTypeHint)
	}

	raw, err := postFormWithBasicAuth(
		c.HTTPClient(),
		oauth2BaseURL(c)+"/revoke",
		values,
		clientID,
		clientSecret,
	)
	if err != nil {
		return nil, err
	}

	return raw, nil
}

// PKCE helpers

// GeneratePKCE generates a code_verifier and code_challenge (S256) pair.
func GeneratePKCE() (verifier, challenge string) {
	const charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~"

	b := make([]byte, 64)
	for i := range b {
		b[i] = charset[rand.IntN(len(charset))]
	}

	verifier = string(b)

	h := sha256.Sum256([]byte(verifier))
	challenge = base64.RawURLEncoding.EncodeToString(h[:])

	return verifier, challenge
}

// IsConsentRedirect returns true when the authorize endpoint responded with
// a 302 redirect to the consent page (as opposed to a redirect carrying an
// authorization code).
func IsConsentRedirect(resp *OAuth2HTTPResponse) bool {
	if resp.StatusCode != http.StatusFound {
		return false
	}

	loc := resp.Header.Get("Location")

	u, err := url.Parse(loc)
	if err != nil {
		return false
	}

	return u.Query().Get("consent_id") != ""
}

// ExtractConsentIDFromResponse extracts the consent_id from an authorize
// response.  It handles the current redirect-based flow (302 to consent page)
// as well as the legacy inline HTML flow (200 with hidden form field).
func ExtractConsentIDFromResponse(resp *OAuth2HTTPResponse) (string, error) {
	if resp.StatusCode == http.StatusFound {
		loc := resp.Header.Get("Location")
		if loc == "" {
			return "", fmt.Errorf("no Location header in redirect response")
		}

		u, err := url.Parse(loc)
		if err != nil {
			return "", fmt.Errorf("cannot parse redirect url: %w", err)
		}

		consentID := u.Query().Get("consent_id")
		if consentID == "" {
			return "", fmt.Errorf("no consent_id in redirect url: %s", loc)
		}

		return consentID, nil
	}

	return ExtractConsentID(resp.Body)
}

// ExtractConsentID extracts the consent_id from a consent HTML page.
func ExtractConsentID(body []byte) (string, error) {
	s := string(body)

	needle := `name="consent_id" value="`

	idx := strings.Index(s, needle)
	if idx == -1 {
		return "", fmt.Errorf("consent_id not found in page")
	}

	start := idx + len(needle)

	end := strings.Index(s[start:], `"`)
	if end == -1 {
		return "", fmt.Errorf("malformed consent_id value")
	}

	return s[start : start+end], nil
}

// OAuth2PerformAuthorizationCodeFlow performs the full authorization code flow
// and returns the token response. This is a convenience function for tests that
// need tokens but are not testing the authorization flow itself.
func OAuth2PerformAuthorizationCodeFlow(
	t testing.TB,
	c *Client,
	clientID, clientSecret, redirectURI string,
) *OAuth2TokenResponse {
	t.Helper()

	return OAuth2PerformAuthorizationCodeFlowWithScopes(
		t,
		c,
		clientID,
		clientSecret,
		redirectURI,
		"openid email profile offline_access",
	)
}

// OAuth2PerformAuthorizationCodeFlowWithScopes performs the authorization code
// flow with the requested OAuth2 scopes.
func OAuth2PerformAuthorizationCodeFlowWithScopes(
	t testing.TB,
	c *Client,
	clientID, clientSecret, redirectURI, scopes string,
) *OAuth2TokenResponse {
	t.Helper()

	verifier, challenge := GeneratePKCE()

	params := url.Values{
		"client_id":             {clientID},
		"redirect_uri":          {redirectURI},
		"response_type":         {"code"},
		"scope":                 {scopes},
		"state":                 {"test-state"},
		"code_challenge":        {challenge},
		"code_challenge_method": {"S256"},
	}

	authResp, err := OAuth2Authorize(c, params)
	require.NoError(t, err)

	var code string

	if IsConsentRedirect(authResp) {
		consentID, err := ExtractConsentIDFromResponse(authResp)
		require.NoError(t, err)

		consentResp, err := OAuth2ConsentApprove(c, consentID)
		require.NoError(t, err)
		require.Equal(t, http.StatusFound, consentResp.StatusCode)

		code, err = OAuth2AuthorizeCodeFromRedirect(consentResp)
		require.NoError(t, err)
	} else {
		require.Equal(t, http.StatusFound, authResp.StatusCode)
		code, err = OAuth2AuthorizeCodeFromRedirect(authResp)
		require.NoError(t, err)
	}

	tokenResp, raw, err := OAuth2TokenWithCode(
		c,
		clientID,
		clientSecret,
		code,
		redirectURI,
		verifier,
	)
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, raw.StatusCode, "token exchange failed: %s", string(raw.Body))
	require.NotNil(t, tokenResp)

	return tokenResp
}