codex/codex-rs/rmcp-client/src/ema_http_client.rs

//! Renewable, endpoint-confined EMA bearers. An unauthorized MCP operation is never replayed.
//! Authentication shares the request deadline, and a rejected bearer cannot evict a newer one.

use std::future::Future;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::MutexGuard;
use std::sync::PoisonError;
use std::time::Duration;

use codex_exec_server::ExecServerError;
use codex_exec_server::HttpClient;
use codex_exec_server::HttpHeader;
use codex_exec_server::HttpRedirectPolicy;
use codex_exec_server::HttpRequestParams;
use codex_exec_server::HttpRequestResponse;
use codex_exec_server::HttpResponseBodyStream;
use futures::future::BoxFuture;
use tokio::sync::OwnedRwLockReadGuard;
use tokio::sync::RwLock;
use tokio::sync::Semaphore;
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
use url::Url;

use crate::EmaAccessToken;
use crate::is_authentication_required_error;
use crate::validate_ema_auth_resource;

/// Supplies resource tokens for one immutable endpoint and pinned enterprise login.
/// Credential ownership is checked when renewing a token, not for cached bearer requests.
pub trait EmaTokenSource: Send + Sync {
    fn endpoint(&self) -> &Url;
    /// Acquires and validates credentials for renewal. Implementations recheck ownership
    /// and a saved grant's generation after exchange before publishing the bearer.
    fn access_token(&self) -> impl Future<Output = anyhow::Result<EmaAccessToken>> + Send;
}

/// Revocable authority held while one EMA connection awaits HTTP responses.
#[derive(Clone)]
pub struct EmaRequestAuthority {
    revoked: CancellationToken,
    is_current: Arc<dyn Fn() -> bool + Send + Sync>,
    responses: Arc<RwLock<()>>,
    policy: Option<codex_config::CloudConfigBundleBinding>,
}

impl Default for EmaRequestAuthority {
    fn default() -> Self {
        Self::new(CancellationToken::new(), || true)
    }
}

impl EmaRequestAuthority {
    /// Couples prompt cancellation with a synchronous check at each HTTP boundary.
    /// The check must be cheap and must not perform credential or network I/O.
    pub fn new(
        revoked: CancellationToken,
        is_current: impl Fn() -> bool + Send + Sync + 'static,
    ) -> Self {
        Self {
            revoked,
            is_current: Arc::new(is_current),
            responses: Arc::default(),
            policy: None,
        }
    }

    /// Policy changes stop new dispatches without cancelling already admitted responses.
    pub fn with_policy(mut self, policy: Option<codex_config::CloudConfigBundleBinding>) -> Self {
        self.policy = policy;
        self
    }

    /// Isolate one connection's credential failure while sharing the owner's response drain.
    pub fn child(&self) -> Self {
        let owner = self.clone();
        Self {
            revoked: self.revoked.child_token(),
            is_current: Arc::new(move || !owner.is_revoked()),
            responses: Arc::clone(&self.responses),
            policy: self.policy.clone(),
        }
    }

    pub fn revoke(&self) {
        self.revoked.cancel();
    }

    pub fn is_revoked(&self) -> bool {
        if !self.revoked.is_cancelled() && !(self.is_current)() {
            self.revoke();
        }
        self.revoked.is_cancelled()
    }

    pub fn check(&self) -> Result<(), ExecServerError> {
        if self.is_revoked() {
            return Err(revoked_authority_error());
        }
        Ok(())
    }

    pub fn check_new_request(&self) -> Result<(), ExecServerError> {
        self.check()?;
        if self.policy.as_ref().is_some_and(|binding| {
            binding.read().status != codex_config::CloudConfigBundleBindingStatus::Current
        }) {
            return Err(ExecServerError::RequestNotSent(
                "enterprise MCP policy is updating or unavailable; retry after it is applied"
                    .into(),
            ));
        }
        Ok(())
    }

    /// Wait until response futures admitted before revocation complete or are cancelled.
    pub async fn quiesce(&self) {
        let _responses = self.responses.clone().write_owned().await;
    }

