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