crank/crates/crank-community-mcp/src/transport.rs

use std::{convert::Infallible, time::Duration};

use axum::{
    Json,
    http::{
        HeaderMap, HeaderValue, StatusCode,
        header::{self, ACCEPT, HeaderName},
    },
    response::{
        IntoResponse, Response,
        sse::{Event, KeepAlive, Sse},
    },
};
use futures_util::stream;
use serde_json::Value;

use crate::jsonrpc::{
    CURRENT_PROTOCOL_VERSION, DEFAULT_PROTOCOL_VERSION, negotiated_protocol_version,
};

pub(super) const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
pub(super) const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";
pub(super) const HEADER_X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
const MAX_REQUEST_ID_LEN: usize = 128;

#[derive(Clone, Copy)]
pub(super) enum ResponseMode {
    Json,
    Sse,
}

#[derive(Clone)]
pub(super) struct AllowedOrigins {
    public_origin: Option<String>,
}

impl AllowedOrigins {
    pub(super) fn new(public_base_url: Option<String>) -> Self {
        Self {
            public_origin: public_base_url.and_then(|value| extract_origin(&value)),
        }
    }

    pub(super) fn is_allowed(&self, origin: &str) -> bool {
        if origin.starts_with("http://localhost")
            || origin.starts_with("http://127.0.0.1")
            || origin.starts_with("https://localhost")
            || origin.starts_with("https://127.0.0.1")
        {
            return true;
        }

        match &self.public_origin {
            Some(public_origin) => public_origin == origin,
            None => false,
        }
    }
}

pub(super) fn validate_origin(
    allowed_origins: &AllowedOrigins,
    headers: &HeaderMap,
) -> Result<(), StatusCode> {
    let Some(origin) = headers.get(header::ORIGIN) else {
        return Ok(());
    };
    let Ok(origin) = origin.to_str() else {
        return Err(StatusCode::FORBIDDEN);
    };

    if allowed_origins.is_allowed(origin) {
        return Ok(());
    }

    Err(StatusCode::FORBIDDEN)
}

pub(super) fn negotiate_post_response_mode(
    headers: &HeaderMap,
) -> Result<ResponseMode, StatusCode> {
    let Some(accept) = headers.get(ACCEPT) else {
        return Err(StatusCode::BAD_REQUEST);
    };
    let Ok(accept) = accept.to_str() else {
        return Err(StatusCode::BAD_REQUEST);
    };

    let mut saw_json = false;
    let mut saw_sse = false;
    let mut preferred = None;

    for part in accept.split(',') {
        let media_type = part
            .split(';')
            .next()
            .unwrap_or_default()
            .trim()
            .to_ascii_lowercase();

        match media_type.as_str() {
            "application/json" => {
                saw_json = true;
                if preferred.is_none() {
                    preferred = Some(ResponseMode::Json);
                }
            }
            "text/event-stream" => {
                saw_sse = true;
                if preferred.is_none() {
                    preferred = Some(ResponseMode::Sse);
                }
            }
            _ => {}
        }
    }

    if saw_json || saw_sse {
        return preferred.ok_or(StatusCode::NOT_ACCEPTABLE);
    }

    Err(StatusCode::NOT_ACCEPTABLE)
}

pub(super) fn validate_get_accept_header(headers: &HeaderMap) -> Result<(), StatusCode> {
    let Some(accept) = headers.get(ACCEPT) else {
        return Err(StatusCode::BAD_REQUEST);
    };
    let Ok(accept) = accept.to_str() else {
        return Err(StatusCode::BAD_REQUEST);
    };

    if accept
        .split(',')
        .map(|part| {
            part.split(';')
                .next()
                .unwrap_or_default()
                .trim()
                .to_ascii_lowercase()
        })
        .any(|media_type| media_type == "text/event-stream")
    {
        return Ok(());
    }

    Err(StatusCode::NOT_ACCEPTABLE)
}

pub(super) fn protocol_version_from_headers(headers: &HeaderMap) -> Result<String, StatusCode> {
    let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else {
        return Ok(DEFAULT_PROTOCOL_VERSION.to_owned());
    };
    let Ok(version) = version.to_str() else {
        return Err(StatusCode::BAD_REQUEST);
    };

    if negotiated_protocol_version(version).is_some() {
        return Ok(version.to_owned());
    }

    Err(StatusCode::BAD_REQUEST)
}

