yew-oauth2 0.14.0

use crate::{
    agent::{
        InnerConfig, LogoutOptions, OAuth2Error,
        client::{Client, LoginContext, expires},
    },
    config::openid::{self, MetadataSource, MetadataUrls},
    context::{Authentication, OAuth2Context},
};
use async_trait::async_trait;
use gloo_utils::window;
use oauth2::TokenResponse as _;
use openidconnect::{
    AuthUrl, AuthorizationCode, ClientId, CsrfToken, EmptyAdditionalClaims, EndpointNotSet,
    EndpointSet, IdTokenClaims, IssuerUrl, JsonWebKeySet, JsonWebKeySetUrl, Nonce,
    PkceCodeChallenge, PkceCodeVerifier, ProviderMetadata, RedirectUrl, RefreshToken, Scope,
    TokenResponse, TokenUrl, UserInfoUrl,
    core::{
        CoreAuthDisplay, CoreAuthenticationFlow, CoreClaimName, CoreClaimType, CoreClient,
        CoreClientAuthMethod, CoreGenderClaim, CoreGrantType, CoreJsonWebKey,
        CoreJweContentEncryptionAlgorithm, CoreJweKeyManagementAlgorithm, CoreResponseMode,
        CoreResponseType, CoreSubjectIdentifierType, CoreTokenResponse,
    },
};
use reqwest::Url;
use serde::{Deserialize, Serialize};
use std::{fmt::Debug, rc::Rc};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct OpenIdLoginState {
    pub pkce_verifier: String,
    pub nonce: String,
}

const DEFAULT_POST_LOGOUT_DIRECT_NAME: &str = "post_logout_redirect_uri";

/// An OpenID Connect based client implementation
#[derive(Clone, Debug)]
pub struct OpenIdClient {
    /// The http client
    http_client: openidconnect::reqwest::Client,
    /// The client
    client: ExtendedClient,
    /// An override for the URL to end the session (logout)
    end_session_url: Option<Url>,
    /// A URL to direct to after the logout was performed
    after_logout_url: Option<String>,
    /// The name of the query parameter sent to the issuer, containing the post-logout redirect URL
    post_logout_redirect_name: Option<String>,
    /// Additional audiences of the ID token which are considered trustworthy
    additional_trusted_audiences: Vec<String>,
    /// Specifies whether the issuer claim must match the expected issuer URL for the provider.
    pub require_issuer_match: bool,
}

/// Additional metadata read from the discovery endpoint
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AdditionalProviderMetadata {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub end_session_endpoint: Option<Url>,
}

impl openidconnect::AdditionalProviderMetadata for AdditionalProviderMetadata {}

pub type ExtendedProviderMetadata = ProviderMetadata<
    AdditionalProviderMetadata,
    CoreAuthDisplay,
    CoreClientAuthMethod,
    CoreClaimName,
    CoreClaimType,
    CoreGrantType,
    CoreJweContentEncryptionAlgorithm,
    CoreJweKeyManagementAlgorithm,
    CoreJsonWebKey,
    CoreResponseMode,
    CoreResponseType,
    CoreSubjectIdentifierType,
>;

pub type ExtendedClient = CoreClient<
    EndpointSet,
    EndpointNotSet,
    EndpointNotSet,
    EndpointNotSet,
    EndpointSet,
    EndpointSet,
>;

#[async_trait(? Send)]
impl Client for OpenIdClient {
    type TokenResponse = CoreTokenResponse;
    type Configuration = openid::Config;
    type LoginState = OpenIdLoginState;
    type SessionState = (
        String,
        Rc<IdTokenClaims<EmptyAdditionalClaims, CoreGenderClaim>>,
    );

    async fn from_config(config: Self::Configuration) -> Result<Self, OAuth2Error> {
        let openid::Config {
            client_id,
            issuer_url,
            metadata_source,
            end_session_url,
            after_logout_url,
            post_logout_redirect_name,
            additional_trusted_audiences,
            require_issuer_match,
        } = config;

        let http_client = openidconnect::reqwest::ClientBuilder::new()
            // Following redirects opens the client up to SSRF vulnerabilities.
            // .redirect(openidconnect::reqwest::redirect::Policy::none())
            .build()
            .map_err(|err| {
                OAuth2Error::Configuration(format!("Failed to build HTTP client: {err}"))
            })?;

        let issuer = IssuerUrl::new(issuer_url)
            .map_err(|err| OAuth2Error::Configuration(format!("invalid issuer URL: {err}")))?;

        let (client, end_session_url) = match metadata_source {
            MetadataSource::Discovery => {
                Self::build_client_from_discovery(&http_client, issuer, client_id, end_session_url)
                    .await?
            }
            MetadataSource::Manual(urls) => {
                Self::build_client_from_urls(&http_client, issuer, client_id, end_session_url, urls)
                    .await?
            }
        };
        Ok(Self {
            http_client,
            client,
            end_session_url,
            after_logout_url,
            post_logout_redirect_name,
            additional_trusted_audiences,
            require_issuer_match,
        })
    }

