1#![allow(clippy::doc_markdown)]
9
10use std::collections::{HashMap, HashSet};
11use std::net::{IpAddr, SocketAddr};
12use std::sync::Arc;
13use std::time::{Duration, Instant};
14
15use axum::{
16 extract::{ConnectInfo, Request, State},
17 http::{header, HeaderValue, StatusCode},
18 middleware::{self, Next},
19 response::{IntoResponse, Response},
20 Router,
21};
22use rmcp::transport::streamable_http_server::{
23 session::local::LocalSessionManager, StreamableHttpServerConfig, StreamableHttpService,
24};
25use tokio::sync::{Mutex, Semaphore};
26use tower_http::limit::RequestBodyLimitLayer;
27use tower_http::timeout::TimeoutLayer;
28
29use crate::config::{HttpPrincipal, HttpSessionAffinity};
30use crate::runtime::{configured_shell_runtime, ShellRuntime};
31use crate::server::{SessionIsolation, WinxService};
32
33pub type BoxError = Box<dyn std::error::Error + Send + Sync>;
34
35const MAX_BODY_BYTES: usize = 64 * 1024 * 1024;
36const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
37const RATE_WINDOW: Duration = Duration::from_secs(60);
38const MAX_RATE_TRACKED_IPS: usize = 4_096;
39const INVALID_AUTH_DELAY: Duration = Duration::from_millis(100);
40pub const DEFAULT_MAX_CONCURRENCY: usize = 32;
41pub const DEFAULT_REQUESTS_PER_MINUTE: u32 = 120;
42
43#[derive(Clone, Debug)]
44pub struct HttpServerOptions {
45 pub bind: String,
46 pub principals: Vec<HttpPrincipal>,
47 pub extra_hosts: Vec<String>,
48 pub allow_query_token: bool,
49 pub allow_non_loopback: bool,
50 pub session_affinity: HttpSessionAffinity,
51 pub max_concurrency: usize,
52 pub requests_per_minute: u32,
53}
54
55impl HttpServerOptions {
56 pub fn single(
57 bind: impl Into<String>,
58 token: impl Into<String>,
59 extra_hosts: Vec<String>,
60 allow_query_token: bool,
61 ) -> Result<Self, BoxError> {
62 let token = token.into();
63 let principal = HttpPrincipal::new("default", token.trim(), false)?;
64 Ok(Self {
65 bind: bind.into(),
66 principals: vec![principal],
67 extra_hosts,
68 allow_query_token,
69 allow_non_loopback: false,
70 session_affinity: HttpSessionAffinity::Workspace,
71 max_concurrency: DEFAULT_MAX_CONCURRENCY,
72 requests_per_minute: DEFAULT_REQUESTS_PER_MINUTE,
73 })
74 }
75}
76
77pub async fn start_http_server(
78 bind: &str,
79 token: String,
80 extra_hosts: Vec<String>,
81 allow_query_token: bool,
82) -> Result<(), BoxError> {
83 let options = HttpServerOptions::single(bind, token, extra_hosts, allow_query_token)?;
84 start_http_server_with_options(options).await
85}
86
87pub async fn start_http_server_with_options(options: HttpServerOptions) -> Result<(), BoxError> {
88 let shell_runtime = configured_shell_runtime().await?;
89 start_http_server_with_runtime(options, shell_runtime).await
90}
91
92pub async fn start_http_server_with_runtime(
93 options: HttpServerOptions,
94 shell_runtime: Arc<dyn ShellRuntime>,
95) -> Result<(), BoxError> {
96 let bind = validate_options(&options)?;
97 let principal_names =
98 options.principals.iter().map(|principal| principal.name().to_string()).collect::<Vec<_>>();
99 let shared_concurrency = Arc::new(Semaphore::new(options.max_concurrency));
100 let shared_rate_limiter = Arc::new(RateLimiter::new(options.requests_per_minute));
101
102 let mut config = StreamableHttpServerConfig::default();
103 config.legacy_session_mode = true;
104 config.allowed_hosts.extend(options.extra_hosts);
105
106 let shared = WinxService::with_runtime(SessionIsolation::Strict, shell_runtime);
107 let mcp_service = StreamableHttpService::new(
108 move || Ok(shared.clone()),
109 Arc::new(LocalSessionManager::default()),
110 config,
111 );
112 let security = Arc::new(RouteSecurity {
113 principals: Arc::new(options.principals),
114 allow_query: options.allow_query_token,
115 session_affinity: options.session_affinity,
116 concurrency: shared_concurrency,
117 rate_limiter: shared_rate_limiter,
118 });
119 let protected_mcp = Router::new()
124 .nest_service("/mcp", mcp_service)
125 .layer(TimeoutLayer::with_status_code(StatusCode::REQUEST_TIMEOUT, REQUEST_TIMEOUT))
126 .layer(middleware::from_fn_with_state(security, enforce_route_security))
127 .layer(RequestBodyLimitLayer::new(MAX_BODY_BYTES));
128 let app = Router::new().merge(protected_mcp);
132
133 let listener = tokio::net::TcpListener::bind(bind).await?;
134 tracing::warn!(
135 bind = %bind,
136 principals = ?principal_names,
137 route = "/mcp",
138 "Winx remote MCP is network-reachable; keep it behind authenticated HTTPS"
139 );
140 if options.allow_query_token {
141 tracing::warn!(
142 "query-token authentication is enabled; URLs may expose credentials in proxy logs/history"
143 );
144 }
145 axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await?;
146 Ok(())
147}
148
149fn validate_options(options: &HttpServerOptions) -> Result<SocketAddr, BoxError> {
150 let bind: SocketAddr = options
151 .bind
152 .parse()
153 .map_err(|error| format!("invalid --bind socket address {:?}: {error}", options.bind))?;
154 if !bind.ip().is_loopback() && !options.allow_non_loopback {
155 return Err(format!(
156 "refusing non-loopback HTTP bind {bind}; pass --allow-non-loopback only behind a trusted firewall/tunnel"
157 )
158 .into());
159 }
160 if options.max_concurrency == 0 {
161 return Err("max_concurrency must be at least 1".into());
162 }
163 if options.requests_per_minute == 0 {
164 return Err("requests_per_minute must be at least 1".into());
165 }
166 if options.principals.is_empty() {
167 return Err("at least one HTTP principal is required".into());
168 }
169 let mut names = HashSet::new();
170 for principal in &options.principals {
171 if !names.insert(principal.name()) {
172 return Err(format!("duplicate HTTP principal: {}", principal.name()).into());
173 }
174 }
175 Ok(bind)
176}
177
178#[derive(Clone)]
179struct RouteSecurity {
180 principals: Arc<Vec<HttpPrincipal>>,
181 allow_query: bool,
182 session_affinity: HttpSessionAffinity,
183 concurrency: Arc<Semaphore>,
184 rate_limiter: Arc<RateLimiter>,
185}
186
187async fn enforce_route_security(
188 State(security): State<Arc<RouteSecurity>>,
189 ConnectInfo(peer): ConnectInfo<SocketAddr>,
190 request: Request,
191 next: Next,
192) -> Response {
193 let path = request.uri().path();
194 if path != "/mcp" && !path.starts_with("/mcp/") {
195 return next.run(request).await;
196 }
197 let started = Instant::now();
198 if !security.rate_limiter.allow(peer.ip()).await {
199 return retry_response(StatusCode::TOO_MANY_REQUESTS, "request rate limit exceeded\n");
200 }
201 let Ok(permit) = security.concurrency.clone().try_acquire_owned() else {
202 return retry_response(StatusCode::SERVICE_UNAVAILABLE, "too many concurrent requests\n");
203 };
204
205 let Some(principal) =
206 authenticate_request(&request, &security.principals, security.allow_query)
207 else {
208 tokio::time::sleep(INVALID_AUTH_DELAY).await;
209 tracing::warn!(
210 peer = %peer,
211 path = %request.uri().path(),
212 "rejected request with missing or invalid HTTP token"
213 );
214 return (StatusCode::UNAUTHORIZED, "missing or invalid token\n").into_response();
215 };
216
217 let principal_name = principal.name().to_string();
218 let method = request.method().clone();
219 let mut request = request;
220 request.extensions_mut().insert(principal);
221 request.extensions_mut().insert(security.session_affinity);
222 let response = next.run(request).await;
223 drop(permit);
224 tracing::info!(
227 target: "winx::usage",
228 event = "http_request",
229 principal = %principal_name,
230 peer = %peer,
231 method = %method,
232 status = response.status().as_u16(),
233 duration_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX),
234 "http request"
235 );
236 response
237}
238
239fn retry_response(status: StatusCode, body: &'static str) -> Response {
240 let mut response = (status, body).into_response();
241 response.headers_mut().insert(header::RETRY_AFTER, HeaderValue::from_static("1"));
242 response
243}
244
245#[derive(Debug)]
246struct RateWindow {
247 started_at: Instant,
248 last_seen: Instant,
249 count: u32,
250}
251
252#[derive(Debug)]
253struct RateLimiter {
254 requests_per_window: u32,
255 windows: Mutex<HashMap<IpAddr, RateWindow>>,
256}
257
258impl RateLimiter {
259 fn new(requests_per_window: u32) -> Self {
260 Self { requests_per_window, windows: Mutex::new(HashMap::new()) }
261 }
262
263 async fn allow(&self, peer: IpAddr) -> bool {
264 let now = Instant::now();
265 let mut windows = self.windows.lock().await;
266 windows.retain(|_, window| now.duration_since(window.last_seen) < RATE_WINDOW * 2);
267 if !windows.contains_key(&peer) && windows.len() >= MAX_RATE_TRACKED_IPS {
268 if let Some(oldest) =
269 windows.iter().min_by_key(|(_, window)| window.last_seen).map(|(ip, _)| *ip)
270 {
271 windows.remove(&oldest);
272 }
273 }
274 let window =
275 windows.entry(peer).or_insert(RateWindow { started_at: now, last_seen: now, count: 0 });
276 if now.duration_since(window.started_at) >= RATE_WINDOW {
277 window.started_at = now;
278 window.count = 0;
279 }
280 window.last_seen = now;
281 if window.count >= self.requests_per_window {
282 return false;
283 }
284 window.count = window.count.saturating_add(1);
285 true
286 }
287}
288
289fn authenticate_request(
290 request: &Request,
291 principals: &[HttpPrincipal],
292 allow_query: bool,
293) -> Option<HttpPrincipal> {
294 let header = bearer_token(request);
295 let query = allow_query.then(|| query_token(request)).flatten();
296 [header, query].into_iter().flatten().find_map(|presented| {
297 principals.iter().find(|principal| principal.matches_token(presented)).cloned()
298 })
299}
300
301fn bearer_token(request: &Request) -> Option<&str> {
302 let value = request.headers().get(header::AUTHORIZATION)?.to_str().ok()?;
303 let (scheme, token) = value.split_once(' ')?;
304 scheme.eq_ignore_ascii_case("bearer").then_some(token.trim())
305}
306
307fn query_token(request: &Request) -> Option<&str> {
308 request.uri().query()?.split('&').find_map(|pair| pair.strip_prefix("token="))
309}
310
311#[cfg(test)]
312mod tests {
313 #![allow(clippy::expect_used)]
314
315 use std::net::{IpAddr, Ipv4Addr};
316
317 use axum::body::Body;
318 use axum::extract::Request;
319
320 use super::{authenticate_request, validate_options, HttpServerOptions, RateLimiter};
321 use crate::config::HttpPrincipal;
322
323 fn request(uri: &str, authorization: Option<&str>) -> Request {
324 let mut builder = Request::builder().uri(uri);
325 if let Some(authorization) = authorization {
326 builder = builder.header(axum::http::header::AUTHORIZATION, authorization);
327 }
328 builder.body(Body::empty()).expect("test request")
329 }
330
331 fn principal(name: &str, token_byte: char) -> HttpPrincipal {
332 let token = token_byte.to_string().repeat(32);
333 HttpPrincipal::new(name, &token, false).expect("valid principal")
334 }
335
336 fn strong_options(bind: &str) -> HttpServerOptions {
337 HttpServerOptions {
338 bind: bind.to_string(),
339 principals: vec![principal("default", 'x')],
340 extra_hosts: Vec::new(),
341 allow_query_token: false,
342 allow_non_loopback: false,
343 session_affinity: crate::config::HttpSessionAffinity::Workspace,
344 max_concurrency: super::DEFAULT_MAX_CONCURRENCY,
345 requests_per_minute: super::DEFAULT_REQUESTS_PER_MINUTE,
346 }
347 }
348
349 #[test]
350 fn bearer_and_opt_in_query_tokens_select_the_matching_principal() {
351 let principals = vec![principal("chatgpt", 'a'), principal("automation", 'b')];
352 let header = request("/mcp", Some(&format!("Bearer {}", "a".repeat(32))));
353 let matched = authenticate_request(&header, &principals, false).expect("header principal");
354 assert_eq!(matched.name(), "chatgpt");
355
356 let query_uri = format!("/mcp?token={}", "b".repeat(32));
357 assert!(authenticate_request(&request(&query_uri, None), &principals, false).is_none());
358 let matched = authenticate_request(&request(&query_uri, None), &principals, true)
359 .expect("query principal");
360 assert_eq!(matched.name(), "automation");
361 }
362
363 #[test]
364 fn non_loopback_binding_requires_explicit_override() {
365 assert!(validate_options(&strong_options("127.0.0.1:8000")).is_ok());
366 assert!(validate_options(&strong_options("0.0.0.0:8000")).is_err());
367
368 let mut public = strong_options("0.0.0.0:8000");
369 public.allow_non_loopback = true;
370 assert!(validate_options(&public).is_ok());
371 }
372
373 #[tokio::test]
374 async fn rate_limiter_rejects_excess_requests() {
375 let limiter = RateLimiter::new(2);
376 let peer = IpAddr::V4(Ipv4Addr::LOCALHOST);
377 assert!(limiter.allow(peer).await);
378 assert!(limiter.allow(peer).await);
379 assert!(!limiter.allow(peer).await);
380 }
381
382 #[test]
383 fn duplicate_principal_names_are_rejected() {
384 let mut options = strong_options("127.0.0.1:8000");
385 options.principals = vec![principal("same", 'a'), principal("same", 'b')];
386 assert!(validate_options(&options).is_err());
387 }
388}