    async fn acquire(&self) -> Result<OwnedRwLockReadGuard<()>, ExecServerError> {
        let response = self.responses.clone().read_owned().await;
        self.check_new_request()?;
        Ok(response)
    }

    async fn cancelled(&self) {
        self.revoked.cancelled().await;
    }
}

/// Adds a renewable resource bearer to the existing runtime-selected HTTP capability.
pub struct EmaAuthenticatedHttpClient<S> {
    inner: Arc<dyn HttpClient>,
    source: S,
    authority: EmaRequestAuthority,
    token: Mutex<Option<CachedToken>>,
    renewal: Semaphore,
}

struct CachedToken {
    value: String,
    refresh_at: Instant,
}

impl CachedToken {
    fn new(token: EmaAccessToken) -> Self {
        // Bound reuse even when the authorization server omits an expiry.
        let lifetime = token
            .expires_in
            .unwrap_or(Duration::from_secs(300))
            .min(Duration::from_secs(300));
        let skew = Duration::from_secs(30).min(lifetime / 2);
        Self {
            value: token.access_token,
            refresh_at: token.received_at + lifetime.saturating_sub(skew),
        }
    }
}

impl<S: EmaTokenSource> EmaAuthenticatedHttpClient<S> {
    pub fn new(
        inner: Arc<dyn HttpClient>,
        source: S,
        access_token: EmaAccessToken,
        authority: EmaRequestAuthority,
    ) -> Self {
        Self {
            inner,
            source,
            authority,
            token: Mutex::new(Some(CachedToken::new(access_token))),
            renewal: Semaphore::new(/*permits*/ 1),
        }
    }

    fn cached_token(&self) -> MutexGuard<'_, Option<CachedToken>> {
        self.token.lock().unwrap_or_else(PoisonError::into_inner)
    }

    async fn authenticate(
        &self,
        params: &mut HttpRequestParams,
    ) -> Result<String, ExecServerError> {
        let expected = self.source.endpoint();
        let target = Url::parse(&params.url).map_err(|_| {
            ExecServerError::HttpRequest("invalid enterprise MCP destination".to_string())
        })?;
        if expected.origin() != target.origin()
            || expected.path() != target.path()
            || expected.query() != target.query()
            || validate_ema_auth_resource(&params.url, /*resource*/ None).is_err()
        {
            return Err(ExecServerError::HttpRequest(
                "refusing to send enterprise bearer outside the configured MCP endpoint"
                    .to_string(),
            ));
        }
        let current_token = || {
            self.cached_token()
                .as_ref()
                .filter(|token| Instant::now() < token.refresh_at)
                .map(|token| token.value.clone())
        };
        let value = if let Some(value) = current_token() {
            value
        } else {
            let _renewal = self.renewal.acquire().await.map_err(|_| {
                ExecServerError::HttpRequest("enterprise token renewal gate closed".to_string())
            })?;
            if let Some(value) = current_token() {
                value
            } else {
                let renewed = self.source.access_token().await.map_err(|error| {
                    if let Err(refusal) = self.authority.check_new_request() {
                        return refusal;
                    }
                    if is_authentication_required_error(&error) {
                        self.authority.revoke();
                    }
                    auth_error(error)
                })?;
                let token = CachedToken::new(renewed);
                if Instant::now() >= token.refresh_at {
                    return Err(ExecServerError::HttpRequest(
                        "enterprise MCP bearer expired during credential validation".to_string(),
                    ));
                }
                let value = token.value.clone();
                *self.cached_token() = Some(token);
                value
            }
        };
        params.redirect_policy = HttpRedirectPolicy::Stop;
        params
            .headers
            .retain(|header| !header.name.eq_ignore_ascii_case("authorization"));
        params.headers.push(HttpHeader {
            name: "authorization".to_string(),
            value: format!("Bearer {value}"),
            value_env_var: None,
        });
        Ok(value)
    }

    fn clear_rejected_token(&self, rejected: &str) {
        let mut token = self.cached_token();
        if token.as_ref().is_some_and(|token| token.value == rejected) {
            *token = None;
        }
    }

    fn validate_response(
        &self,
        response: &HttpRequestResponse,
        token: &str,
    ) -> Result<(), ExecServerError> {
        if response.status == 401 {
            self.clear_rejected_token(token);
            return Err(ExecServerError::AuthenticationRequired(
                "enterprise MCP bearer was rejected; reconnect to renew authorization".to_string(),
            ));
        }
        // An outer redirect wrapper may still follow a response even though the
        // inner HTTP capability received Stop. Do not let it replay this operation.
        if (300..400).contains(&response.status) {
            return Err(ExecServerError::HttpRequest(
                "enterprise MCP redirects are not allowed".to_string(),
            ));
        }
        Ok(())
    }
}

