Skip to main content

winx_code_agent/
http_server.rs

1//! Optional remote MCP transport over Streamable HTTP.
2//!
3//! The endpoint exposes shell and filesystem capabilities, so remote mode is
4//! deliberately fail-closed: loopback bind by default, strong file-backed or
5//! explicit bearer credentials, bounded request concurrency/rate, and isolated
6//! daemon namespaces for independently authenticated principals.
7
8#![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/// Opaque per-request correlation shared by the HTTP middleware and MCP handler.
59/// It contains no client input and is safe to include in metadata-only usage logs.
60#[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    // The timeout sits INSIDE the security middleware: when it fires, the
150    // middleware still receives the 408 response and records the usage event.
151    // With the previous order the timed-out future was dropped before logging,
152    // so 408s were invisible in the usage log.
153    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    // Keep OAuth discovery probes and every unrelated path outside bearer auth.
159    // Winx does not advertise OAuth metadata, so those routes should be ordinary
160    // 404s instead of 401s that encourage clients to start a nonexistent flow.
161    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    // One structured `winx::usage` event per authenticated request. Metadata
306    // only — the body (commands, file contents) is never logged.
307    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}