194c44c
pub mod catalog;
pub mod error;
pub mod settings;
pub mod space_token;
mod auth;
mod net;
mod server;
mod services;
mod tools;
pub use server::McpServer;
pub use services::{McpServices, SharedServices};
use std::net::IpAddr;
use std::sync::Arc;
use axum::Router;
use axum::middleware;
use export::AppState;
use ipnet::IpNet;
use rmcp::transport::StreamableHttpServerConfig;
use rmcp::transport::StreamableHttpService;
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use sha2::{Digest, Sha256};
use uuid::Uuid;
struct InternalKey {
name: String,
hash: String,
}
pub(crate) struct RuntimeInner {
pub(crate) services: SharedServices,
internal_keys: Vec<InternalKey>,
internal_cidrs: Vec<IpNet>,
trusted_proxies: Vec<IpNet>,
system_uid: Option<i64>,
}
#[derive(Clone)]
pub(crate) struct McpRuntime(Arc<RuntimeInner>);
impl McpRuntime {
pub(crate) fn services(&self) -> &SharedServices {
&self.0.services
}
pub(crate) fn system_uid(&self) -> Option<i64> {
self.0.system_uid
}
pub(crate) fn trusted_proxies(&self) -> &[IpNet] {
&self.0.trusted_proxies
}
pub(crate) fn match_internal_key(&self, token: &str) -> Option<&str> {
let hash = sha256_hex(token.as_bytes());
self.0
.internal_keys
.iter()
.find(|key| constant_time_eq(hash.as_bytes(), key.hash.as_bytes()))
.map(|key| key.name.as_str())
}
pub(crate) fn is_internal_ip(&self, ip: Option<IpAddr>) -> bool {
ip.is_some_and(|ip| self.0.internal_cidrs.iter().any(|net| net.contains(&ip)))
}
}
pub async fn router(state: AppState) -> Router {
let config = state.config.mcp.clone();
if !config.enabled() {
return Router::new();
}
let system_uid = match agent::bot::resolve_bot_user(&state, &state.config.ai)
.await
{
Ok(uid) => {
export::authz::set_system_uid(uid);
Some(uid)
}
Err(error) => {
tracing::warn!(%error, "failed to resolve the platform bot user for mcp internal keys");
None
}
};
let internal_keys = config
.internal_keys
.iter()
.filter(|key| key.enabled())
.map(|key| InternalKey {
name: key.name.clone(),
hash: key
.token_hash
.as_ref()
.map(|hash| hash.trim().to_lowercase())
.unwrap_or_else(|| sha256_hex(key.token.as_deref().unwrap_or_default().as_bytes())),
})
.collect();
let runtime = McpRuntime(Arc::new(RuntimeInner {
services: Arc::new(McpServices::new(state)),
internal_keys,
internal_cidrs: config.parsed_internal_cidrs(),
trusted_proxies: config.parsed_trusted_proxies(),
system_uid,
}));
let mut http_config = StreamableHttpServerConfig::default()
.with_legacy_session_mode(false)
.with_json_response(true);
if !config.allowed_hosts.is_empty() {
http_config = http_config.with_allowed_hosts(config.allowed_hosts.clone());
}
let services = runtime.services().clone();
let mcp_service = StreamableHttpService::new(
move || Ok(McpServer::new(services.clone())),
LocalSessionManager::default().into(),
http_config,
);
tracing::info!(
internal_keys = runtime.0.internal_keys.len(),
internal_cidrs = runtime.0.internal_cidrs.len(),
"mcp endpoint ready"
);
Router::new()
.fallback_service(mcp_service)
.layer(middleware::from_fn_with_state(
runtime,
auth::auth_middleware,
))
}
pub(crate) fn sha256_hex(data: &[u8]) -> String {
let digest = Sha256::digest(data);
digest.iter().map(|byte| format!("{byte:02x}")).collect()
}
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
if a.len() != b.len() {
return false;
}
let mut diff: u8 = 0;
for (x, y) in a.iter().zip(b.iter()) {
diff |= x ^ y;
}
diff == 0
}