oauth_proxy.rs

oauth_proxy.rs

/// Oauth2Service is a service that handles the proxying of oauth2 requests in the case of MCP server
/// We use/forward to Auth0 behind the scene but we need to intercept because:
/// * We don't want to support dynamic client registration in auth0, as we should enable it for the whole tenant and it causes security implications
///   *but* we wants users to use it without having to give them our oauth2 client and secret id before hands
///   *thus* we need to fake it in this server and returns our static creds. So it provides a smooth experience. Other implementations seems to do it too
///   There is new rfc/spec in progress to paliate to this https://auth0.com/blog/cimd-vs-dcr-mcp-registration/
///
/// * Auth0 is not Oauth2 MCP compliant by default. It uses `audience` as a parameter in the token request, but MCP/oauth2 expects `resource`
/// * We want users to generate a token for the audience/api of the CORE, *but* it is not possible by default as client check the audience/resource is the same
///   as the hostname of this server. So we intercept the request and replace the audience/resource with the audience of the CORE.
pub struct Oauth2Service {
    original_oauth_server: Uri,
    current_oauth_server: Uri,
    oauth_client_id: String,
    oauth_client_secret: String,
    oauth_audience_override: String,
    http_client: reqwest::Client,

    response_auth_required: (StatusCode, HeaderMap),
    response_oauth2_protected_resources: Arc<Value>,
}

impl Oauth2Service {
    pub fn new(
        original_oauth_server: Uri,
        current_oauth_server: Uri,
        oauth_client_id: String,
        oauth_client_secret: String,
        oauth_audience: String,
    ) -> Arc<Self> {
        let response_auth_required = {
            let headers = HeaderMap::from_iter([
                (CONTENT_TYPE, HeaderValue::from_static("application/json")),
                (WWW_AUTHENTICATE, HeaderValue::from_str(&format!("Bearer error=\"authentification_required\" resource_metadata=\"{}.well-known/oauth-protected-resource\"", current_oauth_server)).expect("Failed to format WWW-Authenticate header"))
            ]);
            (StatusCode::UNAUTHORIZED, headers)
        };

        let response_oauth2_protected_resources = Arc::new(json!({
          "resource": current_oauth_server.to_string(),
          "authorization_servers": [
            current_oauth_server.to_string()
          ],
          "scopes_supported": [
            "email", "offline_access"
          ],
          "bearer_methods_supported": [
            "header"
          ]
        }));

        let this = Self {
            original_oauth_server,
            current_oauth_server,
            oauth_client_id,
            oauth_client_secret,
            oauth_audience_override: oauth_audience,
            http_client: reqwest::Client::builder()
                .timeout(Duration::from_secs(30))
                .connect_timeout(Duration::from_secs(10))
                .redirect(reqwest::redirect::Policy::none())
                .gzip(true)
                .user_agent("Qovery MCP Server")
                .build()
                .expect("Cannot build reqwest client"),
            response_auth_required,
            response_oauth2_protected_resources,
        };
        Arc::new(this)
    }

    pub fn axum_router(self: Arc<Self>) -> Router {
        Router::new()
            .route("/oauth/register", post(Self::oidc_register))
            .route("/oauth/authorize", get(Self::oauth2_authorize))
            .route("/oauth/token", post(Self::oauth2_token_proxy))
            .route(
                "/.well-known/oauth-protected-resource",
                get(Self::oauth2_protected_resources),
            )
            .route("/.well-known/oauth-authorization-server", get(Self::oauth2_auth_server))
            .with_state(self)
    }

    pub fn response_auth_required(&self) -> axum::response::Response {
        self.response_auth_required.clone().into_response()
    }

    async fn oidc_register(
        State(srv): State<Arc<Oauth2Service>>,
        Json(auth_request): Json<RegisterPayload>,
    ) -> Json<Value> {
        // const CALLBACK_URL: [&str; 2] = [
        //     "http://localhost:4242/callback",
        //     "https://claude.ai/api/mcp/auth_callback",
        // ];
        //
        // if !auth_request
        //     .redirect_uris
        //     .iter()
        //     .any(|uri| CALLBACK_URL.iter().any(|callback| callback == uri))
        // {
        //     return Json(json!({
        //       "error": "invalid_redirect_uri",
        //       "description": format!("The only redirect_uri/callback allowed is `{CALLBACK_URL:?}`. It is a limitation of our IDP provider (Auth0). Please pin the port with --callback-port in claude or use a token")
        //     }));
        // }

        let json = json!({
          "client_id": srv.oauth_client_id,
          "client_secret": srv.oauth_client_secret,
          "client_id_issued_at": 1710000000,
          "client_name": auth_request.client_name,
          "redirect_uris": auth_request.redirect_uris,
          "grant_types": auth_request.grant_types,
          "response_types": auth_request.response_types,
          "token_endpoint_auth_method": auth_request.token_endpoint_auth_method,
        });

        Json(json)
    }

