Skip to main content

mcpr_core/
router.rs

1use axum::{
2    body::Bytes,
3    http::{HeaderMap, Method, header},
4};
5use mcpr_protocol as jsonrpc;
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)]
121mod tests {
122    use super::*;
123
124    // ── parse_mcp_body (JSON-RPC 2.0 detection) ──
125
126    #[test]
127    fn detect_mcp_jsonrpc_request() {
128        let body = br#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#;
129        let parsed = parse_mcp_body(&Bytes::from_static(body));
130        assert!(parsed.is_some());
131        let p = parsed.unwrap();
132        assert_eq!(p.method_str(), "tools/list");
133        assert!(!p.is_batch);
134    }
135
136    #[test]
137    fn detect_mcp_jsonrpc_notification() {
138        let body = br#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#;
139        let parsed = parse_mcp_body(&Bytes::from_static(body));
140        assert!(parsed.is_some());
141        assert!(parsed.unwrap().is_notification_only());
142    }
143
144    #[test]
145    fn detect_mcp_jsonrpc_batch() {
146        let body = br#"[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2,"method":"resources/list"}]"#;
147        let parsed = parse_mcp_body(&Bytes::from_static(body));
148        assert!(parsed.is_some());
149        assert!(parsed.unwrap().is_batch);
150    }
151
152    #[test]
153    fn reject_oauth_register_json() {
154        let body = br#"{"client_name":"My App","redirect_uris":["https://example.com/cb"]}"#;
155        assert!(parse_mcp_body(&Bytes::from_static(body)).is_none());
156    }
157
158    #[test]
159    fn reject_form_encoded() {
160        let body = b"grant_type=client_credentials&client_id=abc";
161        assert!(parse_mcp_body(&Bytes::from_static(body)).is_none());
162    }
163
164    #[test]
165    fn reject_empty_body() {
166        assert!(parse_mcp_body(&Bytes::new()).is_none());
167    }
168
169    #[test]
170    fn reject_wrong_jsonrpc_version() {
171        let body = br#"{"jsonrpc":"1.0","id":1,"method":"test"}"#;
172        assert!(parse_mcp_body(&Bytes::from_static(body)).is_none());
173    }
174
175    // ── is_mcp_sse ──
176
177    #[test]
178    fn is_mcp_sse_accept() {
179        let mut headers = HeaderMap::new();
180        headers.insert(header::ACCEPT, "text/event-stream".parse().unwrap());
181        assert!(is_mcp_sse(&headers));
182    }
183
184    #[test]
185    fn is_not_mcp_sse_html() {
186        let mut headers = HeaderMap::new();
187        headers.insert(header::ACCEPT, "text/html".parse().unwrap());
188        assert!(!is_mcp_sse(&headers));
189    }
190
191    #[test]
192    fn is_not_mcp_sse_no_accept() {
193        let headers = HeaderMap::new();
194        assert!(!is_mcp_sse(&headers));
195    }
196
197    // ── is_widget_asset ──
198
199    #[test]
200    fn widget_asset_by_js_ext() {
201        let headers = HeaderMap::new();
202        assert!(is_widget_asset("/assets/main.js", &headers));
203    }
204
205    #[test]
206    fn widget_asset_by_css_ext() {
207        let headers = HeaderMap::new();
208        assert!(is_widget_asset("/styles/app.css", &headers));
209    }
210
211    #[test]
212    fn widget_asset_by_woff2_ext() {
213        let headers = HeaderMap::new();
214        assert!(is_widget_asset("/fonts/inter.woff2", &headers));
215    }
216
217    #[test]
218    fn widget_asset_by_svg_ext() {
219        let headers = HeaderMap::new();
220        assert!(is_widget_asset("/icons/logo.svg", &headers));
221    }
222
223    #[test]
224    fn widget_asset_by_accept_html() {
225        let mut headers = HeaderMap::new();
226        headers.insert(header::ACCEPT, "text/html".parse().unwrap());
227        assert!(is_widget_asset("/some-path", &headers));
228    }
229
230    #[test]
231    fn widget_asset_by_accept_image() {
232        let mut headers = HeaderMap::new();
233        headers.insert(header::ACCEPT, "image/png".parse().unwrap());
234        assert!(is_widget_asset("/logo", &headers));
235    }
236
237    #[test]
238    fn widget_asset_by_accept_javascript() {
239        let mut headers = HeaderMap::new();
240        headers.insert(header::ACCEPT, "application/javascript".parse().unwrap());
241        assert!(is_widget_asset("/bundle", &headers));
242    }
243
244    #[test]
245    fn not_widget_asset_well_known() {
246        let headers = HeaderMap::new();
247        assert!(!is_widget_asset(
248            "/.well-known/oauth-authorization-server",
249            &headers
250        ));
251    }
252
253    #[test]
254    fn not_widget_asset_mcp() {
255        let headers = HeaderMap::new();
256        assert!(!is_widget_asset("/mcp", &headers));
257    }
258
259    #[test]
260    fn not_widget_asset_token() {
261        let headers = HeaderMap::new();
262        assert!(!is_widget_asset("/token", &headers));
263    }
264
265    #[test]
266    fn not_widget_asset_authorize() {
267        let headers = HeaderMap::new();
268        assert!(!is_widget_asset("/authorize", &headers));
269    }
270
271    #[test]
272    fn not_widget_asset_register() {
273        let headers = HeaderMap::new();
274        assert!(!is_widget_asset("/register", &headers));
275    }
276
277    #[test]
278    fn not_widget_asset_json_accept() {
279        let mut headers = HeaderMap::new();
280        headers.insert(header::ACCEPT, "application/json".parse().unwrap());
281        assert!(!is_widget_asset("/some-path", &headers));
282    }
283
284    #[test]
285    fn not_widget_asset_sse_accept() {
286        let mut headers = HeaderMap::new();
287        headers.insert(header::ACCEPT, "text/event-stream".parse().unwrap());
288        assert!(!is_widget_asset("/mcp", &headers));
289    }
290}