fraiseql-auth 2.2.0

//! OAuth2 and OIDC client implementations.

use std::{sync::Arc, time::Duration as StdDuration};

/// Timeout for all outbound OAuth2 / OIDC HTTP requests.
const OAUTH_REQUEST_TIMEOUT: StdDuration = StdDuration::from_secs(30);

use std::fmt::Write as _;

use serde::{Deserialize, Serialize};

use super::{
    super::jwks::{JwksCache, JwksError},
    pkce::PKCEChallenge,
    types::{IdTokenClaims, TokenResponse, UserInfo},
};

/// OIDC provider configuration
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OIDCProviderConfig {
    /// Provider issuer URL
    pub issuer:                   String,
    /// Authorization endpoint
    pub authorization_endpoint:   String,
    /// Token endpoint
    pub token_endpoint:           String,
    /// Userinfo endpoint
    pub userinfo_endpoint:        Option<String>,
    /// JWKS URI for public keys
    pub jwks_uri:                 String,
    /// Scopes supported by provider
    pub scopes_supported:         Vec<String>,
    /// Response types supported
    pub response_types_supported: Vec<String>,
}

impl OIDCProviderConfig {
    /// Create new provider configuration
    pub fn new(
        issuer: String,
        authorization_endpoint: String,
        token_endpoint: String,
        jwks_uri: String,
    ) -> Self {
        Self {
            issuer,
            authorization_endpoint,
            token_endpoint,
            userinfo_endpoint: None,
            jwks_uri,
            scopes_supported: vec![
                "openid".to_string(),
                "profile".to_string(),
                "email".to_string(),
            ],
            response_types_supported: vec!["code".to_string()],
        }
    }
}

/// Result of [`OAuth2Client::authorization_url`].
///
/// The caller MUST store `state` (for CSRF verification at callback), when
/// present the PKCE `pkce.code_verifier` (for token exchange), and when
/// present the `nonce` value (must be verified against the ID token at
/// callback via [`OIDCClient::verify_id_token`]).
#[derive(Debug, Clone)]
pub struct AuthorizationRequest {
    /// The full authorization URL to redirect the user to.
    pub url:   String,
    /// CSRF state value — verify this matches the `state` query param at callback.
    pub state: String,
    /// PKCE challenge, present only when `use_pkce = true`.
    pub pkce:  Option<PKCEChallenge>,
    /// OIDC nonce for replay protection.
    ///
    /// Present only when the authorization URL was generated by
    /// [`OIDCClient::authorization_url`].  The caller must store this value
    /// and pass `Some(&nonce.nonce)` to [`OIDCClient::verify_id_token`] at
    /// callback time.
    pub nonce: Option<super::pkce::NonceParameter>,
}

/// OAuth2 client for authorization code flow.
#[derive(Debug, Clone)]
pub struct OAuth2Client {
    /// Client ID from provider.
    pub client_id:              String,
    /// Client secret from provider.
    client_secret:              String,
    /// Authorization endpoint.
    pub authorization_endpoint: String,
    /// Token endpoint.
    token_endpoint:             String,
    /// Scopes to request.
    pub scopes:                 Vec<String>,
    /// Use PKCE for additional security.
    pub use_pkce:               bool,
    /// HTTP client for token requests.
    http_client:                reqwest::Client,
}

impl OAuth2Client {
    /// Maximum byte size accepted from an OAuth token endpoint response.
    ///
    /// A well-formed token response (access_token, id_token, refresh_token) is a
    /// few kilobytes at most.  1 MiB prevents a malicious provider from sending a
    /// response large enough to exhaust server memory.
    const MAX_OAUTH_RESPONSE_BYTES: usize = 1024 * 1024;

    /// Create new OAuth2 client.
    pub fn new(
        client_id: impl Into<String>,
        client_secret: impl Into<String>,
        authorization_endpoint: impl Into<String>,
        token_endpoint: impl Into<String>,
    ) -> Self {
        Self {
            client_id:              client_id.into(),
            client_secret:          client_secret.into(),
            authorization_endpoint: authorization_endpoint.into(),
            token_endpoint:         token_endpoint.into(),
            scopes:                 vec![
                "openid".to_string(),
                "profile".to_string(),
                "email".to_string(),
            ],
            use_pkce:               false,
            http_client:            reqwest::Client::builder()
                .timeout(OAUTH_REQUEST_TIMEOUT)
                .build()
                .unwrap_or_default(),
        }
    }

