aether-auth 0.1.8

use super::credential_store::McpCredentialStore;
use crate::{OAuthCredentialStorage, OAuthError, OAuthHandler};
use rmcp::transport::auth::{AuthClient, AuthError, AuthorizationManager, OAuthClientConfig};
use std::sync::Arc;

const OAUTH_CLIENT_NAME: &str = "Aether MCP Client";

/// Returns `Ok(Some(manager))` if credentials were found and initialized successfully,
/// `Ok(None)` if no stored credentials exist, or `Err` on failure.
///
/// When `expected_client_id` is given, stored credentials issued to a different client
/// are ignored, forcing a fresh authorization flow.
pub async fn create_auth_manager_from_store(
    server_id: &str,
    base_url: &str,
    expected_client_id: Option<&str>,
    store: Arc<dyn OAuthCredentialStorage>,
) -> Result<Option<AuthorizationManager>, OAuthError> {
    let credential_store =
        McpCredentialStore::new(store, server_id.to_string()).with_expected_client_id(expected_client_id);

    let mut auth_manager = AuthorizationManager::new(base_url).await.map_err(rmcp_err("OAuth init failed"))?;
    auth_manager.set_credential_store(credential_store);

    if auth_manager.initialize_from_store().await.map_err(rmcp_err("credential store load failed"))? {
        Ok(Some(auth_manager))
    } else {
        Ok(None)
    }
}

/// Run a full OAuth authorization flow against an MCP server.
///
/// Discovers the server's authorization metadata, obtains a client (the pre-registered
/// `client_id` when given, dynamic registration otherwise), opens the browser (via the
/// handler), handles the callback, and returns an `AuthClient` ready for authenticated
/// HTTP transport. When `store` is `Some`, the resulting tokens are persisted via the
/// store; when `None`, tokens live only for the lifetime of the returned `AuthClient`.
pub async fn perform_oauth_flow(
    server_id: &str,
    base_url: &str,
    handler: &dyn OAuthHandler,
    client_id: Option<&str>,
    store: Option<Arc<dyn OAuthCredentialStorage>>,
) -> Result<AuthClient<reqwest::Client>, OAuthError> {
    let mut manager = AuthorizationManager::new(base_url).await.map_err(rmcp_err("OAuth init failed"))?;

    if let Some(store) = store {
        manager.set_credential_store(
            McpCredentialStore::new(store, server_id.to_string()).with_expected_client_id(client_id),
        );
    }

    let metadata = manager.discover_metadata().await.map_err(rmcp_err("OAuth metadata discovery failed"))?;
    manager.set_metadata(metadata);

    let scopes = manager.select_scopes(None, &[]);
    let scope_refs = scopes.iter().map(String::as_str).collect::<Vec<_>>();
    let client_config = match client_id {
        Some(client_id) => OAuthClientConfig::new(client_id, handler.redirect_uri()),
        None => manager
            .register_client(OAUTH_CLIENT_NAME, handler.redirect_uri(), &scope_refs)
            .await
            .map_err(rmcp_err("dynamic client registration failed"))?,
    };
    manager.configure_client(client_config).map_err(rmcp_err("OAuth client configuration failed"))?;

    let auth_url =
        manager.get_authorization_url(&scope_refs).await.map_err(rmcp_err("get_authorization_url failed"))?;

    // Some authorization servers (e.g. Sentry) bake `resource` into their
    // authorization_endpoint metadata. rmcp then adds its own `resource` param,
    // producing a duplicate that the server rejects with "invalid_target".
    let callback = handler.authorize(&dedupe_query_params(&auth_url)).await?;

    manager
        .exchange_code_for_token(&callback.code, &callback.state)
        .await
        .map_err(rmcp_err("token exchange failed"))?;

    Ok(AuthClient::new(reqwest::Client::default(), manager))
}

fn rmcp_err(context: &'static str) -> impl Fn(AuthError) -> OAuthError {
    move |e| OAuthError::Rmcp(format!("{context}: {e}"))
}

fn dedupe_query_params(url_str: &str) -> String {
    let Ok(mut url) = url::Url::parse(url_str) else {
        return url_str.to_string();
    };
    let pairs: std::collections::HashMap<String, String> =
        url.query_pairs().map(|(k, v)| (k.into_owned(), v.into_owned())).collect();
    url.query_pairs_mut().clear().extend_pairs(&pairs);
    url.to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn dedupes_duplicate_resource_param() {
        let input = "https://example.com/oauth/authorize?resource=https%3A%2F%2Fa%2Fb&response_type=code&client_id=x&resource=https%3A%2F%2Fa%2Fb";
        let out = dedupe_query_params(input);
        let url = url::Url::parse(&out).unwrap();
        let resources: Vec<_> = url.query_pairs().filter(|(k, _)| k == "resource").collect();
        assert_eq!(resources.len(), 1);
        assert_eq!(resources[0].1, "https://a/b");
        assert!(url.query_pairs().any(|(k, _)| k == "response_type"));
        assert!(url.query_pairs().any(|(k, _)| k == "client_id"));
    }

    #[test]
    fn preserves_unique_params() {
        let input = "https://example.com/?resource=x&other=y";
        let out = dedupe_query_params(input);
        let url = url::Url::parse(&out).unwrap();
        let pairs: Vec<_> = url.query_pairs().collect();
        assert_eq!(pairs.len(), 2);
    }
}