fn auth_error(error: anyhow::Error) -> ExecServerError {
    if let Some(error) = error.downcast_ref::<crate::EmaTokenSourceError>() {
        return if error.is_authentication_required() {
            ExecServerError::AuthenticationRequired(error.to_string())
        } else {
            ExecServerError::HttpRequest(error.to_string())
        };
    }
    if is_authentication_required_error(&error) {
        ExecServerError::AuthenticationRequired(
            "enterprise MCP authentication is required".to_string(),
        )
    } else {
        ExecServerError::HttpRequest(
            "enterprise MCP token acquisition or credential validation failed".to_string(),
        )
    }
}

impl<S: EmaTokenSource> HttpClient for EmaAuthenticatedHttpClient<S> {
    fn http_request(
        &self,
        mut params: HttpRequestParams,
    ) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
        Box::pin(async move {
            let deadline = params
                .timeout_ms
                .map(|timeout_ms| Instant::now() + Duration::from_millis(timeout_ms));
            let request = async {
                let _authority = self.authority.acquire().await?;
                let token = self.authenticate(&mut params).await?;
                self.authority.check_new_request()?;
                if let Some(deadline) = deadline {
                    params.timeout_ms = Some(remaining_timeout(deadline)?);
                }
                let response = self.inner.http_request(params).await?;
                self.authority.check()?;
                self.validate_response(&response, &token)?;
                Ok(response)
            };
            tokio::select! {
                biased;
                () = self.authority.cancelled() => Err(revoked_authority_error()),
                () = deadline_elapsed(deadline) => Err(request_timeout_error()),
                result = request => result,
            }
        })
    }

    fn http_request_stream(
        &self,
        mut params: HttpRequestParams,
    ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
        Box::pin(async move {
            let deadline = params
                .timeout_ms
                .map(|timeout_ms| Instant::now() + Duration::from_millis(timeout_ms));
            let request = async {
                let _authority = self.authority.acquire().await?;
                let token = self.authenticate(&mut params).await?;
                self.authority.check_new_request()?;
                if let Some(deadline) = deadline {
                    params.timeout_ms = Some(remaining_timeout(deadline)?);
                }
                let response = self.inner.http_request_stream(params).await?;
                self.authority.check()?;
                self.validate_response(&response.0, &token)?;
                Ok((
                    response.0,
                    response.1.cancel_on(self.authority.revoked.clone()),
                ))
            };
            tokio::select! {
                biased;
                () = self.authority.cancelled() => Err(revoked_authority_error()),
                () = deadline_elapsed(deadline) => Err(request_timeout_error()),
                result = request => result,
            }
        })
    }
}

async fn deadline_elapsed(deadline: Option<Instant>) {
    match deadline {
        Some(deadline) => tokio::time::sleep_until(deadline).await,
        None => std::future::pending().await,
    }
}

fn remaining_timeout(deadline: Instant) -> Result<u64, ExecServerError> {
    let remaining = deadline
        .checked_duration_since(Instant::now())
        .filter(|remaining| !remaining.is_zero())
        .ok_or_else(request_timeout_error)?;
    Ok(u64::try_from(remaining.as_millis())
        .unwrap_or(u64::MAX)
        .max(1))
}

fn request_timeout_error() -> ExecServerError {
    ExecServerError::HttpRequest("enterprise MCP request timed out".to_string())
}

fn revoked_authority_error() -> ExecServerError {
    ExecServerError::AuthenticationRequired(
        "enterprise MCP authorization changed; reconnect before using this server".to_string(),
    )
}

#[cfg(test)]
#[path = "ema_http_client_tests.rs"]
mod tests;