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

winx_code_agent/
http_server.rs
