Skip to main content

dynoxide/mcp/
mod.rs

1//! MCP (Model Context Protocol) server for Dynoxide.
2//!
3//! Exposes DynamoDB operations as MCP tools, allowing coding agents to
4//! interact with the local DynamoDB emulator through structured tool calls.
5
6mod auth;
7mod errors;
8mod server;
9
10pub use auth::{
11    AuthError, AuthMode, ResolvedAuth, first_run_message, is_loopback_host, resolve_auth,
12};
13pub use server::{McpConfig, McpServer};
14
15use crate::Database;
16use std::sync::Arc;
17
18/// Start the MCP server over stdio transport.
19///
20/// This blocks until the client disconnects. All logging goes to stderr;
21/// stdout is reserved for the JSON-RPC transport.
22pub async fn serve_stdio(
23    db: Database,
24    config: McpConfig,
25) -> Result<(), Box<dyn std::error::Error>> {
26    use rmcp::{ServiceExt, transport::stdio};
27
28    let server = McpServer::with_config(Arc::new(db), config)
29        .serve(stdio())
30        .await
31        .map_err(|e| format!("MCP server error: {e}"))?;
32
33    server.waiting().await?;
34    Ok(())
35}
36
37/// Transport-level options for the MCP HTTP server.
38///
39/// Distinct from [`McpConfig`], which is per-session server behaviour cloned
40/// into every connection. These are resolved once at startup: where to bind,
41/// the bearer-token auth mode, and any operator-added `Host`/`Origin` allowlist
42/// entries beyond loopback.
43#[derive(Clone, Debug)]
44pub struct HttpOptions {
45    pub host: String,
46    pub port: u16,
47    pub auth: AuthMode,
48    /// Extra hosts to accept beyond the loopback default. Each entry also adds
49    /// a matching `http://<host>` origin. Empty preserves loopback-only.
50    pub extra_allowed_hosts: Vec<String>,
51}
52
53/// Wrap a bare IPv6 literal in brackets for a URL authority. Leaves IPv4,
54/// hostnames, and already-bracketed literals untouched.
55fn bracket_ipv6(host: &str) -> String {
56    if host.contains(':') && !host.starts_with('[') {
57        format!("[{host}]")
58    } else {
59        host.to_string()
60    }
61}
62
63/// Format a `host:port` bind address, bracketing bare IPv6 literals.
64fn format_bind_addr(host: &str, port: u16) -> String {
65    format!("{}:{}", bracket_ipv6(host), port)
66}
67
68/// Start the MCP server over Streamable HTTP transport.
69///
70/// Binds to `opts.host:opts.port` and serves MCP at `/mcp`, behind the
71/// bearer-token auth layer.
72pub async fn serve_http(
73    db: Database,
74    opts: HttpOptions,
75    config: McpConfig,
76) -> Result<(), Box<dyn std::error::Error>> {
77    serve_http_with_shutdown(db, opts, config, None).await
78}
79
80/// Start the MCP server over Streamable HTTP with an external shutdown signal.
81pub async fn serve_http_with_shutdown(
82    db: Database,
83    opts: HttpOptions,
84    config: McpConfig,
85    shutdown: Option<tokio_util::sync::CancellationToken>,
86) -> Result<(), Box<dyn std::error::Error>> {
87    use rmcp::transport::streamable_http_server::{
88        StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
89    };
90
91    let ct = shutdown.unwrap_or_default();
92    let db = Arc::new(db);
93
94    // Loopback defaults, plus any operator-added hosts (for non-loopback binds
95    // reached by name). Each added host gets a matching http:// origin.
96    let mut allowed_hosts: Vec<String> = vec!["localhost".into(), "127.0.0.1".into(), "::1".into()];
97    let mut allowed_origins: Vec<String> =
98        vec!["http://localhost".into(), "http://127.0.0.1".into()];
99    for host in &opts.extra_allowed_hosts {
100        allowed_hosts.push(host.clone());
101        allowed_origins.push(format!("http://{}", bracket_ipv6(host)));
102    }
103
104    // load-bearing: rmcp's config struct is #[non_exhaustive], so struct-literal
105    // init does not compile. Keep the field-reassign block.
106    #[allow(clippy::field_reassign_with_default)]
107    let http_config = {
108        let mut c = StreamableHttpServerConfig::default();
109        c.stateful_mode = false;
110        c.json_response = true;
111        c.sse_keep_alive = None;
112        c.cancellation_token = ct.clone();
113        // Explicit DNS rebinding defences. Stating the lists protects against an
114        // rmcp default flip. Native clients pass because rmcp skips Origin
115        // validation when the header is absent (rmcp 1.6.0 tower.rs:385-387).
116        // Bearer-token auth (the .layer below) is the primary control once the
117        // bind widens; this allowlist is defense-in-depth for browser origins.
118        c.allowed_hosts = allowed_hosts;
119        c.allowed_origins = allowed_origins;
120        c
121    };
122
123    let service: StreamableHttpService<McpServer, LocalSessionManager> = StreamableHttpService::new(
124        {
125            let db = db.clone();
126            let config = config.clone();
127            move || Ok(McpServer::with_config(db.clone(), config.clone()))
128        },
129        Default::default(),
130        http_config,
131    );
132
133    // Auth runs outside rmcp's Host/Origin checks: unauthenticated callers get
134    // 401 regardless of Host; a token holder spoofing Host still hits 403.
135    let router = axum::Router::new().nest_service("/mcp", service).layer(
136        axum::middleware::from_fn_with_state(opts.auth.clone(), auth::enforce),
137    );
138    let addr = format_bind_addr(&opts.host, opts.port);
139    let listener = tokio::net::TcpListener::bind(&addr).await?;
140    let local_addr = listener.local_addr()?;
141
142    eprintln!("MCP HTTP server listening on http://{local_addr}/mcp");
143
144    // Don't use with_graceful_shutdown for the MCP transport: persistent
145    // Streamable-HTTP sessions held open by MCP clients (Claude Code, Cursor,
146    // etc.) don't close on their own, so the drain phase can hang
147    // indefinitely. Race the cancellation against the serve future and drop
148    // the serve future on cancel; the listener closes, this function returns,
149    // and the surrounding tokio::join! in main.rs proceeds so the process
150    // exits.
151    use std::future::IntoFuture;
152    let serve_fut = axum::serve(listener, router).into_future();
153    tokio::select! {
154        res = serve_fut => res?,
155        _ = ct.cancelled_owned() => {}
156    }
157
158    Ok(())
159}