    /// Set scopes for request.
    pub fn with_scopes(mut self, scopes: Vec<String>) -> Self {
        self.scopes = scopes;
        self
    }

    /// Enable PKCE protection.
    pub const fn with_pkce(mut self, enabled: bool) -> Self {
        self.use_pkce = enabled;
        self
    }

    /// Generate authorization URL.
    ///
    /// Returns an [`AuthorizationRequest`] containing the URL, the CSRF state
    /// value (must be stored and verified at callback), and an optional PKCE
    /// challenge (when `use_pkce = true`; the `code_verifier` must be stored
    /// and sent during token exchange).
    pub fn authorization_url(&self, redirect_uri: &str) -> AuthorizationRequest {
        let state = uuid::Uuid::new_v4().to_string();
        let scope = self.scopes.join(" ");

        let mut url = format!(
            "{}?client_id={}&redirect_uri={}&response_type=code&scope={}&state={}",
            self.authorization_endpoint,
            urlencoding::encode(&self.client_id),
            urlencoding::encode(redirect_uri),
            urlencoding::encode(&scope),
            urlencoding::encode(&state),
        );

        let pkce = if self.use_pkce {
            let challenge = PKCEChallenge::new();
            let _ = write!(
                url,
                "&code_challenge={}&code_challenge_method=S256",
                urlencoding::encode(&challenge.code_challenge),
            );
            Some(challenge)
        } else {
            None
        };

        AuthorizationRequest {
            url,
            state,
            pkce,
            nonce: None,
        }
    }

    // 1 MiB

    /// Post a form request to the token endpoint and parse the response.
    async fn post_token_request(&self, params: &[(&str, &str)]) -> Result<TokenResponse, String> {
        let response = self
            .http_client
            .post(&self.token_endpoint)
            .form(params)
            .send()
            .await
            .map_err(|e| format!("Token request failed: {e}"))?;

        // Read the entire body once so we can apply a size cap regardless of
        // whether the response is a success or an error.
        let status = response.status();
        let body_bytes = response
            .bytes()
            .await
            .map_err(|e| format!("Failed to read token response body: {e}"))?;

        if !status.is_success() {
            let capped = &body_bytes[..body_bytes.len().min(Self::MAX_OAUTH_RESPONSE_BYTES)];
            let body = String::from_utf8_lossy(capped);
            return Err(format!("Token endpoint returned error: {body}"));
        }

        if body_bytes.len() > Self::MAX_OAUTH_RESPONSE_BYTES {
            return Err(format!(
                "Token response body too large ({} bytes, max {})",
                body_bytes.len(),
                Self::MAX_OAUTH_RESPONSE_BYTES
            ));
        }

        serde_json::from_slice::<TokenResponse>(&body_bytes)
            .map_err(|e| format!("Failed to parse token response: {e}"))
    }

    /// Exchange authorization code for tokens.
    ///
    /// # Errors
    ///
    /// Returns an error if the HTTP request to the token endpoint fails or the response
    /// cannot be parsed as a `TokenResponse`.
    pub async fn exchange_code(
        &self,
        code: &str,
        redirect_uri: &str,
    ) -> Result<TokenResponse, String> {
        let params = [
            ("grant_type", "authorization_code"),
            ("code", code),
            ("client_id", self.client_id.as_str()),
            ("client_secret", self.client_secret.as_str()),
            ("redirect_uri", redirect_uri),
        ];
        self.post_token_request(&params).await
    }

    /// Refresh access token using a refresh token.
    ///
    /// # Errors
    ///
    /// Propagates errors from the token endpoint request (network failure,
    /// non-2xx HTTP status, oversized response body, or JSON parse error).
    pub async fn refresh_token(&self, refresh_token: &str) -> Result<TokenResponse, String> {
        let params = [
            ("grant_type", "refresh_token"),
            ("refresh_token", refresh_token),
            ("client_id", self.client_id.as_str()),
            ("client_secret", self.client_secret.as_str()),
        ];
        self.post_token_request(&params).await
    }
}

/// OIDC client for OpenID Connect flow.
#[derive(Debug)]
pub struct OIDCClient {
    /// Provider configuration.
    pub config:     OIDCProviderConfig,
    /// Client ID.
    pub client_id:  String,
    /// Client secret — retained for token revocation and introspection endpoints.
    #[allow(dead_code)] // Reason: retained for token revocation and introspection endpoints
    client_secret: String,
    /// JWKS key cache for ID token signature verification.
    pub jwks_cache: Arc<JwksCache>,
    /// HTTP client for userinfo requests.
    http_client:    reqwest::Client,
}