pub(super) fn validate_session_protocol_version(
    headers: &HeaderMap,
    negotiated_session_version: &str,
) -> Result<(), StatusCode> {
    let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else {
        return Ok(());
    };
    let Ok(version) = version.to_str() else {
        return Err(StatusCode::BAD_REQUEST);
    };

    if negotiated_protocol_version(version).is_none() {
        return Err(StatusCode::BAD_REQUEST);
    }

    if version != negotiated_session_version {
        return Err(StatusCode::BAD_REQUEST);
    }

    Ok(())
}

pub(super) fn session_id_from_headers(headers: &HeaderMap) -> Result<Option<String>, StatusCode> {
    let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) else {
        return Ok(None);
    };
    let Ok(session_id) = session_id.to_str() else {
        return Err(StatusCode::BAD_REQUEST);
    };

    Ok(Some(session_id.to_owned()))
}

pub(super) fn json_response(
    status: StatusCode,
    payload: Value,
    session_id: Option<&str>,
    protocol_version: Option<&str>,
) -> Response {
    let mut response = (status, Json(payload)).into_response();

    if let Some(session_id) = session_id {
        response.headers_mut().insert(
            HEADER_MCP_SESSION_ID,
            HeaderValue::from_str(session_id)
                .unwrap_or_else(|_| HeaderValue::from_static("invalid")),
        );
    }

    if let Some(protocol_version) = protocol_version {
        response.headers_mut().insert(
            HEADER_MCP_PROTOCOL_VERSION,
            HeaderValue::from_str(protocol_version)
                .unwrap_or_else(|_| HeaderValue::from_static(CURRENT_PROTOCOL_VERSION)),
        );
    }

    response
}

pub(super) fn transport_response(
    status: StatusCode,
    payload: Value,
    response_mode: ResponseMode,
    session_id: Option<&str>,
    protocol_version: Option<&str>,
) -> Response {
    if status == StatusCode::OK && matches!(response_mode, ResponseMode::Sse) {
        let payload = payload.to_string();
        let stream = stream::once(async move { Ok(Event::default().data(payload)) });

        return sse_response(status, stream, session_id, protocol_version);
    }

    json_response(status, payload, session_id, protocol_version)
}

pub(super) fn with_request_id_header(mut response: Response, request_id: &str) -> Response {
    if let Ok(value) = HeaderValue::from_str(request_id) {
        response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
    }
    response
}

pub(super) fn sse_response<S>(
    status: StatusCode,
    stream: S,
    session_id: Option<&str>,
    protocol_version: Option<&str>,
) -> Response
where
    S: futures_util::stream::Stream<Item = Result<Event, Infallible>> + Send + 'static,
{
    let mut response = (
        status,
        Sse::new(stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15))),
    )
        .into_response();

    if let Some(session_id) = session_id {
        response.headers_mut().insert(
            HEADER_MCP_SESSION_ID,
            HeaderValue::from_str(session_id)
                .unwrap_or_else(|_| HeaderValue::from_static("invalid")),
        );
    }

    if let Some(protocol_version) = protocol_version {
        response.headers_mut().insert(
            HEADER_MCP_PROTOCOL_VERSION,
            HeaderValue::from_str(protocol_version)
                .unwrap_or_else(|_| HeaderValue::from_static(CURRENT_PROTOCOL_VERSION)),
        );
    }

    response
}

pub(super) fn resolve_request_id(headers: &HeaderMap) -> String {
    headers
        .get(&HEADER_X_REQUEST_ID)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| is_valid_request_id(value))
        .map(ToOwned::to_owned)
        .unwrap_or_else(|| uuid::Uuid::now_v7().to_string())
}

pub(super) fn is_valid_request_id(value: &str) -> bool {
    !value.is_empty()
        && value.len() <= MAX_REQUEST_ID_LEN
        && value
            .bytes()
            .all(|byte| matches!(byte, 0x21..=0x7e) && byte != b',' && byte != b';')
}

pub(super) fn extract_origin(url: &str) -> Option<String> {
    let mut parts = url.split('/');
    let scheme = parts.next()?;
    let empty = parts.next()?;
    let authority = parts.next()?;

    if empty.is_empty() {
        return Some(format!("{scheme}//{authority}"));
    }

    None
}