use std::sync::Arc;
use axum::Router as AxumRouter;
use axum::body::Body;
use axum::extract::State;
use axum::http::{HeaderMap, HeaderName, HeaderValue, Request, StatusCode, header};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use axum::routing::MethodRouter;
use serde_json::Value;
use super::config::Streaming;
use super::hardening::{self, enforce_deadline, enforce_host, enforce_origin, enforce_rate_limit};
use super::limits::RateLimiter;
use super::precondition::{authenticate, read_body};
use super::{
Anonymous, AuthPolicy, Authenticator, EventTail, JSON_MEDIA_TYPE, Refusal, SecurityHeaders,
ServerConfig, SessionStore,
};
use crate::adapters::primary::transport::origin::OriginPolicy;
use crate::domain::ProtocolVersion;
use crate::domain::jsonrpc::{Frame, WireId, decode};
use crate::domain::mcp::McpSessionId;
use crate::domain::registry::AuthContext;
const SESSION_HEADER: HeaderName = HeaderName::from_static("mcp-session-id");
const LAST_EVENT_ID: HeaderName = HeaderName::from_static("last-event-id");
const PROTOCOL_HEADER: HeaderName = HeaderName::from_static("mcp-protocol-version");
#[derive(Debug, Clone)]
struct SecurityLayer {
content_security_policy: HeaderValue,
strict_transport_security: HeaderValue,
x_frame_options: HeaderValue,
}
impl SecurityLayer {
fn parse(headers: &SecurityHeaders) -> Result<Self, InvalidHeaderValue> {
Ok(Self {
content_security_policy: HeaderValue::from_str(headers.content_security_policy())?,
strict_transport_security: HeaderValue::from_str(headers.strict_transport_security())?,
x_frame_options: HeaderValue::from_str(headers.x_frame_options())?,
})
}
}
#[derive(Debug, thiserror::Error)]
#[error("a configured security header is not a valid HTTP header value: {0}")]
pub struct InvalidHeaderValue(axum::http::header::InvalidHeaderValue);
impl From<axum::http::header::InvalidHeaderValue> for InvalidHeaderValue {
fn from(source: axum::http::header::InvalidHeaderValue) -> Self {
Self(source)
}
}
pub(crate) struct RouterState<S, H> {
config: ServerConfig,
sessions: Arc<SessionStore<S>>,
handler: H,
policy: AuthPolicy,
authenticator: Arc<dyn Authenticator + Send + Sync>,
origin: OriginPolicy,
rate_limiter: RateLimiter,
}
impl<S, H> RouterState<S, H> {
pub(crate) const fn origin(&self) -> &OriginPolicy {
&self.origin
}
pub(crate) const fn config(&self) -> &ServerConfig {
&self.config
}
pub(crate) const fn rate_limiter(&self) -> &RateLimiter {
&self.rate_limiter
}
}
#[derive(Debug, Clone, Copy)]
enum Header<'h> {
Absent,
Value(&'h str),
Unreadable,
}
impl<'h> Header<'h> {
fn read(headers: &'h HeaderMap, name: &HeaderName) -> Self {
let Some(raw) = headers.get(name) else {
return Self::Absent;
};
raw.to_str().map_or(Self::Unreadable, Self::Value)
}
}
pub trait EventSource: Send + Sync + 'static {
fn tail(&self) -> &EventTail<Value>;
}
impl EventSource for EventTail<Value> {
fn tail(&self) -> &EventTail<Value> {
self
}
}
#[derive(Debug)]
pub enum SessionHandle<S> {
Establishing,
Stateless,
Live {
id: McpSessionId,
session: Arc<S>,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Unsessioned {
Establishing,
Stateless,
NeedsSession,
}
#[derive(Debug)]
pub struct Inbound<'a, S> {
pub context: &'a AuthContext,
pub session: SessionHandle<S>,
pub body: &'a [u8],
pub protocol: Option<ProtocolVersion>,
pub headers: &'a HeaderMap,
pub streaming: Streaming,
}
pub trait RequestHandler<S>: Send + Sync + 'static {
fn unsessioned(&self, body: &[u8], headers: &HeaderMap) -> Unsessioned;
fn dispatch(
&self,
call: Inbound<'_, S>,
) -> impl Future<Output = Result<Response, Refusal>> + Send;
}
#[derive(Debug, thiserror::Error)]
pub enum RouterError {
#[error("a configured security header is not a legal HTTP header value: {0}")]
Header(#[source] InvalidHeaderValue),
#[error(
"with_authenticator was wired but the policy is still AuthPolicy::Public; \
a real authenticator paired with the default public policy fails open \
silently. Add .with_policy(AuthPolicy::Require(scopes)) after \
.with_authenticator(...), or drop the .with_authenticator(...) call if \
the endpoint is meant to be anonymous."
)]
UnprotectedAuthenticator,
}
impl From<InvalidHeaderValue> for RouterError {
fn from(value: InvalidHeaderValue) -> Self {
Self::Header(value)
}
}
pub struct McpRouter<S, H> {
config: ServerConfig,
sessions: Arc<SessionStore<S>>,
handler: H,
policy: AuthPolicy,
authenticator: Arc<dyn Authenticator + Send + Sync>,
origin: OriginPolicy,
rate_limiter: RateLimiter,
authenticator_was_explicit: bool,
}
impl<S, H> core::fmt::Debug for McpRouter<S, H> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("McpRouter")
.field("endpoint", &self.config.endpoint())
.field("streaming", &self.config.streaming())
.field("policy", &self.policy)
.finish_non_exhaustive()
}
}
impl<S, H> McpRouter<S, H>
where
S: EventSource,
H: RequestHandler<S>,
{
#[must_use]
pub fn new(config: ServerConfig, sessions: Arc<SessionStore<S>>, handler: H) -> Self {
Self {
rate_limiter: RateLimiter::new(*config.rate_limit()),
origin: OriginPolicy::strict(None),
config,
sessions,
handler,
policy: AuthPolicy::Public,
authenticator: Arc::new(Anonymous),
authenticator_was_explicit: false,
}
}
#[must_use]
pub fn with_origin_policy(mut self, origin: OriginPolicy) -> Self {
self.origin = origin;
self
}
#[must_use]
pub fn with_rate_limiter(mut self, limiter: RateLimiter) -> Self {
self.rate_limiter = limiter;
self
}
#[must_use]
pub fn with_authenticator(
mut self,
authenticator: Arc<dyn Authenticator + Send + Sync>,
) -> Self {
self.authenticator = authenticator;
self.authenticator_was_explicit = true;
self
}
#[must_use]
pub const fn with_policy(mut self, policy: AuthPolicy) -> Self {
self.policy = policy;
self
}
pub fn into_router(self) -> Result<AxumRouter, RouterError> {
if self.authenticator_was_explicit && matches!(self.policy, AuthPolicy::Public) {
return Err(RouterError::UnprotectedAuthenticator);
}
let endpoint = self.config.endpoint().as_str().to_owned();
let security = self.config.security().clone();
let max_concurrency = self.config.max_concurrency();
let state = Arc::new(RouterState {
config: self.config,
sessions: Arc::clone(&self.sessions),
handler: self.handler,
policy: self.policy,
authenticator: Arc::clone(&self.authenticator),
origin: self.origin,
rate_limiter: self.rate_limiter,
});
let methods: MethodRouter<Arc<RouterState<S, H>>> = axum::routing::post(post_request)
.get(open_stream)
.delete(terminate_session);
let mut router = AxumRouter::new().route(&endpoint, methods);
if let Some(metadata) = state.config.resource_metadata() {
router = router.route(
&metadata.resource().metadata_path(),
axum::routing::get(serve_document),
);
}
let router = router
.layer(axum::middleware::from_fn_with_state(
Arc::clone(&state),
enforce_deadline,
))
.layer(axum::middleware::from_fn_with_state(
Arc::clone(&state),
enforce_origin,
))
.layer(axum::middleware::from_fn_with_state(
Arc::clone(&state),
enforce_host,
))
.layer(axum::middleware::from_fn_with_state(
Arc::clone(&state),
enforce_rate_limit,
))
.with_state(state);
Ok(with_security(
hardening::apply_layers(router, max_concurrency),
&security,
)?)
}
}
pub fn with_security(
router: AxumRouter,
headers: &SecurityHeaders,
) -> Result<AxumRouter, InvalidHeaderValue> {
let layer = Arc::new(SecurityLayer::parse(headers)?);
Ok(router.layer(axum::middleware::from_fn_with_state(
layer,
apply_security_headers,
)))
}
async fn apply_security_headers(
State(policy): State<Arc<SecurityLayer>>,
request: Request<Body>,
next: Next,
) -> Response {
let mut response = next.run(request).await;
let out = response.headers_mut();
out.insert(
header::CONTENT_SECURITY_POLICY,
policy.content_security_policy.clone(),
);
out.insert(
header::STRICT_TRANSPORT_SECURITY,
policy.strict_transport_security.clone(),
);
out.insert(header::X_FRAME_OPTIONS, policy.x_frame_options.clone());
response
}
async fn post_request<S, H>(
State(state): State<Arc<RouterState<S, H>>>,
headers: HeaderMap,
body: Body,
) -> Response
where
S: EventSource,
H: RequestHandler<S>,
{
let context = match authenticate(state.authenticator.as_ref(), &state.policy, &headers).await {
Ok(context) => context,
Err(refusal) => return refuse_with_challenge(state.config.challenge().as_deref(), refusal),
};
let Ok(limit) = usize::try_from(state.config.max_body().bytes()) else {
return refuse_with_challenge(
state.config.challenge().as_deref(),
Refusal::Handler("the configured body limit does not fit this platform".to_owned()),
);
};
let body = match read_body(body, limit).await {
Ok(body) => body,
Err(refusal) => return refuse_with_challenge(state.config.challenge().as_deref(), refusal),
};
let request_id = request_id_of(&body);
let session = match Header::read(&headers, &SESSION_HEADER) {
Header::Absent => match state.handler.unsessioned(&body, &headers) {
Unsessioned::Establishing => SessionHandle::Establishing,
Unsessioned::Stateless => SessionHandle::Stateless,
Unsessioned::NeedsSession => {
return refuse_with_id_and_challenge(
state.config.challenge().as_deref(),
Refusal::MissingSession,
request_id.as_ref(),
);
}
},
Header::Value(_) | Header::Unreadable => match resolve_session(&state.sessions, &headers) {
Ok(session) => session,
Err(refusal) => {
return refuse_with_id_and_challenge(
state.config.challenge().as_deref(),
refusal,
request_id.as_ref(),
);
}
},
};
let protocol = match read_protocol(&headers) {
Ok(protocol) => protocol,
Err(_) if matches!(session, SessionHandle::Stateless) => None,
Err(refusal) => {
return refuse_with_id_and_challenge(
state.config.challenge().as_deref(),
refusal,
request_id.as_ref(),
);
}
};
let call = Inbound {
context: &context,
session,
body: &body,
protocol,
headers: &headers,
streaming: state.config.streaming(),
};
match state.handler.dispatch(call).await {
Ok(response) => stamp_media_type(response),
Err(refusal) => refuse_with_id_and_challenge(
state.config.challenge().as_deref(),
refusal,
request_id.as_ref(),
),
}
}
fn stamp_media_type(response: Response) -> Response {
if !response.status().is_success() {
return response;
}
let already_framed = response
.headers()
.get(header::CONTENT_TYPE)
.is_some_and(|value| {
let raw = value.as_bytes();
raw.starts_with(b"text/event-stream") || raw.starts_with(b"application/json")
});
if already_framed {
return response;
}
let mut response = response;
response.headers_mut().insert(
header::CONTENT_TYPE,
header::HeaderValue::from_static(JSON_MEDIA_TYPE),
);
response
}
async fn serve_document<S, H>(State(state): State<Arc<RouterState<S, H>>>) -> Response
where
S: EventSource,
H: RequestHandler<S>,
{
let Some(metadata) = state.config.resource_metadata() else {
return StatusCode::NOT_FOUND.into_response();
};
let Ok(body) = serde_json::to_vec(metadata) else {
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
};
(
StatusCode::OK,
[(header::CONTENT_TYPE, "application/json")],
body,
)
.into_response()
}
async fn open_stream<S, H>(
State(state): State<Arc<RouterState<S, H>>>,
headers: HeaderMap,
) -> Response
where
S: EventSource,
H: RequestHandler<S>,
{
if let Err(refusal) = authenticate(state.authenticator.as_ref(), &state.policy, &headers).await
{
return refuse_with_challenge(state.config.challenge().as_deref(), refusal);
}
let session = match resolve_session(&state.sessions, &headers) {
Ok(session) => session,
Err(refusal) => return refuse_with_challenge(state.config.challenge().as_deref(), refusal),
};
let SessionHandle::Live { session, .. } = session else {
return refuse_with_challenge(
state.config.challenge().as_deref(),
Refusal::Handler("the session lookup produced no session".to_owned()),
);
};
let cursor = match Header::read(&headers, &LAST_EVENT_ID) {
Header::Absent => None,
Header::Value(text) => Some(text),
Header::Unreadable => {
return refuse_with_challenge(
state.config.challenge().as_deref(),
Refusal::MalformedCursor,
);
}
};
match session
.tail()
.plan(cursor)
.into_response(state.config.keep_alive())
{
Ok(response) => response,
Err(error) => refuse_with_challenge(
state.config.challenge().as_deref(),
Refusal::Handler(error.to_string()),
),
}
}
async fn terminate_session<S, H>(
State(state): State<Arc<RouterState<S, H>>>,
headers: HeaderMap,
) -> Response
where
S: EventSource,
H: RequestHandler<S>,
{
if let Err(refusal) = authenticate(state.authenticator.as_ref(), &state.policy, &headers).await
{
return refuse_with_challenge(state.config.challenge().as_deref(), refusal);
}
let id = match session_id(&headers) {
Ok(Some(id)) => id,
Ok(None) => {
return refuse_with_challenge(
state.config.challenge().as_deref(),
Refusal::MissingSession,
);
}
Err(refusal) => return refuse_with_challenge(state.config.challenge().as_deref(), refusal),
};
match state.sessions.terminate(&id) {
Ok(Some(_terminated)) => StatusCode::OK.into_response(),
Ok(None) => {
refuse_with_challenge(state.config.challenge().as_deref(), Refusal::UnknownSession)
}
Err(error) => refuse_with_challenge(
state.config.challenge().as_deref(),
Refusal::Handler(error.to_string()),
),
}
}
fn session_id(headers: &HeaderMap) -> Result<Option<McpSessionId>, Refusal> {
match Header::read(headers, &SESSION_HEADER) {
Header::Absent => Ok(None),
Header::Unreadable => Err(Refusal::MalformedSessionId),
Header::Value(text) => text
.parse::<McpSessionId>()
.map(Some)
.map_err(|_| Refusal::MalformedSessionId),
}
}
fn resolve_session<S: EventSource>(
sessions: &SessionStore<S>,
headers: &HeaderMap,
) -> Result<SessionHandle<S>, Refusal> {
let Some(id) = session_id(headers)? else {
return Err(Refusal::MissingSession);
};
let found = sessions
.lookup(&id)
.map_err(|error| Refusal::Handler(error.to_string()))?;
found.map_or_else(
|| Err(Refusal::UnknownSession),
|session| Ok(SessionHandle::Live { id, session }),
)
}
fn read_protocol(headers: &HeaderMap) -> Result<Option<ProtocolVersion>, Refusal> {
let raw = match Header::read(headers, &PROTOCOL_HEADER) {
Header::Absent => return Ok(None),
Header::Unreadable => {
return Err(Refusal::ProtocolHeader {
reason: String::from("the value is not visible ASCII"),
});
}
Header::Value(value) => value,
};
raw.parse::<ProtocolVersion>()
.map(Some)
.map_err(|_| Refusal::ProtocolHeader {
reason: raw.to_owned(),
})
}
fn request_id_of(body: &[u8]) -> Option<WireId<'static>> {
let decoded = decode(body).ok()?;
let Frame::Single(envelope) = decoded else {
return None;
};
envelope.id().map(WireId::to_owned)
}
fn refuse_with_challenge(challenge: Option<&str>, refusal: Refusal) -> Response {
refuse_with_id_and_challenge(challenge, refusal, None)
}
fn refuse_with_id_and_challenge(
challenge: Option<&str>,
refusal: Refusal,
id: Option<&WireId<'_>>,
) -> Response {
let status = refusal.status();
let mut response = refusal.into_response_with_id(id);
if status == StatusCode::UNAUTHORIZED
&& let Some(challenge) = challenge
&& let Ok(value) = HeaderValue::from_str(challenge)
{
response
.headers_mut()
.insert(header::WWW_AUTHENTICATE, value);
}
response
}
#[cfg(test)]
mod tests;