    fn set_redirect_uri(mut self, url: Url) -> Self {
        self.client = self.client.set_redirect_uri(RedirectUrl::from_url(url));
        self
    }

    fn make_login_context(
        &self,
        config: &InnerConfig,
        redirect_url: Url,
    ) -> Result<LoginContext<Self::LoginState>, OAuth2Error> {
        let client = self
            .client
            .clone()
            .set_redirect_uri(RedirectUrl::from_url(redirect_url));

        let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256();

        let mut req = client.authorize_url(
            CoreAuthenticationFlow::AuthorizationCode,
            CsrfToken::new_random,
            Nonce::new_random,
        );

        for scope in &config.scopes {
            req = req.add_scope(Scope::new(scope.clone()));
        }

        if let Some(audience) = &config.audience {
            req = req.add_extra_param("audience".to_string(), audience);
        }

        let (url, state, nonce) = req.set_pkce_challenge(pkce_challenge).url();

        Ok(LoginContext {
            url,
            csrf_token: state.secret().clone(),
            state: OpenIdLoginState {
                pkce_verifier: pkce_verifier.secret().clone(),
                nonce: nonce.secret().clone(),
            },
        })
    }

    async fn exchange_code(
        &self,
        code: String,
        state: Self::LoginState,
    ) -> Result<(OAuth2Context, Self::SessionState), OAuth2Error> {
        let pkce_verifier = PkceCodeVerifier::new(state.pkce_verifier);

        let result = self
            .client
            .exchange_code(AuthorizationCode::new(code))
            .set_pkce_verifier(pkce_verifier)
            .request_async(&self.http_client)
            .await
            .map_err(|err| OAuth2Error::LoginResult(format!("failed to exchange code: {err}")))?;

        log::debug!("Exchange code result: {:?}", result);

        let id_token = result.extra_fields().id_token().ok_or_else(|| {
            OAuth2Error::LoginResult("Server did not return an ID token".to_string())
        })?;

        let claims = Rc::new(
            id_token
                .clone()
                .into_claims(
                    &self
                        .client
                        .id_token_verifier()
                        .require_issuer_match(self.require_issuer_match)
                        .set_other_audience_verifier_fn(|aud| {
                            self.additional_trusted_audiences.contains(aud)
                        }),
                    &Nonce::new(state.nonce),
                )
                .map_err(|err| {
                    OAuth2Error::LoginResult(format!("failed to verify ID token: {err}"))
                })?,
        );

        Ok((
            OAuth2Context::Authenticated(Authentication {
                access_token: result.access_token().secret().to_string(),
                id_token: result.id_token().map(|t| t.to_string()),
                refresh_token: result.refresh_token().map(|t| t.secret().to_string()),
                expires: expires(result.expires_in()),
                claims: Some(claims.clone()),
            }),
            (id_token.to_string(), claims),
        ))
    }

    async fn exchange_refresh_token(
        &self,
        refresh_token: String,
        session_state: Self::SessionState,
    ) -> Result<(OAuth2Context, Self::SessionState), OAuth2Error> {
        let result = self
            .client
            .exchange_refresh_token(&RefreshToken::new(refresh_token))
            .request_async(&self.http_client)
            .await
            .map_err(|err| {
                OAuth2Error::Refresh(format!("failed to exchange refresh token: {err}"))
            })?;

        Ok((
            OAuth2Context::Authenticated(Authentication {
                access_token: result.access_token().secret().to_string(),
                id_token: result.id_token().map(|t| t.to_string()),
                refresh_token: result.refresh_token().map(|t| t.secret().to_string()),
                expires: expires(result.expires_in()),
                claims: Some(session_state.1.clone()),
            }),
            session_state,
        ))
    }