impl OIDCClient {
    /// Maximum byte size for a userinfo endpoint response.
    ///
    /// Userinfo payloads carry a small set of JWT-derived claims.
    /// 1 `MiB` is generous while blocking allocation-bomb responses.
    const MAX_USERINFO_RESPONSE_BYTES: usize = 1024 * 1024;

    // 1 MiB

    /// Create new OIDC client with JWKS caching.
    ///
    /// The JWKS cache TTL defaults to 1 hour.
    ///
    /// # Errors
    ///
    /// Returns [`JwksError`] if `config.jwks_uri` is not a valid HTTPS URL
    /// (HTTP is allowed only for localhost).
    pub fn new(
        config: OIDCProviderConfig,
        client_id: impl Into<String>,
        client_secret: impl Into<String>,
    ) -> Result<Self, JwksError> {
        let jwks_cache = Arc::new(JwksCache::new(&config.jwks_uri, StdDuration::from_secs(3600))?);
        Ok(Self {
            config,
            client_id: client_id.into(),
            client_secret: client_secret.into(),
            jwks_cache,
            http_client: reqwest::Client::builder()
                .timeout(OAUTH_REQUEST_TIMEOUT)
                .build()
                .unwrap_or_default(),
        })
    }

    /// Create OIDC client with a pre-built JWKS cache (for testing).
    pub fn with_jwks_cache(
        config: OIDCProviderConfig,
        client_id: impl Into<String>,
        client_secret: impl Into<String>,
        jwks_cache: Arc<JwksCache>,
    ) -> Self {
        Self {
            config,
            client_id: client_id.into(),
            client_secret: client_secret.into(),
            jwks_cache,
            http_client: reqwest::Client::builder()
                .timeout(OAUTH_REQUEST_TIMEOUT)
                .build()
                .unwrap_or_default(),
        }
    }

    /// Generate an OIDC authorization URL with a fresh nonce for replay protection.
    ///
    /// This extends the standard OAuth2 flow by appending a `nonce` parameter to
    /// the authorization URL. The returned [`AuthorizationRequest::nonce`] **must**
    /// be stored (e.g. in the encrypted session state) and passed to
    /// [`verify_id_token`](Self::verify_id_token) at callback time.
    ///
    /// PKCE is always enabled for OIDC flows started via this method.
    pub fn authorization_url(&self, redirect_uri: &str) -> AuthorizationRequest {
        let state = uuid::Uuid::new_v4().to_string();
        let scope = self.config.scopes_supported.join(" ");
        let nonce = super::pkce::NonceParameter::new();
        let challenge = PKCEChallenge::new();

        let url = format!(
            "{}?client_id={}&redirect_uri={}&response_type=code&scope={}&state={}\
             &nonce={}&code_challenge={}&code_challenge_method=S256",
            self.config.authorization_endpoint,
            urlencoding::encode(&self.client_id),
            urlencoding::encode(redirect_uri),
            urlencoding::encode(&scope),
            urlencoding::encode(&state),
            urlencoding::encode(&nonce.nonce),
            urlencoding::encode(&challenge.code_challenge),
        );

        AuthorizationRequest {
            url,
            state,
            pkce: Some(challenge),
            nonce: Some(nonce),
        }
    }