    // We intercept to inject correct audience/resource and ask to contact Auth0 to get the token
    async fn oauth2_authorize(State(srv): State<Arc<Oauth2Service>>, req: Parts) -> Redirect {
        info!("Oauth2 authorize request: {:?}", req);
        let query = req.uri.path_and_query().and_then(|p| p.query()).unwrap_or("");
        let mut params: HashMap<String, String> = serde_urlencoded::from_str(query).unwrap_or_default();
        params.insert("resource".to_string(), srv.oauth_audience_override.clone());
        params.insert("audience".to_string(), srv.oauth_audience_override.clone());

        Redirect::to(
            format!(
                "{}authorize?{}",
                srv.original_oauth_server,
                serde_urlencoded::to_string(&params).unwrap_or_default()
            )
            .as_str(),
        )
    }

    /// OAuth2 flow is done/user authentificated, we fetch the token from Auth0 with the correct audience/resource
    /// We proxy and not redirect as the info is inside the body of the request.
    async fn oauth2_token_proxy(
        State(srv): State<Arc<Oauth2Service>>,
        mut req: Parts,
        body: Bytes,
    ) -> Result<(StatusCode, http::HeaderMap, Bytes), (StatusCode, String)> {
        trace!("Token proxy request: {:?} {:?}", req, body);
        const EMPTY_HEADER: fn() -> HeaderValue = || HeaderValue::from_static("");

        let mut params: HashMap<String, String> =
            serde_urlencoded::from_bytes(&body).map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;

        params.insert("resource".to_string(), srv.oauth_audience_override.clone());
        params.insert("audience".to_string(), srv.oauth_audience_override.clone());

        let (content_type, auth) = {
            let mut headers = std::mem::take(&mut req.headers);
            (
                headers.remove(CONTENT_TYPE).unwrap_or(EMPTY_HEADER()),
                headers.remove(AUTHORIZATION).unwrap_or(EMPTY_HEADER()),
            )
        };
        let mut resp = srv
            .http_client
            .post(format!("{}oauth/token", srv.original_oauth_server))
            .header(CONTENT_TYPE, content_type)
            .header(AUTHORIZATION, auth)
            .body(serde_urlencoded::to_string(&params).unwrap_or_default())
            .send()
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

        let status = resp.status();
        let headers = std::mem::take(resp.headers_mut());
        let body = resp
            .bytes()
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

        Ok((status, headers, body))
    }

    /// Advertise the oauth2 config of this server to the client.
    /// The client uses this to know where to authenticate with the oauth2 server
    async fn oauth2_protected_resources(State(srv): State<Arc<Oauth2Service>>, req: Request) -> Json<Arc<Value>> {
        trace!("Oauth2 protected resource request: {:?}", req);
        Json(srv.response_oauth2_protected_resources.clone())
    }

    /// We fetch the information about the auth server from Auth0.
    /// And modify only the part we need for the client to contact us on specific endpoints
    /// ie.:curl https://auth.qovery.com/.well-known/oauth-authorization-server | jq .
    pub async fn oauth2_auth_server(
        State(srv): State<Arc<Oauth2Service>>,
    ) -> Result<Json<Value>, (StatusCode, String)> {
        let response = srv
            .http_client
            .get(format!(
                "{}.well-known/oauth-authorization-server",
                srv.original_oauth_server
            ))
            .send()
            .await
            .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;

        trace!("Oauth2 auth server response: {:?}", response);
        let mut json: Value = response
            .json()
            .await
            .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;

        *json.get_mut("token_endpoint").unwrap_or(&mut Value::Null) =
            Value::String(format!("{}oauth/token", srv.current_oauth_server));
        *json.get_mut("authorization_endpoint").unwrap_or(&mut Value::Null) =
            Value::String(format!("{}oauth/authorize", srv.current_oauth_server));
        *json.get_mut("registration_endpoint").unwrap_or(&mut Value::Null) =
            Value::String(format!("{}oauth/register", srv.current_oauth_server));

        Ok(Json(json))
    }
}

#[derive(Default, Debug, Clone, PartialEq, Deserialize)]
struct RegisterPayload {
    pub client_name: String,
    pub grant_types: Vec<String>,
    pub redirect_uris: Vec<String>,
    pub response_types: Vec<String>,
    pub token_endpoint_auth_method: String,
}