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