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, VecDeque};
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, session::SessionManager as _, StreamableHttpServerConfig,
27    StreamableHttpService,
28};
29use serde::Serialize;
30use tokio::sync::{Mutex, Semaphore};
31use tower_http::limit::RequestBodyLimitLayer;
32use tower_http::timeout::TimeoutLayer;
33
34use crate::config::{HttpPrincipal, HttpSessionAffinity};
35use crate::runtime::{configured_shell_runtime, ShellRuntime};
36use crate::server::{SessionIsolation, WinxService};
37
38mod observability;
39use observability::{observe_request_entry, PermitBody};
40
41pub type BoxError = Box<dyn std::error::Error + Send + Sync>;
42
43const MAX_BODY_BYTES: usize = 64 * 1024 * 1024;
44const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
45const RATE_WINDOW: Duration = Duration::from_secs(60);
46const MAX_RATE_TRACKED_IPS: usize = 4_096;
47const INVALID_AUTH_DELAY: Duration = Duration::from_millis(100);
48pub const DEFAULT_MAX_CONCURRENCY: usize = 32;
49pub const DEFAULT_REQUESTS_PER_MINUTE: u32 = 120;
50/// Fixed-window request budget per authenticated principal. Behind one reverse
51/// proxy every client shares a single source IP, so the IP window alone cannot
52/// keep one busy principal from starving the others.
53pub const DEFAULT_PRINCIPAL_REQUESTS_PER_MINUTE: u32 = 120;
54
55#[derive(Debug)]
56struct HealthState {
57    ready: AtomicBool,
58}
59
60#[derive(Serialize)]
61struct HealthResponse {
62    status: &'static str,
63    ready: bool,
64}
65
66/// Opaque per-request correlation shared by the HTTP middleware and MCP handler.
67/// It contains no client input and is safe to include in metadata-only usage logs.
68#[derive(Clone, Debug)]
69pub(crate) struct RequestCorrelation(String);
70
71impl RequestCorrelation {
72    fn new() -> Self {
73        Self(format!("r_{:032x}", rand::rng().random::<u128>()))
74    }
75
76    pub(crate) fn as_str(&self) -> &str {
77        &self.0
78    }
79}
80
81#[derive(Clone, Debug)]
82pub struct HttpServerOptions {
83    pub bind: String,
84    pub principals: Vec<HttpPrincipal>,
85    pub extra_hosts: Vec<String>,
86    pub allow_query_token: bool,
87    pub allow_non_loopback: bool,
88    pub session_affinity: HttpSessionAffinity,
89    pub max_concurrency: usize,
90    pub requests_per_minute: u32,
91    pub principal_requests_per_minute: u32,
92}
93
94impl HttpServerOptions {
95    pub fn single(
96        bind: impl Into<String>,
97        token: impl Into<String>,
98        extra_hosts: Vec<String>,
99        allow_query_token: bool,
100    ) -> Result<Self, BoxError> {
101        let token = token.into();
102        let principal = HttpPrincipal::new("default", token.trim(), false)?;
103        Ok(Self {
104            bind: bind.into(),
105            principals: vec![principal],
106            extra_hosts,
107            allow_query_token,
108            allow_non_loopback: false,
109            session_affinity: HttpSessionAffinity::Workspace,
110            max_concurrency: DEFAULT_MAX_CONCURRENCY,
111            requests_per_minute: DEFAULT_REQUESTS_PER_MINUTE,
112            principal_requests_per_minute: DEFAULT_PRINCIPAL_REQUESTS_PER_MINUTE,
113        })
114    }
115}
116
117pub async fn start_http_server(
118    bind: &str,
119    token: String,
120    extra_hosts: Vec<String>,
121    allow_query_token: bool,
122) -> Result<(), BoxError> {
123    let options = HttpServerOptions::single(bind, token, extra_hosts, allow_query_token)?;
124    start_http_server_with_options(options).await
125}
126
127pub async fn start_http_server_with_options(options: HttpServerOptions) -> Result<(), BoxError> {
128    let shell_runtime = configured_shell_runtime().await?;
129    start_http_server_with_runtime(options, shell_runtime).await
130}
131
132pub async fn start_http_server_with_runtime(
133    options: HttpServerOptions,
134    shell_runtime: Arc<dyn ShellRuntime>,
135) -> Result<(), BoxError> {
136    let bind = validate_options(&options)?;
137    let principal_names =
138        options.principals.iter().map(|principal| principal.name().to_string()).collect::<Vec<_>>();
139    let shared_concurrency = Arc::new(Semaphore::new(options.max_concurrency));
140    let shared_rate_limiter = Arc::new(RateLimiter::new(options.requests_per_minute));
141    let principal_rate_limiter = Arc::new(RateLimiter::new(options.principal_requests_per_minute));
142
143    let mut config = StreamableHttpServerConfig::default();
144    config.legacy_session_mode = true;
145    config.allowed_hosts.extend(options.extra_hosts);
146
147    let shared = WinxService::with_runtime(SessionIsolation::Strict, shell_runtime);
148    let session_manager = Arc::new(LocalSessionManager::default());
149    let mcp_service =
150        StreamableHttpService::new(move || Ok(shared.clone()), session_manager.clone(), config);
151    let security = Arc::new(RouteSecurity {
152        principals: Arc::new(options.principals),
153        allow_query: options.allow_query_token,
154        session_affinity: options.session_affinity,
155        concurrency: shared_concurrency,
156        rate_limiter: shared_rate_limiter,
157        principal_rate_limiter,
158        session_owners: Arc::new(SessionOwners::default()),
159        session_manager,
160    });
161    // The timeout sits INSIDE the security middleware: when it fires, the
162    // middleware still receives the 408 response and records the usage event.
163    // With the previous order the timed-out future was dropped before logging,
164    // so 408s were invisible in the usage log.
165    let protected_mcp = Router::new()
166        .nest_service("/mcp", mcp_service)
167        .layer(TimeoutLayer::with_status_code(StatusCode::REQUEST_TIMEOUT, REQUEST_TIMEOUT))
168        .layer(middleware::from_fn_with_state(security, enforce_route_security))
169        .layer(RequestBodyLimitLayer::new(MAX_BODY_BYTES));
170    // Keep OAuth discovery probes and every unrelated path outside bearer auth.
171    // Winx does not advertise OAuth metadata, so those routes should be ordinary
172    // 404s instead of 401s that encourage clients to start a nonexistent flow.
173    let health_state = Arc::new(HealthState { ready: AtomicBool::new(true) });
174    let health_routes = Router::new()
175        .route("/healthz", get(healthz))
176        .route("/readyz", get(readyz))
177        .with_state(Arc::clone(&health_state));
178    let app = Router::new().merge(health_routes).merge(protected_mcp);
179
180    let listener = tokio::net::TcpListener::bind(bind).await?;
181    tracing::warn!(
182        bind = %bind,
183        principals = ?principal_names,
184        route = "/mcp",
185        "Winx remote MCP is network-reachable; keep it behind authenticated HTTPS"
186    );
187    if options.allow_query_token {
188        tracing::warn!(
189            "query-token authentication is enabled; URLs may expose credentials in proxy logs/history"
190        );
191    }
192    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
193        .with_graceful_shutdown(shutdown_signal(health_state))
194        .await?;
195    Ok(())
196}
197
198async fn healthz() -> impl IntoResponse {
199    (StatusCode::OK, Json(HealthResponse { status: "ok", ready: true }))
200}
201
202async fn readyz(State(state): State<Arc<HealthState>>) -> impl IntoResponse {
203    let ready = state.ready.load(Ordering::Acquire);
204    let status = if ready { StatusCode::OK } else { StatusCode::SERVICE_UNAVAILABLE };
205    (status, Json(HealthResponse { status: if ready { "ready" } else { "draining" }, ready }))
206}
207
208async fn shutdown_signal(state: Arc<HealthState>) {
209    #[cfg(unix)]
210    {
211        let terminate = async {
212            match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
213                Ok(mut signal) => {
214                    signal.recv().await;
215                }
216                Err(error) => {
217                    tracing::warn!(%error, "failed to install SIGTERM handler");
218                    std::future::pending::<()>().await;
219                }
220            }
221        };
222        tokio::select! {
223            result = tokio::signal::ctrl_c() => {
224                if let Err(error) = result {
225                    tracing::warn!(%error, "failed while waiting for Ctrl-C");
226                }
227            }
228            () = terminate => {}
229        }
230    }
231    #[cfg(not(unix))]
232    if let Err(error) = tokio::signal::ctrl_c().await {
233        tracing::warn!(%error, "failed while waiting for Ctrl-C");
234    }
235
236    state.ready.store(false, Ordering::Release);
237    tracing::info!("HTTP shutdown requested; draining in-flight requests");
238}
239
240fn validate_options(options: &HttpServerOptions) -> Result<SocketAddr, BoxError> {
241    let bind: SocketAddr = options
242        .bind
243        .parse()
244        .map_err(|error| format!("invalid --bind socket address {:?}: {error}", options.bind))?;
245    if !bind.ip().is_loopback() && !options.allow_non_loopback {
246        return Err(format!(
247            "refusing non-loopback HTTP bind {bind}; pass --allow-non-loopback only behind a trusted firewall/tunnel"
248        )
249        .into());
250    }
251    if options.max_concurrency == 0 {
252        return Err("max_concurrency must be at least 1".into());
253    }
254    if options.requests_per_minute == 0 {
255        return Err("requests_per_minute must be at least 1".into());
256    }
257    if options.principal_requests_per_minute == 0 {
258        return Err("principal_requests_per_minute must be at least 1".into());
259    }
260    if options.principals.is_empty() {
261        return Err("at least one HTTP principal is required".into());
262    }
263    let mut names = HashSet::new();
264    for principal in &options.principals {
265        if !names.insert(principal.name()) {
266            return Err(format!("duplicate HTTP principal: {}", principal.name()).into());
267        }
268    }
269    Ok(bind)
270}
271
272#[derive(Clone)]
273struct RouteSecurity {
274    principals: Arc<Vec<HttpPrincipal>>,
275    allow_query: bool,
276    session_affinity: HttpSessionAffinity,
277    concurrency: Arc<Semaphore>,
278    rate_limiter: Arc<RateLimiter<IpAddr>>,
279    principal_rate_limiter: Arc<RateLimiter<String>>,
280    session_owners: Arc<SessionOwners>,
281    session_manager: Arc<LocalSessionManager>,
282}
283
284/// Upper bound on remembered transport sessions. Past it the oldest entry is
285/// evicted and its SDK session closed at the same time, so an id never
286/// outlives the record that says who may use it.
287const MAX_TRACKED_SESSIONS: usize = 65_536;
288
289/// Which principal created each legacy MCP transport session (`Mcp-Session-Id`).
290///
291/// The SDK's session manager is shared by every principal and treats a known
292/// session id as sufficient to resume its stream or delete it. Winx records
293/// the creator when the `initialize` response mints the id and refuses any
294/// later request that presents the id under another principal's token.
295#[derive(Default)]
296struct SessionOwners {
297    inner: std::sync::Mutex<SessionOwnerTable>,
298}
299
300#[derive(Default)]
301struct SessionOwnerTable {
302    owners: HashMap<String, String>,
303    order: VecDeque<String>,
304}
305
306impl SessionOwners {
307    /// Only the recorded creator may present a minted `session_id`. An id this
308    /// process never minted is not an SDK session at all: modern stateless
309    /// clients may send one as a conversation identity, and the SDK answers
310    /// 404 itself if such an id reaches a legacy session path.
311    fn owner_matches(&self, session_id: &str, principal: &str) -> bool {
312        let table = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
313        table.owners.get(session_id).is_none_or(|owner| owner == principal)
314    }
315
316    /// Remember the creator of a freshly minted id. Returns the ids evicted
317    /// to stay within [`MAX_TRACKED_SESSIONS`]; the caller closes those SDK
318    /// sessions so a forgotten id can never be presented by anyone.
319    fn record(&self, session_id: &str, principal: &str) -> Vec<String> {
320        let mut table = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
321        if table.owners.contains_key(session_id) {
322            return Vec::new();
323        }
324        table.owners.insert(session_id.to_string(), principal.to_string());
325        table.order.push_back(session_id.to_string());
326        let mut evicted = Vec::new();
327        while table.order.len() > MAX_TRACKED_SESSIONS {
328            if let Some(oldest) = table.order.pop_front() {
329                table.owners.remove(&oldest);
330                evicted.push(oldest);
331            }
332        }
333        evicted
334    }
335
336    fn forget(&self, session_id: &str) {
337        let mut table = self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
338        if table.owners.remove(session_id).is_some() {
339            table.order.retain(|id| id != session_id);
340        }
341    }
342}
343
344fn session_id_header(headers: &header::HeaderMap) -> Option<String> {
345    headers
346        .get("mcp-session-id")
347        .and_then(|value| value.to_str().ok())
348        .map(str::trim)
349        .filter(|value| !value.is_empty())
350        .map(ToString::to_string)
351}
352
353async fn enforce_route_security(
354    State(security): State<Arc<RouteSecurity>>,
355    ConnectInfo(peer): ConnectInfo<SocketAddr>,
356    request: Request,
357    next: Next,
358) -> Response {
359    let path = request.uri().path();
360    if path != "/mcp" && !path.starts_with("/mcp/") {
361        return next.run(request).await;
362    }
363    let (started, method, correlation) = observe_request_entry(peer, request.method());
364    if !security.rate_limiter.allow(peer.ip()).await {
365        log_rejection(peer, None, &correlation, "ip_rate_limit", StatusCode::TOO_MANY_REQUESTS);
366        return retry_response(StatusCode::TOO_MANY_REQUESTS, "request rate limit exceeded\n");
367    }
368    let Ok(permit) = security.concurrency.clone().try_acquire_owned() else {
369        log_rejection(peer, None, &correlation, "concurrency", StatusCode::SERVICE_UNAVAILABLE);
370        return retry_response(StatusCode::SERVICE_UNAVAILABLE, "too many concurrent requests\n");
371    };
372
373    let Some(principal) =
374        authenticate_request(&request, &security.principals, security.allow_query)
375    else {
376        tokio::time::sleep(INVALID_AUTH_DELAY).await;
377        tracing::warn!(
378            peer = %peer,
379            path = %request.uri().path(),
380            "rejected request with missing or invalid HTTP token"
381        );
382        log_rejection(peer, None, &correlation, "auth", StatusCode::UNAUTHORIZED);
383        return (StatusCode::UNAUTHORIZED, "missing or invalid token\n").into_response();
384    };
385
386    let principal_name = principal.name().to_string();
387    if !security.principal_rate_limiter.allow(principal_name.clone()).await {
388        log_rejection(
389            peer,
390            Some(&principal_name),
391            &correlation,
392            "principal_rate_limit",
393            StatusCode::TOO_MANY_REQUESTS,
394        );
395        return retry_response(StatusCode::TOO_MANY_REQUESTS, "principal rate limit exceeded\n");
396    }
397    // A transport session id belongs to the principal that created it. Any
398    // other authenticated principal sees it as unknown, exactly like a stale
399    // id: the SDK would otherwise let it resume the stream or close it.
400    let presented_session = session_id_header(request.headers());
401    if let Some(session_id) = presented_session.as_deref() {
402        if !security.session_owners.owner_matches(session_id, &principal_name) {
403            tracing::warn!(
404                peer = %peer,
405                principal = %principal_name,
406                "rejected request presenting another principal's MCP session id"
407            );
408            log_rejection(
409                peer,
410                Some(&principal_name),
411                &correlation,
412                "foreign_session",
413                StatusCode::NOT_FOUND,
414            );
415            return (StatusCode::NOT_FOUND, "session not found\n").into_response();
416        }
417    }
418    let mut request = request;
419    request.extensions_mut().insert(principal);
420    request.extensions_mut().insert(security.session_affinity);
421    request.extensions_mut().insert(correlation.clone());
422    let response = next.run(request).await;
423    if response.status().is_success() {
424        if let Some(minted) = session_id_header(response.headers()) {
425            for evicted in security.session_owners.record(&minted, &principal_name) {
426                let manager = security.session_manager.clone();
427                tokio::spawn(async move {
428                    let id: rmcp::transport::common::server_side_http::SessionId =
429                        Arc::from(evicted.as_str());
430                    let _ = manager.close_session(&id).await;
431                });
432            }
433        }
434        if method == axum::http::Method::DELETE {
435            if let Some(session_id) = presented_session.as_deref() {
436                security.session_owners.forget(session_id);
437            }
438        }
439    }
440    // Preserve the legacy event and duration as the response-ready observation.
441    // Body termination is a separate metadata-only event from PermitBody.
442    let response_ready = Instant::now();
443    tracing::info!(
444        target: "winx::usage",
445        event = "http_request",
446        phase = "response_ready",
447        principal = %principal_name,
448        request_id = %correlation.as_str(),
449        peer = %peer,
450        method = %method,
451        status = response.status().as_u16(),
452        duration_ms = u64::try_from(response_ready.saturating_duration_since(started).as_millis())
453            .unwrap_or(u64::MAX),
454        "http request"
455    );
456    // The permit travels with the body, including while a streaming response is
457    // open. An already-empty body can now log completion after response-ready.
458    response.map(|body| {
459        axum::body::Body::new(PermitBody::new(body, permit).observe(
460            correlation,
461            started,
462            response_ready,
463        ))
464    })
465}
466
467/// One redacted `winx::usage` event per admission rejection, so overload and
468/// auth failures are visible in reports instead of vanishing before the
469/// per-request event.
470fn log_rejection(
471    peer: SocketAddr,
472    principal: Option<&str>,
473    correlation: &RequestCorrelation,
474    reason: &str,
475    status: StatusCode,
476) {
477    tracing::info!(
478        target: "winx::usage",
479        event = "http_rejected",
480        request_id = %correlation.as_str(),
481        principal = principal.unwrap_or(""),
482        peer = %peer,
483        reason,
484        status = status.as_u16(),
485        "http request rejected"
486    );
487}
488fn retry_response(status: StatusCode, body: &'static str) -> Response {
489    let mut response = (status, body).into_response();
490    response.headers_mut().insert(header::RETRY_AFTER, HeaderValue::from_static("1"));
491    response
492}
493
494#[derive(Debug)]
495struct RateWindow {
496    started_at: Instant,
497    last_seen: Instant,
498    count: u32,
499}
500
501#[derive(Debug)]
502/// Fixed-window limiter keyed by source IP (pre-auth abuse control) or by
503/// authenticated principal (fairness behind a shared proxy).
504struct RateLimiter<K> {
505    requests_per_window: u32,
506    windows: Mutex<HashMap<K, RateWindow>>,
507}
508
509impl<K: std::hash::Hash + Eq + Clone> RateLimiter<K> {
510    fn new(requests_per_window: u32) -> Self {
511        Self { requests_per_window, windows: Mutex::new(HashMap::new()) }
512    }
513
514    async fn allow(&self, peer: K) -> bool {
515        let now = Instant::now();
516        let mut windows = self.windows.lock().await;
517        windows.retain(|_, window| now.duration_since(window.last_seen) < RATE_WINDOW * 2);
518        if !windows.contains_key(&peer) && windows.len() >= MAX_RATE_TRACKED_IPS {
519            if let Some(oldest) =
520                windows.iter().min_by_key(|(_, window)| window.last_seen).map(|(ip, _)| ip.clone())
521            {
522                windows.remove(&oldest);
523            }
524        }
525        let window =
526            windows.entry(peer).or_insert(RateWindow { started_at: now, last_seen: now, count: 0 });
527        if now.duration_since(window.started_at) >= RATE_WINDOW {
528            window.started_at = now;
529            window.count = 0;
530        }
531        window.last_seen = now;
532        if window.count >= self.requests_per_window {
533            return false;
534        }
535        window.count = window.count.saturating_add(1);
536        true
537    }
538}
539
540fn authenticate_request(
541    request: &Request,
542    principals: &[HttpPrincipal],
543    allow_query: bool,
544) -> Option<HttpPrincipal> {
545    let header = bearer_token(request);
546    let query = allow_query.then(|| query_token(request)).flatten();
547    [header, query].into_iter().flatten().find_map(|presented| {
548        principals.iter().find(|principal| principal.matches_token(presented)).cloned()
549    })
550}
551
552fn bearer_token(request: &Request) -> Option<&str> {
553    let value = request.headers().get(header::AUTHORIZATION)?.to_str().ok()?;
554    let (scheme, token) = value.split_once(' ')?;
555    scheme.eq_ignore_ascii_case("bearer").then_some(token.trim())
556}
557
558fn query_token(request: &Request) -> Option<&str> {
559    request.uri().query()?.split('&').find_map(|pair| pair.strip_prefix("token="))
560}
561
562#[cfg(test)]
563mod admission_tests {
564    #![allow(clippy::expect_used)]
565    use std::sync::Arc;
566
567    use tokio::sync::Semaphore;
568
569    use super::{PermitBody, RateLimiter};
570
571    #[tokio::test]
572    async fn principal_windows_are_independent_of_each_other() {
573        let limiter: RateLimiter<String> = RateLimiter::new(2);
574        assert!(limiter.allow("left".to_string()).await);
575        assert!(limiter.allow("left".to_string()).await);
576        assert!(!limiter.allow("left".to_string()).await, "left exhausted its window");
577        assert!(limiter.allow("right".to_string()).await, "right is unaffected");
578    }
579
580    #[tokio::test]
581    async fn permit_is_released_only_when_the_body_finishes() {
582        let semaphore = Arc::new(Semaphore::new(1));
583        let permit = semaphore.clone().try_acquire_owned().expect("permit");
584        let body =
585            axum::body::Body::new(PermitBody::new(axum::body::Body::from("payload"), permit));
586        assert_eq!(semaphore.available_permits(), 0, "held while the body is pending");
587        let bytes = axum::body::to_bytes(body, usize::MAX).await.expect("read body");
588        assert_eq!(&bytes[..], b"payload");
589        assert_eq!(semaphore.available_permits(), 1, "released with the consumed body");
590
591        let permit = semaphore.clone().try_acquire_owned().expect("permit again");
592        let body =
593            axum::body::Body::new(PermitBody::new(axum::body::Body::from("abandoned"), permit));
594        drop(body);
595        assert_eq!(semaphore.available_permits(), 1, "released when a client disconnects");
596    }
597}
598
599#[cfg(test)]
600mod session_owner_tests {
601    use super::SessionOwners;
602
603    #[test]
604    fn minted_session_ids_belong_to_their_first_creator() {
605        let owners = SessionOwners::default();
606        assert!(owners.owner_matches("conv-a", "left"), "unminted ids are not SDK sessions");
607        assert!(owners.record("s1", "left").is_empty(), "nothing evicted");
608        assert!(owners.owner_matches("s1", "left"));
609        assert!(!owners.owner_matches("s1", "right"), "another principal");
610        assert!(owners.record("s1", "right").is_empty(), "nothing evicted");
611        assert!(owners.owner_matches("s1", "left"), "first creator wins");
612        owners.forget("s1");
613        assert!(owners.owner_matches("s1", "right"), "closed ids are plain unknown ids again");
614    }
615
616    #[test]
617    fn eviction_hands_back_the_ids_whose_sdk_sessions_must_close() {
618        let owners = SessionOwners::default();
619        assert!(owners.record("oldest", "left").is_empty(), "nothing evicted");
620        let mut evicted = Vec::new();
621        for index in 0..super::MAX_TRACKED_SESSIONS {
622            evicted.extend(owners.record(&format!("s{index}"), "left"));
623        }
624        assert_eq!(evicted, vec!["oldest".to_string()]);
625        assert!(owners.owner_matches(&format!("s{}", super::MAX_TRACKED_SESSIONS - 1), "left"));
626    }
627}
628
629#[cfg(test)]
630mod tests {
631    #![allow(clippy::expect_used)]
632
633    use std::net::{IpAddr, Ipv4Addr};
634    use std::sync::atomic::{AtomicBool, Ordering};
635    use std::sync::Arc;
636
637    use axum::body::Body;
638    use axum::extract::{Request, State};
639    use axum::http::StatusCode;
640    use axum::response::IntoResponse as _;
641
642    use super::{
643        authenticate_request, readyz, validate_options, HealthState, HttpServerOptions, RateLimiter,
644    };
645    use crate::config::HttpPrincipal;
646
647    fn request(uri: &str, authorization: Option<&str>) -> Request {
648        let mut builder = Request::builder().uri(uri);
649        if let Some(authorization) = authorization {
650            builder = builder.header(axum::http::header::AUTHORIZATION, authorization);
651        }
652        builder.body(Body::empty()).expect("test request")
653    }
654
655    fn principal(name: &str, token_byte: char) -> HttpPrincipal {
656        let token = token_byte.to_string().repeat(32);
657        HttpPrincipal::new(name, &token, false).expect("valid principal")
658    }
659
660    fn strong_options(bind: &str) -> HttpServerOptions {
661        HttpServerOptions {
662            bind: bind.to_string(),
663            principals: vec![principal("default", 'x')],
664            extra_hosts: Vec::new(),
665            allow_query_token: false,
666            allow_non_loopback: false,
667            session_affinity: crate::config::HttpSessionAffinity::Workspace,
668            max_concurrency: super::DEFAULT_MAX_CONCURRENCY,
669            requests_per_minute: super::DEFAULT_REQUESTS_PER_MINUTE,
670            principal_requests_per_minute: super::DEFAULT_PRINCIPAL_REQUESTS_PER_MINUTE,
671        }
672    }
673
674    #[test]
675    fn bearer_and_opt_in_query_tokens_select_the_matching_principal() {
676        let principals = vec![principal("chatgpt", 'a'), principal("automation", 'b')];
677        let header = request("/mcp", Some(&format!("Bearer {}", "a".repeat(32))));
678        let matched = authenticate_request(&header, &principals, false).expect("header principal");
679        assert_eq!(matched.name(), "chatgpt");
680
681        let query_uri = format!("/mcp?token={}", "b".repeat(32));
682        assert!(authenticate_request(&request(&query_uri, None), &principals, false).is_none());
683        let matched = authenticate_request(&request(&query_uri, None), &principals, true)
684            .expect("query principal");
685        assert_eq!(matched.name(), "automation");
686    }
687
688    #[test]
689    fn non_loopback_binding_requires_explicit_override() {
690        assert!(validate_options(&strong_options("127.0.0.1:8000")).is_ok());
691        assert!(validate_options(&strong_options("0.0.0.0:8000")).is_err());
692
693        let mut public = strong_options("0.0.0.0:8000");
694        public.allow_non_loopback = true;
695        assert!(validate_options(&public).is_ok());
696    }
697
698    #[tokio::test]
699    async fn rate_limiter_rejects_excess_requests() {
700        let limiter = RateLimiter::new(2);
701        let peer = IpAddr::V4(Ipv4Addr::LOCALHOST);
702        assert!(limiter.allow(peer).await);
703        assert!(limiter.allow(peer).await);
704        assert!(!limiter.allow(peer).await);
705    }
706
707    #[test]
708    fn duplicate_principal_names_are_rejected() {
709        let mut options = strong_options("127.0.0.1:8000");
710        options.principals = vec![principal("same", 'a'), principal("same", 'b')];
711        assert!(validate_options(&options).is_err());
712    }
713
714    #[tokio::test]
715    async fn readiness_turns_unavailable_while_draining() {
716        let state = Arc::new(HealthState { ready: AtomicBool::new(true) });
717        let ready = readyz(State(Arc::clone(&state))).await.into_response();
718        assert_eq!(ready.status(), StatusCode::OK);
719
720        state.ready.store(false, Ordering::Release);
721        let draining = readyz(State(state)).await.into_response();
722        assert_eq!(draining.status(), StatusCode::SERVICE_UNAVAILABLE);
723    }
724}
725
726#[cfg(test)]
727mod observability_tests;
winx_code_agent/http_server.rs

winx_code_agent/
http_server.rs
