Skip to main content

mcpr_core/proxy/
router.rs

1use crate::protocol as jsonrpc;
2use axum::{
3    body::Bytes,
4    http::{HeaderMap, Method, header},
5};
6
7/// Classified request type for type-separate dispatch.
8pub enum ClassifiedRequest {
9    /// OAuth callback relay page
10    OAuthCallback,
11    /// Widget HTML page: /widgets/{name}.html
12    WidgetHtml { name: String },
13    /// Widget list: /widgets or /widgets/
14    WidgetList,
15    /// Static widget asset (JS, CSS, images, fonts)
16    WidgetAsset,
17    /// MCP JSON-RPC POST with parsed body
18    McpPost { parsed: jsonrpc::ParsedBody },
19    /// MCP SSE GET (text/event-stream)
20    McpSse,
21    /// Everything else → forward to upstream
22    Passthrough,
23}
24
25/// Classify an incoming request for type-separate dispatch.
26pub fn classify(
27    method: &Method,
28    path: &str,
29    headers: &HeaderMap,
30    body: &Bytes,
31    has_widgets: bool,
32) -> ClassifiedRequest {
33    if *method == Method::GET && path == "/oauth/callback" {
34        return ClassifiedRequest::OAuthCallback;
35    }
36
37    if *method == Method::GET {
38        if let Some(name) = path
39            .strip_prefix("/widgets/")
40            .and_then(|s| s.strip_suffix(".html"))
41        {
42            return ClassifiedRequest::WidgetHtml {
43                name: name.to_string(),
44            };
45        }
46        if path == "/widgets" || path == "/widgets/" {
47            return ClassifiedRequest::WidgetList;
48        }
49    }
50
51    if *method == Method::GET && has_widgets && is_widget_asset(path, headers) {
52        return ClassifiedRequest::WidgetAsset;
53    }
54
55    if *method == Method::POST
56        && let Some(parsed) = parse_mcp_body(body)
57    {
58        return ClassifiedRequest::McpPost { parsed };
59    }
60
61    if *method == Method::GET && is_mcp_sse(headers) {
62        return ClassifiedRequest::McpSse;
63    }
64
65    ClassifiedRequest::Passthrough
66}
67
68/// Check if a POST body is a valid JSON-RPC 2.0 message (MCP request).
69fn parse_mcp_body(body: &Bytes) -> Option<jsonrpc::ParsedBody> {
70    jsonrpc::parse_body(body)
71}
72
73/// Check if a GET request is an MCP SSE call based on accept header.
74fn is_mcp_sse(headers: &HeaderMap) -> bool {
75    headers
76        .get(header::ACCEPT)
77        .and_then(|v| v.to_str().ok())
78        .map(|a| a.contains("text/event-stream"))
79        .unwrap_or(false)
80}
81
82/// Check if a request is for a static widget asset.
83/// Uses both file extension and Accept header to decide.
84fn is_widget_asset(path: &str, headers: &HeaderMap) -> bool {
85    let ext = path.rsplit('.').next().unwrap_or("");
86    if matches!(
87        ext,
88        "js" | "mjs"
89            | "css"
90            | "html"
91            | "svg"
92            | "png"
93            | "jpg"
94            | "jpeg"
95            | "gif"
96            | "ico"
97            | "woff"
98            | "woff2"
99            | "ttf"
100            | "eot"
101            | "map"
102            | "webp"
103    ) {
104        return true;
105    }
106
107    if let Some(accept) = headers.get(header::ACCEPT).and_then(|v| v.to_str().ok())
108        && (accept.contains("text/html")
109            || accept.contains("text/css")
110            || accept.contains("image/")
111            || accept.contains("font/")
112            || accept.contains("application/javascript"))
113    {
114        return true;
115    }
116
117    false
118}
119
120#[cfg(test)]
121#[allow(non_snake_case)]
122mod tests {
123    use super::*;
124
125    // ── parse_mcp_body (JSON-RPC 2.0 detection) ──
126
127    #[test]
128    fn parse_mcp_body__jsonrpc_request() {
129        let body = br#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#;
130        let parsed = parse_mcp_body(&Bytes::from_static(body));
131        assert!(parsed.is_some());
132        let p = parsed.unwrap();
133        assert_eq!(p.method_str(), "tools/list");
134        assert!(!p.is_batch);
135    }
136
137    #[test]
138    fn parse_mcp_body__jsonrpc_notification() {
139        let body = br#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#;
140        let parsed = parse_mcp_body(&Bytes::from_static(body));
141        assert!(parsed.is_some());
142        assert!(parsed.unwrap().is_notification_only());
143    }
144
145    #[test]
146    fn parse_mcp_body__jsonrpc_batch() {
147        let body = br#"[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2,"method":"resources/list"}]"#;
148        let parsed = parse_mcp_body(&Bytes::from_static(body));
149        assert!(parsed.is_some());
150        assert!(parsed.unwrap().is_batch);
151    }
152
153    #[test]
154    fn parse_mcp_body__rejects_oauth_register() {
155        let body = br#"{"client_name":"My App","redirect_uris":["https://example.com/cb"]}"#;
156        assert!(parse_mcp_body(&Bytes::from_static(body)).is_none());
157    }
158
159    #[test]
160    fn parse_mcp_body__rejects_form_encoded() {
161        let body = b"grant_type=client_credentials&client_id=abc";
162        assert!(parse_mcp_body(&Bytes::from_static(body)).is_none());
163    }
164
165    #[test]
166    fn parse_mcp_body__rejects_empty() {
167        assert!(parse_mcp_body(&Bytes::new()).is_none());
168    }
169
170    #[test]
171    fn parse_mcp_body__rejects_wrong_version() {
172        let body = br#"{"jsonrpc":"1.0","id":1,"method":"test"}"#;
173        assert!(parse_mcp_body(&Bytes::from_static(body)).is_none());
174    }
175
176    // ── is_mcp_sse ──
177
178    #[test]
179    fn is_mcp_sse__accepts_event_stream() {
180        let mut headers = HeaderMap::new();
181        headers.insert(header::ACCEPT, "text/event-stream".parse().unwrap());
182        assert!(is_mcp_sse(&headers));
183    }
184
185    #[test]
186    fn is_mcp_sse__rejects_html() {
187        let mut headers = HeaderMap::new();
188        headers.insert(header::ACCEPT, "text/html".parse().unwrap());
189        assert!(!is_mcp_sse(&headers));
190    }
191
192    #[test]
193    fn is_mcp_sse__rejects_no_accept() {
194        let headers = HeaderMap::new();
195        assert!(!is_mcp_sse(&headers));
196    }
197
198    // ── is_widget_asset ──
199
200    #[test]
201    fn is_widget_asset__js_extension() {
202        let headers = HeaderMap::new();
203        assert!(is_widget_asset("/assets/main.js", &headers));
204    }
205
206    #[test]
207    fn is_widget_asset__css_extension() {
208        let headers = HeaderMap::new();
209        assert!(is_widget_asset("/styles/app.css", &headers));
210    }
211
212    #[test]
213    fn is_widget_asset__woff2_extension() {
214        let headers = HeaderMap::new();
215        assert!(is_widget_asset("/fonts/inter.woff2", &headers));
216    }
217
218    #[test]
219    fn is_widget_asset__svg_extension() {
220        let headers = HeaderMap::new();
221        assert!(is_widget_asset("/icons/logo.svg", &headers));
222    }
223
224    #[test]
225    fn is_widget_asset__accept_html() {
226        let mut headers = HeaderMap::new();
227        headers.insert(header::ACCEPT, "text/html".parse().unwrap());
228        assert!(is_widget_asset("/some-path", &headers));
229    }
230
231    #[test]
232    fn is_widget_asset__accept_image() {
233        let mut headers = HeaderMap::new();
234        headers.insert(header::ACCEPT, "image/png".parse().unwrap());
235        assert!(is_widget_asset("/logo", &headers));
236    }
237
238    #[test]
239    fn is_widget_asset__accept_javascript() {
240        let mut headers = HeaderMap::new();
241        headers.insert(header::ACCEPT, "application/javascript".parse().unwrap());
242        assert!(is_widget_asset("/bundle", &headers));
243    }
244
245    #[test]
246    fn is_widget_asset__rejects_well_known() {
247        let headers = HeaderMap::new();
248        assert!(!is_widget_asset(
249            "/.well-known/oauth-authorization-server",
250            &headers
251        ));
252    }
253
254    #[test]
255    fn is_widget_asset__rejects_mcp() {
256        let headers = HeaderMap::new();
257        assert!(!is_widget_asset("/mcp", &headers));
258    }
259
260    #[test]
261    fn is_widget_asset__rejects_token() {
262        let headers = HeaderMap::new();
263        assert!(!is_widget_asset("/token", &headers));
264    }
265
266    #[test]
267    fn is_widget_asset__rejects_authorize() {
268        let headers = HeaderMap::new();
269        assert!(!is_widget_asset("/authorize", &headers));
270    }
271
272    #[test]
273    fn is_widget_asset__rejects_register() {
274        let headers = HeaderMap::new();
275        assert!(!is_widget_asset("/register", &headers));
276    }
277
278    #[test]
279    fn is_widget_asset__rejects_json_accept() {
280        let mut headers = HeaderMap::new();
281        headers.insert(header::ACCEPT, "application/json".parse().unwrap());
282        assert!(!is_widget_asset("/some-path", &headers));
283    }
284
285    #[test]
286    fn is_widget_asset__rejects_sse_accept() {
287        let mut headers = HeaderMap::new();
288        headers.insert(header::ACCEPT, "text/event-stream".parse().unwrap());
289        assert!(!is_widget_asset("/mcp", &headers));
290    }
291}