rust_web_server/mcp/
gateway.rs

1//! A minimal MCP gateway: one endpoint in front of several upstream MCP
2//! servers.
3//!
4//! [`McpGateway::build`] connects to each [`Upstream`], lists its tools and
5//! registers each one as `<prefix>__<tool>` on an ordinary [`McpServer`];
6//! calling it forwards the call to that upstream. The gateway's caller
7//! authenticates to the gateway (put [`McpServer::require_bearer_tokens`] or
8//! [`McpServer::require_oauth`] on the server `build` returns); upstreams
9//! get only the headers configured for them — the caller's token is never
10//! passed on. [`McpGateway::allow`] decides which caller (the `sub` of
11//! [`McpContext::auth_claims`](super::McpContext::auth_claims)) sees which tools, and nothing is allowed
12//! unless a rule says so.
13//!
14//! Minimal on purpose: tools are listed once, at `build`; upstream resources,
15//! prompts and notifications are not federated; each upstream speaks the
16//! protocol version it negotiated with the gateway.
17
18use std::sync::atomic::{AtomicU64, Ordering};
19use std::sync::{Arc, Mutex};
20
21use super::{json_escape, json_rpc, McpContent, McpServer, ToolAnnotations};
22use crate::http_client::Client;
23
24/// The protocol version the gateway asks upstreams for.
25const PROTOCOL_VERSION: &str = "2025-11-25";
26
27/// Separates an upstream's prefix from its tool names.
28pub const SEPARATOR: &str = "__";
29
30/// `tools/list` pages followed per upstream before giving up.
31const MAX_LIST_PAGES: usize = 100;
32
33/// An upstream MCP server (Streamable HTTP).
34#[derive(Clone, Debug)]
35pub struct Upstream {
36    prefix: String,
37    url: String,
38    headers: Vec<(String, String)>,
39    timeout_ms: u64,
40}
41
42impl Upstream {
43    /// The server at `url` (its MCP endpoint, e.g. `https://mcp.example.com/mcp`),
44    /// whose tools appear as `<prefix>__<tool>`. `prefix` is letters, digits,
45    /// `-` and single `_`s.
46    pub fn new(prefix: impl Into<String>, url: impl Into<String>) -> Self {
47        Upstream { prefix: prefix.into(), url: url.into(), headers: Vec::new(), timeout_ms: 30_000 }
48    }
49
50    /// Send this header to the upstream on every request — its credentials,
51    /// e.g. `.header("Authorization", "Bearer <upstream token>")`.
52    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
53        self.headers.push((name.into(), value.into()));
54        self
55    }
56
57    /// Give up on an upstream request after this long (default 30 s).
58    pub fn timeout_ms(mut self, timeout_ms: u64) -> Self {
59        self.timeout_ms = timeout_ms;
60        self
61    }
62}
63
64/// Builds a gateway [`McpServer`]; see the [module docs](self).
65pub struct McpGateway {
66    name: String,
67    version: String,
68    upstreams: Vec<Upstream>,
69    rules: Vec<(String, Vec<String>)>,
70}
71
72impl McpGateway {
73    /// A gateway announcing itself as `name`/`version`.
74    pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
75        McpGateway { name: name.into(), version: version.into(), upstreams: Vec::new(), rules: Vec::new() }
76    }
77
78    /// Federate `upstream`'s tools.
79    pub fn upstream(mut self, upstream: Upstream) -> Self {
80        self.upstreams.push(upstream);
81        self
82    }
83
84    /// Let `identity` see and call the tools matching `tools`: a full name
85    /// (`github__search_issues`), a prefix ending in `*` (`github__*`), or
86    /// `*` for all. `identity` is the caller's `sub` claim; `*` means any
87    /// identified caller. Rules add up; a caller no rule names gets no tools.
88    pub fn allow<I, S>(mut self, identity: impl Into<String>, tools: I) -> Self
89    where
90        I: IntoIterator<Item = S>,
91        S: Into<String>,
92    {
93        self.rules.push((identity.into(), tools.into_iter().map(Into::into).collect()));
94        self
95    }
96
97    /// Connect to every upstream, list its tools and return the gateway.
98    /// Fails, naming the upstream, if one can't be reached or answers
99    /// something that isn't MCP. Add authentication to the result.
100    pub fn build(self) -> Result<McpServer, String> {
101        let mut prefixes: Vec<&str> = Vec::new();
102        for upstream in &self.upstreams {
103            if !valid_prefix(&upstream.prefix) {
104                return Err(format!("upstream prefix \"{}\": use letters, digits, '-' and single '_'", upstream.prefix));
105            }
106            if prefixes.contains(&upstream.prefix.as_str()) {
107                return Err(format!("upstream prefix \"{}\" is used twice", upstream.prefix));
108            }
109            prefixes.push(&upstream.prefix);
110        }
111
112        let rules = Arc::new(self.rules);
113        let mut server = McpServer::new(self.name, self.version).tool_filter(move |ctx, tool| {
114            let identity = ctx.auth_claims.as_deref().and_then(|claims| top_level_str(claims, "sub"));
115            identity.is_some_and(|identity| allowed(&rules, &identity, tool))
116        });
117        for upstream in self.upstreams {
118            let connection = Arc::new(Connection::new(upstream));
119            for tool in connection.list_tools()? {
120                server = register(server, &connection, tool);
121            }
122        }
123        Ok(server)
124    }
125}
126
127/// One upstream tool, as `tools/list` described it.
128struct UpstreamTool {
129    name: String,
130    description: String,
131    title: Option<String>,
132    input_schema: String,
133    output_schema: Option<String>,
134    annotations: Option<ToolAnnotations>,
135}
136
137fn register(server: McpServer, connection: &Arc<Connection>, tool: UpstreamTool) -> McpServer {
138    let exposed = format!("{}{SEPARATOR}{}", connection.upstream.prefix, tool.name);
139    let forward = Arc::clone(connection);
140    let upstream_name = tool.name.clone();
141    let mut server = server.tool_with_context(&exposed, &tool.description, &tool.input_schema, move |_ctx, args| {
142        forward.call_tool(&upstream_name, args).map(McpContent::raw_result)
143    });
144    if let Some(title) = tool.title {
145        server = server.tool_title(&exposed, title);
146    }
147    if let Some(schema) = tool.output_schema {
148        server = server.tool_output_schema(&exposed, &schema);
149    }
150    if let Some(annotations) = tool.annotations {
151        server = server.tool_annotations(&exposed, annotations);
152    }
153    server
154}
155
156fn valid_prefix(prefix: &str) -> bool {
157    !prefix.is_empty()
158        && !prefix.contains(SEPARATOR)
159        && !prefix.ends_with('_')
160        && prefix.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
161}
162
163/// Does a rule for `identity` (or `*`) match `tool`?
164fn allowed(rules: &[(String, Vec<String>)], identity: &str, tool: &str) -> bool {
165    rules
166        .iter()
167        .filter(|(who, _)| who == identity || who == "*")
168        .flat_map(|(_, patterns)| patterns)
169        .any(|pattern| match pattern.strip_suffix('*') {
170            Some(prefix) => tool.starts_with(prefix),
171            None => pattern == tool,
172        })
173}
174
175/// The string member `key` of the top level of JSON object `json`.
176fn top_level_str(json: &str, key: &str) -> Option<String> {
177    let raw = top_level_raw(json, key)?;
178    raw.starts_with('"').then(|| json_rpc::extract_str(&format!(r#"{{"v":{raw}}}"#), "v")).flatten()
179}
180
181/// The raw JSON of member `key` of the top level of JSON object `json`.
182fn top_level_raw(json: &str, key: &str) -> Option<String> {
183    json_rpc::object_members(json)?.into_iter().find(|(k, _)| k == key).map(|(_, v)| v)
184}
185
186/// The gateway's session with one upstream.
187struct Connection {
188    upstream: Upstream,
189    client: Client,
190    /// `(Mcp-Session-Id, negotiated protocol version)`, once initialized.
191    session: Mutex<Option<Session>>,
192    next_id: AtomicU64,
193}
194
195#[derive(Clone)]
196struct Session {
197    id: Option<String>,
198    version: String,
199}
200
201impl Connection {
202    fn new(upstream: Upstream) -> Self {
203        let client = Client::new().timeout_ms(upstream.timeout_ms).max_redirects(0);
204        Connection { upstream, client, session: Mutex::new(None), next_id: AtomicU64::new(1) }
205    }
206
207    fn name(&self) -> String {
208        format!("upstream {} ({})", self.upstream.prefix, self.upstream.url)
209    }
210
211    /// Every tool, following `nextCursor`.
212    fn list_tools(&self) -> Result<Vec<UpstreamTool>, String> {
213        let mut tools = Vec::new();
214        let mut cursor: Option<String> = None;
215        for _ in 0..MAX_LIST_PAGES {
216            let params = match &cursor {
217                Some(cursor) => format!(r#"{{"cursor":"{}"}}"#, json_escape(cursor)),
218                None => "{}".to_string(),
219            };
220            let result = self.request("tools/list", &params).map_err(|e| format!("{}: {e}", self.name()))?;
221            let list = top_level_raw(&result, "tools").ok_or_else(|| format!("{}: tools/list without tools", self.name()))?;
222            for tool in json_rpc::split_array_elements(&list) {
223                tools.push(parse_tool(&tool).ok_or_else(|| format!("{}: malformed tool {tool}", self.name()))?);
224            }
225            cursor = top_level_str(&result, "nextCursor");
226            if cursor.is_none() {
227                return Ok(tools);
228            }
229        }
230        Err(format!("{}: more than {MAX_LIST_PAGES} pages of tools", self.name()))
231    }
232
233    /// Calls `name` upstream: the raw `tools/call` result, or why it failed.
234    fn call_tool(&self, name: &str, arguments: &str) -> Result<String, String> {
235        let params = format!(r#"{{"name":"{}","arguments":{arguments}}}"#, json_escape(name));
236        self.request("tools/call", &params).map_err(|e| format!("{}: {e}", self.upstream.prefix))
237    }
238
239    /// Sends one request in the current session (initializing first if
240    /// needed, and once more if the upstream has forgotten the session).
241    fn request(&self, method: &str, params: &str) -> Result<String, String> {
242        let session = self.session()?;
243        match self.send(&session, method, params)? {
244            Reply::Result(result) => Ok(result),
245            Reply::SessionGone => {
246                self.forget(&session);
247                let session = self.session()?;
248                match self.send(&session, method, params)? {
249                    Reply::Result(result) => Ok(result),
250                    Reply::SessionGone => Err("session not found twice in a row".to_string()),
251                }
252            }
253        }
254    }
255
256    fn session(&self) -> Result<Session, String> {
257        let mut guard = self.session.lock().map_err(|_| "session lock poisoned".to_string())?;
258        if let Some(session) = guard.as_ref() {
259            return Ok(session.clone());
260        }
261        let session = self.initialize()?;
262        *guard = Some(session.clone());
263        Ok(session)
264    }
265
266    fn forget(&self, stale: &Session) {
267        if let Ok(mut guard) = self.session.lock() {
268            if guard.as_ref().is_some_and(|current| current.id == stale.id) {
269                *guard = None;
270            }
271        }
272    }
273
274    fn initialize(&self) -> Result<Session, String> {
275        let params = format!(
276            r#"{{"protocolVersion":"{PROTOCOL_VERSION}","capabilities":{{}},"clientInfo":{{"name":"rws-mcp-gateway","version":"{}"}}}}"#,
277            env!("CARGO_PKG_VERSION"),
278        );
279        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
280        let response = self.post(None, PROTOCOL_VERSION, &envelope(id, "initialize", &params))?;
281        if !(200..300).contains(&response.status()) {
282            return Err(format!("initialize answered HTTP {}", response.status()));
283        }
284        let session_id = response.header("mcp-session-id").map(str::to_string);
285        let result = parse_reply(response.header("content-type").unwrap_or(""), response.bytes(), id)?;
286        let version = top_level_str(&result, "protocolVersion").unwrap_or_else(|| PROTOCOL_VERSION.to_string());
287        let session = Session { id: session_id, version };
288        let initialized = r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#;
289        self.post(session.id.as_deref(), &session.version, initialized)?;
290        Ok(session)
291    }
292
293    fn send(&self, session: &Session, method: &str, params: &str) -> Result<Reply, String> {
294        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
295        let response = self.post(session.id.as_deref(), &session.version, &envelope(id, method, params))?;
296        if response.status() == 404 && session.id.is_some() {
297            return Ok(Reply::SessionGone);
298        }
299        if !(200..300).contains(&response.status()) {
300            return Err(format!("answered HTTP {}", response.status()));
301        }
302        parse_reply(response.header("content-type").unwrap_or(""), response.bytes(), id).map(Reply::Result)
303    }
304
305    /// One POST to the upstream, with its own headers only.
306    fn post(&self, session_id: Option<&str>, version: &str, body: &str) -> Result<crate::http_client::Response, String> {
307        let mut request = self
308            .client
309            .post(&self.upstream.url)
310            .header("Content-Type", "application/json")
311            .header("Accept", "application/json, text/event-stream")
312            .header("MCP-Protocol-Version", version);
313        if let Some(id) = session_id {
314            request = request.header("Mcp-Session-Id", id);
315        }
316        for (name, value) in &self.upstream.headers {
317            request = request.header(name, value);
318        }
319        request.body(body.as_bytes().to_vec()).send().map_err(|e| format!("unreachable: {e}"))
320    }
321}
322
323enum Reply {
324    Result(String),
325    SessionGone,
326}
327
328fn envelope(id: u64, method: &str, params: &str) -> String {
329    format!(r#"{{"jsonrpc":"2.0","id":{id},"method":"{method}","params":{params}}}"#)
330}
331
332fn parse_tool(json: &str) -> Option<UpstreamTool> {
333    let members = json_rpc::object_members(json)?;
334    let raw = |key: &str| members.iter().find(|(k, _)| k == key).map(|(_, v)| v.clone());
335    let string = |key: &str| top_level_str(json, key);
336    let name = string("name")?;
337    let input_schema = raw("inputSchema").filter(|schema| schema.starts_with('{'))?;
338    let annotations = raw("annotations").filter(|a| a.starts_with('{')).map(|a| {
339        let hint = |key: &str| top_level_raw(&a, key).and_then(|v| v.parse::<bool>().ok());
340        ToolAnnotations {
341            read_only_hint: hint("readOnlyHint"),
342            destructive_hint: hint("destructiveHint"),
343            idempotent_hint: hint("idempotentHint"),
344            open_world_hint: hint("openWorldHint"),
345        }
346    });
347    Some(UpstreamTool {
348        name,
349        description: string("description").unwrap_or_default(),
350        title: string("title"),
351        input_schema,
352        output_schema: raw("outputSchema").filter(|schema| schema.starts_with('{')),
353        annotations,
354    })
355}
356
357/// The `result` of the JSON-RPC response with `id` in an upstream reply —
358/// a JSON body, or an SSE stream whose events may include notifications and
359/// requests before the response. A JSON-RPC `error` becomes its message.
360#[doc(hidden)]
361pub fn parse_reply(content_type: &str, body: &[u8], id: u64) -> Result<String, String> {
362    let body = std::str::from_utf8(body).map_err(|_| "reply is not UTF-8".to_string())?;
363    let media_type = content_type.split(';').next().unwrap_or("").trim();
364    let id = id.to_string();
365    if media_type.eq_ignore_ascii_case("text/event-stream") {
366        let mut data = String::new();
367        for line in body.lines().chain(std::iter::once("")) {
368            if let Some(value) = line.strip_prefix("data:") {
369                if !data.is_empty() {
370                    data.push('\n');
371                }
372                data.push_str(value.strip_prefix(' ').unwrap_or(value));
373            } else if line.is_empty() && !data.is_empty() {
374                let event = std::mem::take(&mut data);
375                if top_level_raw(&event, "id").as_deref() == Some(id.as_str()) && top_level_raw(&event, "method").is_none() {
376                    return response_result(&event);
377                }
378            }
379        }
380        Err("the event stream ended without a response".to_string())
381    } else {
382        if top_level_raw(body, "id").as_deref() != Some(id.as_str()) {
383            return Err("reply is not the response to this request".to_string());
384        }
385        response_result(body)
386    }
387}
388
389fn response_result(response: &str) -> Result<String, String> {
390    if let Some(result) = top_level_raw(response, "result") {
391        return Ok(result);
392    }
393    let error = top_level_raw(response, "error").ok_or_else(|| "response has neither result nor error".to_string())?;
394    Err(top_level_str(&error, "message").unwrap_or(error))
395}
396
397/// Entry points for `fuzz_api`.
398#[doc(hidden)]
399pub mod fuzzing {
400    /// The tools a `tools/list` result describes.
401    pub fn tools(result: &str) -> usize {
402        let Some(list) = super::top_level_raw(result, "tools") else { return 0 };
403        super::json_rpc::split_array_elements(&list).iter().filter_map(|tool| super::parse_tool(tool)).count()
404    }
405}
406
407#[cfg(test)]
408mod tests;
rust_web_server/mcp/gateway.rs

rust_web_server/mcp/
gateway.rs