    fn logout(&self, session_state: Self::SessionState, options: LogoutOptions) {
        if let Some(url) = &self.end_session_url {
            let mut url = url.clone();

            let name = self
                .post_logout_redirect_name
                .as_deref()
                .unwrap_or(DEFAULT_POST_LOGOUT_DIRECT_NAME);

            url.query_pairs_mut()
                .append_pair("id_token_hint", &session_state.0);

            if let Some(after) = options
                .target
                .map(|url| url.to_string())
                .or_else(|| self.after_logout_url())
            {
                url.query_pairs_mut().append_pair(name, &after);
            }

            log::info!("Navigating to: {url}");

            window().location().replace(url.as_str()).ok();
        } else {
            log::warn!("Found no session end URL");
        }
    }
}

impl OpenIdClient {
    fn after_logout_url(&self) -> Option<String> {
        if let Some(after) = &self.after_logout_url {
            if Url::parse(after).is_ok() {
                // test if this is an absolute URL
                return Some(after.to_string());
            }

            window()
                .location()
                .href()
                .ok()
                .and_then(|url| {
                    Url::parse(&url)
                        .and_then(|current| current.join(after))
                        .ok()
                })
                .map(|u| u.to_string())
        } else {
            window().location().href().ok()
        }
    }
    async fn build_client_from_discovery(
        http_client: &openidconnect::reqwest::Client,
        issuer: IssuerUrl,
        client_id: String,
        end_session_url: Option<String>,
    ) -> Result<(ExtendedClient, Option<Url>), OAuth2Error> {
        let metadata = ExtendedProviderMetadata::discover_async(issuer, http_client)
            .await
            .map_err(|err| {
                OAuth2Error::Configuration(format!("Failed to discover client: {err}"))
            })?;

        // Extract the URIs we MUST have
        let auth_uri = metadata.authorization_endpoint().clone();

        let token_uri = metadata
            .token_endpoint()
            .ok_or_else(|| {
                OAuth2Error::Configuration("Provider missing required token endpoint".into())
            })?
            .clone();

        let user_info_uri = metadata
            .userinfo_endpoint()
            .ok_or_else(|| {
                OAuth2Error::Configuration("Provider missing required auth info endpoint".into())
            })?
            .clone();
        let end_session_url = end_session_url
            .map(|url| Url::parse(&url))
            .transpose()
            .map_err(|err| {
                OAuth2Error::Configuration(format!("Unable to parse end_session_url: {err}"))
            })?
            .or_else(|| metadata.additional_metadata().end_session_endpoint.clone());

        Ok((
            CoreClient::from_provider_metadata(metadata, ClientId::new(client_id), None)
                .set_auth_uri(auth_uri)
                .set_token_uri(token_uri)
                .set_user_info_url(user_info_uri),
            end_session_url,
        ))
    }

    async fn build_client_from_urls(
        http_client: &openidconnect::reqwest::Client,
        issuer: IssuerUrl,
        client_id: String,
        end_session_url: Option<String>,
        urls: MetadataUrls,
    ) -> Result<(ExtendedClient, Option<Url>), OAuth2Error> {
        let auth_uri = AuthUrl::new(urls.auth)
            .map_err(|err| OAuth2Error::Configuration(format!("invalid auth URL: {err}")))?;

        let token_uri = TokenUrl::new(urls.token)
            .map_err(|err| OAuth2Error::Configuration(format!("invalid token URL: {err}")))?;

        let jwks_uri = JsonWebKeySetUrl::new(urls.jwks)
            .map_err(|err| OAuth2Error::Configuration(format!("invalid jwks URL: {err}")))?;

        let jwks = JsonWebKeySet::fetch_async(&jwks_uri, http_client)
            .await
            .map_err(|err| OAuth2Error::Configuration(format!("Could not fetch jwks: {err}")))?;

        let user_info_uri = UserInfoUrl::new(urls.user_info).map_err(|err| {
            OAuth2Error::Configuration(format!("Unable to parse user_info_url: {err}"))
        })?;

        let end_session_url = end_session_url
            .map(|url| Url::parse(&url))
            .transpose()
            .map_err(|err| {
                OAuth2Error::Configuration(format!("Unable to parse end_session_url: {err}"))
            })?;

        Ok((
            CoreClient::new(ClientId::new(client_id), issuer, jwks)
                .set_auth_uri(auth_uri)
                .set_token_uri(token_uri)
                .set_user_info_url(user_info_uri),
            end_session_url,
        ))
    }
}