    /// Verify an ID token's JWT signature and claims.
    ///
    /// Decodes the JWT header to extract the `kid`, fetches the matching public
    /// key from the JWKS cache, then validates signature, issuer, audience, and
    /// required claims.
    ///
    /// **Nonce**: when `expected_nonce` is `Some`, the token's `nonce` claim must
    /// match exactly.  When it is `None` but the token *contains* a `nonce` claim,
    /// validation still succeeds — callers that generated the authorization URL
    /// via [`authorization_url`](Self::authorization_url) MUST pass the stored
    /// nonce here.
    ///
    /// **`max_age`**: when `max_age_secs` is `Some`, the token's `auth_time` claim
    /// is required and must be within `max_age_secs` seconds of the current time.
    /// This prevents accepting tokens from sessions that were authenticated too
    /// long ago (RFC 6749 §3.1.2.1 / OIDC Core §3.1.2.1).
    ///
    /// # Errors
    ///
    /// Returns an error if the token is malformed, the signature is invalid,
    /// claims validation fails, the nonce doesn't match, or the `auth_time` /
    /// `max_age` constraint is violated.
    pub async fn verify_id_token(
        &self,
        id_token: &str,
        expected_nonce: Option<&str>,
        max_age_secs: Option<u64>,
    ) -> Result<IdTokenClaims, String> {
        // 1. Decode header to get kid
        let header = jsonwebtoken::decode_header(id_token)
            .map_err(|e| format!("Invalid JWT header: {e}"))?;
        let kid = header.kid.ok_or("JWT missing 'kid' in header")?;

        // 2. Get key from JWKS cache
        let key = self
            .jwks_cache
            .get_key(&kid)
            .await
            .map_err(|e| format!("JWKS fetch error: {e}"))?
            .ok_or_else(|| format!("No key found for kid '{kid}'"))?;

        // 3. Build validation criteria
        let mut validation = jsonwebtoken::Validation::new(header.alg);
        validation.set_issuer(&[&self.config.issuer]);
        validation.set_audience(&[&self.client_id]);
        validation.set_required_spec_claims(&["exp", "iat", "iss", "aud", "sub"]);

        // 4. Decode and validate
        let token_data = jsonwebtoken::decode::<IdTokenClaims>(id_token, &key, &validation)
            .map_err(|e| format!("ID token validation failed: {e}"))?;
        let claims = token_data.claims;

        // 5. Verify nonce using constant-time comparison (replay protection — RFC 6749 §10.12 /
        //    OIDC Core §3.1.3.7).
        if let Some(expected) = expected_nonce {
            super::claims_validator::validate_nonce_claim(&claims, expected)
                .map_err(|e| e.to_string())?;
        }

        // 6. Validate auth_time against max_age (OIDC Core §3.1.2.1).
        if let Some(max_age) = max_age_secs {
            let now_secs = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map_or(i64::MAX, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX));
            super::claims_validator::validate_auth_time_claim(&claims, max_age, now_secs)
                .map_err(|e| e.to_string())?;
        }

