1use axum::{
2 body::Bytes,
3 http::{HeaderMap, Method, header},
4};
5use mcpr_protocol as jsonrpc;
6
7pub enum ClassifiedRequest {
9 OAuthCallback,
11 WidgetHtml { name: String },
13 WidgetList,
15 WidgetAsset,
17 McpPost { parsed: jsonrpc::ParsedBody },
19 McpSse,
21 Passthrough,
23}
24
25pub 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
68fn parse_mcp_body(body: &Bytes) -> Option<jsonrpc::ParsedBody> {
70 jsonrpc::parse_body(body)
71}
72
73fn 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
82fn 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 #[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 #[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 #[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}