        Ok(claims)
    }

    /// Fetch user information from the provider's userinfo endpoint.
    ///
    /// # Errors
    ///
    /// Returns an error if no userinfo endpoint is configured, the HTTP request
    /// fails, or the response cannot be parsed.
    pub async fn get_userinfo(&self, access_token: &str) -> Result<UserInfo, String> {
        let endpoint = self
            .config
            .userinfo_endpoint
            .as_ref()
            .ok_or("No userinfo endpoint configured for this provider")?;

        let response = self
            .http_client
            .get(endpoint)
            .bearer_auth(access_token)
            .send()
            .await
            .map_err(|e| format!("Userinfo request failed: {e}"))?;

        if !response.status().is_success() {
            return Err(format!("Userinfo endpoint returned {}", response.status()));
        }

        let body = response
            .bytes()
            .await
            .map_err(|e| format!("Failed to read userinfo response: {e}"))?;
        if body.len() > Self::MAX_USERINFO_RESPONSE_BYTES {
            return Err(format!(
                "Userinfo response too large ({} bytes, max {})",
                body.len(),
                Self::MAX_USERINFO_RESPONSE_BYTES
            ));
        }
        serde_json::from_slice::<UserInfo>(&body)
            .map_err(|e| format!("Failed to parse userinfo response: {e}"))
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used)] // Reason: test code, panics acceptable
    #![allow(missing_docs)] // Reason: test helpers

    use super::*;

    #[test]
    fn oauth_response_cap_constant_is_reasonable() {
        assert_eq!(OAuth2Client::MAX_OAUTH_RESPONSE_BYTES, 1024 * 1024);
    }

    #[test]
    fn oauth_response_error_body_is_capped() {
        // Simulate what post_token_request does with an oversized error body.
        let cap = OAuth2Client::MAX_OAUTH_RESPONSE_BYTES;
        let oversized: Vec<u8> = vec![b'e'; cap + 1_000];
        let capped = &oversized[..oversized.len().min(cap)];
        let text = String::from_utf8_lossy(capped).into_owned();
        assert_eq!(text.len(), cap, "body must be capped at MAX_OAUTH_RESPONSE_BYTES");
    }

    // ── S25-H1: OAuth2/OIDC client timeout ────────────────────────────────────

    #[test]
    fn oauth_request_timeout_is_set() {
        let secs = OAUTH_REQUEST_TIMEOUT.as_secs();
        assert!(secs > 0 && secs <= 120, "OAuth timeout should be 1–120 s, got {secs}");
    }

    #[test]
    fn oauth2_client_new_creates_instance() {
        let client = OAuth2Client::new(
            "client_id",
            "client_secret",
            "https://example.com/auth",
            "https://example.com/token",
        );
        assert_eq!(client.client_id, "client_id");
    }

    #[test]
    fn oidc_client_new_creates_instance() {
        let config = OIDCProviderConfig {
            issuer:                   "https://example.com".to_string(),
            authorization_endpoint:   "https://example.com/auth".to_string(),
            token_endpoint:           "https://example.com/token".to_string(),
            userinfo_endpoint:        None,
            jwks_uri:                 "https://example.com/.well-known/jwks.json".to_string(),
            scopes_supported:         vec!["openid".to_string()],
            response_types_supported: vec!["code".to_string()],
        };
        let client = OIDCClient::new(config, "client_id", "client_secret").unwrap();
        assert_eq!(client.client_id, "client_id");
    }

    // ── S26: OIDCClient userinfo response size cap ────────────────────────────

    #[test]
    fn oidc_userinfo_cap_constant_is_reasonable() {
        const { assert!(OIDCClient::MAX_USERINFO_RESPONSE_BYTES >= 64 * 1024) }
        const { assert!(OIDCClient::MAX_USERINFO_RESPONSE_BYTES <= 100 * 1024 * 1024) }
    }

    #[tokio::test]
    async fn oidc_userinfo_oversized_response_is_rejected() {
        use wiremock::{
            Mock, MockServer, ResponseTemplate,
            matchers::{method, path},
        };

        let mock_server = MockServer::start().await;
        let oversized = vec![b'x'; OIDCClient::MAX_USERINFO_RESPONSE_BYTES + 1];
        Mock::given(method("GET"))
            .and(path("/userinfo"))
            .respond_with(ResponseTemplate::new(200).set_body_bytes(oversized))
            .mount(&mock_server)
            .await;

        let config = OIDCProviderConfig {
            issuer:                   mock_server.uri(),
            authorization_endpoint:   format!("{}/auth", mock_server.uri()),
            token_endpoint:           format!("{}/token", mock_server.uri()),
            userinfo_endpoint:        Some(format!("{}/userinfo", mock_server.uri())),
            jwks_uri:                 format!("{}/.well-known/jwks.json", mock_server.uri()),
            scopes_supported:         vec!["openid".to_string()],
            response_types_supported: vec!["code".to_string()],
        };
        let client = OIDCClient::new(config, "client_id", "secret").unwrap();

        let result = client.get_userinfo("dummy_token").await;
        assert!(result.is_err(), "oversized userinfo response must be rejected, got: {result:?}");
        let msg = result.unwrap_err();
        assert!(msg.contains("too large"), "error must mention size limit: {msg}");
    }

    #[tokio::test]
    async fn oidc_userinfo_within_limit_proceeds_to_parse() {
        use wiremock::{
            Mock, MockServer, ResponseTemplate,
            matchers::{method, path},
        };

        let mock_server = MockServer::start().await;
        // Valid but minimal payload — will fail at JSON parse (missing fields),
        // proving the size gate was passed.
        Mock::given(method("GET"))
            .and(path("/userinfo"))
            .respond_with(ResponseTemplate::new(200).set_body_bytes(b"{}".to_vec()))
            .mount(&mock_server)
            .await;

        let config = OIDCProviderConfig {
            issuer:                   mock_server.uri(),
            authorization_endpoint:   format!("{}/auth", mock_server.uri()),
            token_endpoint:           format!("{}/token", mock_server.uri()),
            userinfo_endpoint:        Some(format!("{}/userinfo", mock_server.uri())),
            jwks_uri:                 format!("{}/.well-known/jwks.json", mock_server.uri()),
            scopes_supported:         vec!["openid".to_string()],
            response_types_supported: vec!["code".to_string()],
        };
        let client = OIDCClient::new(config, "client_id", "secret").unwrap();

        let result = client.get_userinfo("dummy_token").await;
        // Must fail at JSON parse (missing required fields), not at size gate
        assert!(
            result.is_err(),
            "expected Err when userinfo JSON is missing required fields, got: {result:?}"
        );
        let msg = result.unwrap_err();
        assert!(
            !msg.contains("too large"),
            "size gate must not trigger for small payload: {msg}"
        );
    }
}