From f6535e94cbd39036f45769d80ec46c9a476ce8cd Mon Sep 17 00:00:00 2001 From: ajianaz Date: Fri, 16 Oct 2026 14:32:08 +0700 Subject: [PATCH] =?UTF-8?q?fix(cli):=20bring=20the=20MCP=20transport=20to?= =?UTF-8?q?=20parity=20with=20REST=20=E2=80=94=20constant-time=20auth,=20f?= =?UTF-8?q?ailed-auth=20limiter,=20body=20ceiling,=20resilient=20accept?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 3 + crates/tole-cli/src/auth.rs | 156 ++++++++++++++++++ crates/tole-cli/src/lib.rs | 1 + crates/tole-cli/src/mcp_http.rs | 163 +++++++++++++++---- crates/tole-cli/src/serve.rs | 93 ++--------- crates/tole-cli/tests/mcp_http_hardening.rs | 168 ++++++++++++++++++++ 6 files changed, 470 insertions(+), 114 deletions(-) create mode 100644 crates/tole-cli/src/auth.rs create mode 100644 crates/tole-cli/tests/mcp_http_hardening.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a9ec9c0..bb13e3a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `tole doctor` (#378): tiered provider diagnostics — `offline` (wiring), `catalog` (live `/models`, model must be served) and `full` (chat, streaming and a real tool-call loop) — each blocker carries its fix command, and keys / URL credentials are never printed. Exits non-zero when the requested tier is not reached; `--tier full` is the release-validation gate. - Per-user-turn token stats (#379): `tole status` lists input / output / cache-read / cache-write / reasoning tokens per turn and the status JSON (`GET /sessions/{id}/status`, `tole_session_status`) gains an additive `usage_report.turns` array. Derived from the durable usage ledger, so old sessions read the same way and nothing is stored twice; a value no step reported is `n/a` / `null`, never 0, and no elapsed time is invented. +### Fixed +- `tole serve --transport mcp` is hardened to parity with the REST transport: the bearer token is compared in constant time and its scheme is case-insensitive; failed authentication is rate-limited per source IP (10 per minute, then `429`, same policy as REST); request bodies over 1 MiB are refused with `413` (declared or streamed, never buffered whole); a transient `accept()` error is logged and no longer ends the server and every session; and the no-op `timeout()` around the auth check is gone. The token comparison, header matching and limiter now live in one shared `auth` module used by both transports. + ### Changed - Memory loop injection hygiene (#384): the recalled batch is hard-capped at 24 items / 64 KiB and deduplicated; a malformed recall payload (non-array, or a hit without content) now fails the whole batch visibly instead of silently injecting fewer items; `tole serve`/MCP sessions report recall failures on stderr like the CLI does. diff --git a/crates/tole-cli/src/auth.rs b/crates/tole-cli/src/auth.rs new file mode 100644 index 0000000..6e07a9e --- /dev/null +++ b/crates/tole-cli/src/auth.rs @@ -0,0 +1,156 @@ +//! Bearer-token authentication shared by both `tole serve` transports +//! (REST and MCP over Streamable HTTP): constant-time comparison, header +//! matching and a per-IP failed-auth limiter, so the two faces cannot drift +//! apart again (the MCP face used to compare with `==`, never rate-limited +//! failures and had no body ceiling). + +use std::collections::HashMap; +use std::net::IpAddr; +use std::sync::Mutex; + +/// Failed-auth attempts allowed per source IP per window before the IP is +/// refused until the window rolls over (simple fixed-window limiter). +pub const AUTH_WINDOW_SECS: u64 = 60; +pub const MAX_AUTH_FAILURES: u32 = 10; + +/// Constant-time equality for equal-length slices; length differences are +/// not secret here (tokens have a fixed shape). +pub fn ct_eq(a: &[u8], b: &[u8]) -> bool { + a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0 +} + +/// Does an `Authorization` header value carry the bearer `token`? The scheme +/// is case-insensitive (RFC 7235) and surrounding whitespace is ignored. +pub fn bearer_matches(header: Option<&str>, token: &str) -> bool { + header + .and_then(|a| a.split_once(' ')) + .filter(|(scheme, _)| scheme.eq_ignore_ascii_case("bearer")) + .map(|(_, t)| ct_eq(t.trim().as_bytes(), token.as_bytes())) + .unwrap_or(false) +} + +fn now_secs() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// `(window_start_epoch, failure_count)` per source IP. +#[derive(Default)] +pub struct AuthLimiter { + failures: Mutex>, +} + +impl AuthLimiter { + pub fn new() -> Self { + Self::default() + } + + /// Has this IP used up its failures in the current window? The lock is + /// held only for the lookup, never across I/O. + pub fn is_limited(&self, ip: IpAddr) -> bool { + self.is_limited_at(ip, now_secs()) + } + + pub fn record_failure(&self, ip: IpAddr) { + self.record_failure_at(ip, now_secs()); + } + + fn lock(&self) -> std::sync::MutexGuard<'_, HashMap> { + // A poisoned lock must not disable the limiter. + self.failures.lock().unwrap_or_else(|p| p.into_inner()) + } + + fn is_limited_at(&self, ip: IpAddr, now: u64) -> bool { + matches!( + self.lock().get(&ip).map(|(window, count)| { + now.saturating_sub(*window) <= AUTH_WINDOW_SECS && *count >= MAX_AUTH_FAILURES + }), + Some(true) + ) + } + + fn record_failure_at(&self, ip: IpAddr, now: u64) { + let mut fails = self.lock(); + // Periodic sweep: drop entries whose window has fully expired, which + // bounds memory growth from unique source IPs. + fails.retain(|_, (window, _)| now.saturating_sub(*window) <= AUTH_WINDOW_SECS * 4); + let entry = fails.entry(ip).or_insert((now, 0)); + if now.saturating_sub(entry.0) > AUTH_WINDOW_SECS { + entry.0 = now; + entry.1 = 0; + } + entry.1 += 1; + } + + #[cfg(test)] + fn tracked(&self) -> usize { + self.lock().len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ip(n: u8) -> IpAddr { + IpAddr::from([10, 0, 0, n]) + } + + #[test] + fn ct_eq_matches_only_equal_slices() { + assert!(ct_eq(b"secret", b"secret")); + assert!(!ct_eq(b"secret", b"secreT")); + assert!(!ct_eq(b"secret", b"secre")); + assert!(ct_eq(b"", b"")); + } + + #[test] + fn bearer_matching_is_case_insensitive_and_exact() { + assert!(bearer_matches(Some("Bearer t0ken"), "t0ken")); + assert!(bearer_matches(Some("bearer t0ken "), "t0ken")); + assert!(bearer_matches(Some("BEARER t0ken"), "t0ken")); + assert!(!bearer_matches(Some("Bearer t0kex"), "t0ken")); + assert!(!bearer_matches(Some("Bearer t0ken2"), "t0ken")); + assert!(!bearer_matches(Some("Basic t0ken"), "t0ken")); + assert!(!bearer_matches(Some("t0ken"), "t0ken")); + assert!(!bearer_matches(None, "t0ken")); + } + + #[test] + fn the_limiter_refuses_after_the_cap_and_recovers_after_the_window() { + let l = AuthLimiter::new(); + let t0 = 1_000; + for _ in 0..MAX_AUTH_FAILURES - 1 { + l.record_failure_at(ip(1), t0); + } + assert!(!l.is_limited_at(ip(1), t0), "one below the cap"); + l.record_failure_at(ip(1), t0); + assert!(l.is_limited_at(ip(1), t0)); + assert!( + l.is_limited_at(ip(1), t0 + AUTH_WINDOW_SECS), + "still inside" + ); + assert!( + !l.is_limited_at(ip(1), t0 + AUTH_WINDOW_SECS + 1), + "window rolled over" + ); + // Other IPs are unaffected. + assert!(!l.is_limited_at(ip(2), t0)); + // A new failure after the window starts a fresh count. + l.record_failure_at(ip(1), t0 + AUTH_WINDOW_SECS + 5); + assert!(!l.is_limited_at(ip(1), t0 + AUTH_WINDOW_SECS + 5)); + } + + #[test] + fn stale_entries_are_swept_so_memory_stays_bounded() { + let l = AuthLimiter::new(); + for n in 0..50u8 { + l.record_failure_at(ip(n), 0); + } + assert_eq!(l.tracked(), 50); + l.record_failure_at(ip(200), AUTH_WINDOW_SECS * 4 + 1); + assert_eq!(l.tracked(), 1, "expired windows are dropped"); + } +} diff --git a/crates/tole-cli/src/lib.rs b/crates/tole-cli/src/lib.rs index 8ec164d..c64996b 100644 --- a/crates/tole-cli/src/lib.rs +++ b/crates/tole-cli/src/lib.rs @@ -4,6 +4,7 @@ //! drift from production behavior (CodeCora scan 2026-09-18). pub mod approvals; pub mod approver; +pub mod auth; pub mod config; pub mod config_apply; pub mod config_trust; diff --git a/crates/tole-cli/src/mcp_http.rs b/crates/tole-cli/src/mcp_http.rs index 874d775..21bfed5 100644 --- a/crates/tole-cli/src/mcp_http.rs +++ b/crates/tole-cli/src/mcp_http.rs @@ -18,19 +18,36 @@ use std::sync::Arc; use tole_cli::session_tools::SessionToolState; use tole_core::memory::MemoryConfig; -/// Collect a request body under `limit`. Any read error or timeout is an -/// `Err` — never silently an empty body (#277). +/// Largest request body accepted: the same ceiling as the REST transport. +const MAX_BODY_BYTES: usize = 1024 * 1024; + +/// Why a request body was not accepted. +#[derive(Debug, PartialEq, Eq)] +enum BodyError { + /// More than `MAX_BODY_BYTES` (declared or streamed): answered with 413. + TooLarge, + /// Read error or timeout: answered with 400. + Unreadable, +} + +/// Collect a request body under a time limit AND a size ceiling. Any read +/// error or timeout is an `Err` — never silently an empty body (#277) — and +/// a body over `max_bytes` is rejected as it streams in, so neither a +/// declared nor a chunked oversized body is ever buffered whole. async fn collect_body_bounded( body: B, limit: std::time::Duration, -) -> Result + max_bytes: usize, +) -> Result where B: hyper::body::Body, + B::Error: Into>, { - use http_body_util::BodyExt; - match tokio::time::timeout(limit, body.collect()).await { + use http_body_util::{BodyExt, LengthLimitError, Limited}; + match tokio::time::timeout(limit, Limited::new(body, max_bytes).collect()).await { Ok(Ok(c)) => Ok(c.to_bytes()), - _ => Err(()), + Ok(Err(e)) if e.downcast_ref::().is_some() => Err(BodyError::TooLarge), + _ => Err(BodyError::Unreadable), } } @@ -117,8 +134,22 @@ pub async fn run_mcp_http( let conn_sem = Arc::new(tokio::sync::Semaphore::new(MAX_CONNECTIONS)); let token = token.to_string(); + // Failed-auth limiter shared by every connection (same policy as REST). + let limiter = Arc::new(tole_cli::auth::AuthLimiter::new()); loop { - let (stream, _peer) = listener.accept().await?; + // A transient accept error (EMFILE, ECONNABORTED, ...) must not end + // the server and drop every session: report it and keep serving, + // with a short pause so resource exhaustion does not spin. + let (stream, peer) = match listener.accept().await { + Ok(x) => x, + Err(e) => { + eprintln!("tole serve (mcp): accept failed (continuing): {e}"); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + continue; + } + }; + let peer_ip = peer.ip(); + let limiter = Arc::clone(&limiter); let Ok(permit) = conn_sem.clone().try_acquire_owned() else { // At capacity: close immediately, same semantics as REST. drop(stream); @@ -132,26 +163,28 @@ pub async fn run_mcp_http( let io = TokioIo::new(stream); let svc_for_conn = svc.clone(); let token_for_conn = token.clone(); + let limiter_for_conn = Arc::clone(&limiter); let hyper_service = hyper::service::service_fn(move |req: hyper::Request| { let svc = svc_for_conn.clone(); let token = token_for_conn.clone(); + let limiter = Arc::clone(&limiter_for_conn); async move { use http_body_util::BodyExt; // Bearer-token gate BEFORE the MCP service sees - // anything (#137 auth surface), under the IO - // timeout so a trickling client can't pin the - // task (#190, mirrors the REST face). - let auth = tokio::time::timeout(SERVE_IO_TIMEOUT, async { - req.headers() - .get(hyper::header::AUTHORIZATION) - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.strip_prefix("Bearer ")) - .map(|t| t == token) - .unwrap_or(false) - }) - .await - .unwrap_or(false); + // anything (#137 auth surface): constant-time + // comparison, and a per-IP limiter on failures — + // both shared with the REST transport. The header + // read itself is bounded by `header_read_timeout` + // on the connection below. + let limited = limiter.is_limited(peer_ip); + let auth = !limited + && tole_cli::auth::bearer_matches( + req.headers() + .get(hyper::header::AUTHORIZATION) + .and_then(|v| v.to_str().ok()), + &token, + ); // Unified response body type for both arms. type RespBody = http_body_util::combinators::BoxBody< hyper::body::Bytes, @@ -163,12 +196,16 @@ pub async fn run_mcp_http( .boxed() } if !auth { + let (status, msg) = if limited { + (429, "{\"error\":\"too many failed auth attempts\"}") + } else { + limiter.record_failure(peer_ip); + (401, "{\"error\":\"unauthorized\"}") + }; let resp = hyper::Response::builder() - .status(401) + .status(status) .header("content-type", "application/json") - .body(box_full(hyper::body::Bytes::from( - "{\"error\":\"unauthorized\"}", - ))) + .body(box_full(hyper::body::Bytes::from(msg))) .map_err(|e| std::io::Error::other(e.to_string()))?; return Ok::<_, std::io::Error>(resp); } @@ -180,20 +217,46 @@ pub async fn run_mcp_http( // SSE RESPONSE streaming below is untouched. use tower_service::Service as _; let (parts, incoming) = req.into_parts(); - // A read error/timeout REJECTS the request: a - // truncated body must never execute as a valid - // empty one (#277). - let Ok(body_bytes) = collect_body_bounded(incoming, SERVE_IO_TIMEOUT).await - else { + // A declared body over the ceiling is refused before + // a single byte is read (chunked bodies are caught + // by the limit while streaming, below). + let declared = parts + .headers + .get(hyper::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()); + if declared.is_some_and(|n| n > MAX_BODY_BYTES) { let resp = hyper::Response::builder() - .status(400) + .status(413) .header("content-type", "application/json") .body(box_full(hyper::body::Bytes::from( - "{\"error\":\"request body read failed\"}", + "{\"error\":\"request body too large\"}", ))) .map_err(|e| std::io::Error::other(e.to_string()))?; return Ok::<_, std::io::Error>(resp); - }; + } + // A read error/timeout REJECTS the request: a + // truncated body must never execute as a valid + // empty one (#277). + let body_bytes = + match collect_body_bounded(incoming, SERVE_IO_TIMEOUT, MAX_BODY_BYTES) + .await + { + Ok(b) => b, + Err(e) => { + let (status, msg) = if e == BodyError::TooLarge { + (413, "{\"error\":\"request body too large\"}") + } else { + (400, "{\"error\":\"request body read failed\"}") + }; + let resp = hyper::Response::builder() + .status(status) + .header("content-type", "application/json") + .body(box_full(hyper::body::Bytes::from(msg))) + .map_err(|e| std::io::Error::other(e.to_string()))?; + return Ok::<_, std::io::Error>(resp); + } + }; let full_req: http::Request> = http::Request::from_parts(parts, http_body_util::Full::new(body_bytes)); let mut svc = svc; @@ -232,7 +295,7 @@ pub async fn run_mcp_http( #[cfg(test)] mod body_tests { - use super::collect_body_bounded; + use super::{collect_body_bounded, BodyError}; use hyper::body::{Body, Frame}; use std::pin::Pin; use std::task::{Context, Poll}; @@ -277,8 +340,13 @@ mod body_tests { stall: false, }, Duration::from_secs(5), + 1024, )); - assert!(r.is_err(), "a body read error must reject"); + assert_eq!( + r, + Err(BodyError::Unreadable), + "a body read error must reject" + ); } #[test] @@ -289,8 +357,13 @@ mod body_tests { stall: true, }, Duration::from_millis(50), + 1024, )); - assert!(r.is_err(), "a body read timeout must reject"); + assert_eq!( + r, + Err(BodyError::Unreadable), + "a body read timeout must reject" + ); } #[test] @@ -298,7 +371,27 @@ mod body_tests { let r = rt().block_on(collect_body_bounded( http_body_util::Full::new(hyper::body::Bytes::from("ok")), Duration::from_secs(5), + 1024, )); assert_eq!(r.unwrap().as_ref(), b"ok"); } + + #[test] + fn a_body_over_the_ceiling_is_rejected_while_streaming() { + let big = hyper::body::Bytes::from(vec![b'x'; 2048]); + let r = rt().block_on(collect_body_bounded( + http_body_util::Full::new(big), + Duration::from_secs(5), + 1024, + )); + assert_eq!(r, Err(BodyError::TooLarge)); + // Exactly at the ceiling is fine. + let exact = hyper::body::Bytes::from(vec![b'x'; 1024]); + let r = rt().block_on(collect_body_bounded( + http_body_util::Full::new(exact), + Duration::from_secs(5), + 1024, + )); + assert_eq!(r.unwrap().len(), 1024); + } } diff --git a/crates/tole-cli/src/serve.rs b/crates/tole-cli/src/serve.rs index 5fd384d..cc4852f 100644 --- a/crates/tole-cli/src/serve.rs +++ b/crates/tole-cli/src/serve.rs @@ -23,9 +23,8 @@ use anyhow::{Context, Result}; use serde_json::json; -use std::collections::HashMap; use std::io::{BufRead, BufReader, Read, Write}; -use std::net::{IpAddr, TcpListener, TcpStream}; +use std::net::{TcpListener, TcpStream}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -54,10 +53,6 @@ pub struct ServeConfig { /// stack + a socket; capping bounds thread/socket exhaustion (CodeCora /// scan-3 Wave-2 hardening). const MAX_CONNECTIONS: usize = 32; -/// Failed-auth attempts allowed per source IP per window before the IP is -/// dropped until the window rolls over (simple fixed-window limiter). -const AUTH_WINDOW_SECS: u64 = 60; -const MAX_AUTH_FAILURES: u32 = 10; struct State { sessions: SharedSessions, @@ -72,9 +67,9 @@ struct State { sessions_dir: Option, turnend: Vec, live_connections: std::sync::atomic::AtomicUsize, - /// (window_start_epoch, failure_count) per source IP — fixed-window - /// auth-failure limiter (brute-force hardening). - auth_failures: Mutex>, + /// Per-source-IP failed-auth limiter (brute-force hardening), shared + /// logic with the MCP transport. + auth_failures: tole_cli::auth::AuthLimiter, } /// Read/write ceilings per connection: a client that opens a socket and @@ -103,7 +98,7 @@ pub fn run_serve(cfg: ServeConfig) -> Result<()> { sessions_dir: cfg.sessions_dir.clone(), turnend: cfg.turnend.clone(), live_connections: std::sync::atomic::AtomicUsize::new(0), - auth_failures: Mutex::new(HashMap::new()), + auth_failures: tole_cli::auth::AuthLimiter::new(), }); eprintln!( "tole serve: listening on http://{addr} ({} allow pattern(s), plan_mode={})", @@ -295,13 +290,6 @@ fn read_request(stream: &TcpStream) -> Result { }) } -/// Constant-time byte comparison for secrets (a short-circuiting `==` -/// leaks the length of the matching prefix through timing). The length -/// itself is not secret here. -fn ct_eq(a: &[u8], b: &[u8]) -> bool { - a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0 -} - fn reason_phrase(status: u16) -> &'static str { match status { 200 => "OK", @@ -368,65 +356,20 @@ fn handle_conn(mut stream: TcpStream, state: Arc) { let peer_ip = stream.peer_addr().ok().map(|a| a.ip()); // Rate-limit: refuse when this IP has too many recent auth failures. - // The lock is dropped before any respond call (CodeCora: a stalled - // client must not hold the auth_failures mutex during a write). - if let Some(ip) = peer_ip { - let rate_limited = { - let fails = state - .auth_failures - .lock() - .unwrap_or_else(|p| p.into_inner()); - matches!( - fails.get(&ip).map(|(window, count)| { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0) - .saturating_sub(*window) - <= AUTH_WINDOW_SECS - && *count >= MAX_AUTH_FAILURES - }), - Some(true) - ) - }; - if rate_limited { - respond( - &mut stream, - 429, - &json!({"error": "too many failed auth attempts"}), - ); - return; - } + if peer_ip.is_some_and(|ip| state.auth_failures.is_limited(ip)) { + respond( + &mut stream, + 429, + &json!({"error": "too many failed auth attempts"}), + ); + return; } let authorized = req.path == "/health" - || req - .authorization - .as_deref() - .and_then(|a| a.split_once(' ')) - .filter(|(scheme, _)| scheme.eq_ignore_ascii_case("bearer")) - .map(|(_, t)| ct_eq(t.trim().as_bytes(), state.token.as_bytes())) - .unwrap_or(false); + || tole_cli::auth::bearer_matches(req.authorization.as_deref(), &state.token); if !authorized { if let Some(ip) = peer_ip { - let mut fails = state - .auth_failures - .lock() - .unwrap_or_else(|p| p.into_inner()); - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0); - // Periodic sweep: drop entries whose window has fully - // expired — bounds memory growth from unique source IPs - // (CodeCora scan round-2 on this PR). - fails.retain(|_, (window, _)| now.saturating_sub(*window) <= AUTH_WINDOW_SECS * 4); - let entry = fails.entry(ip).or_insert((now, 0)); - if now.saturating_sub(entry.0) > AUTH_WINDOW_SECS { - entry.0 = now; - entry.1 = 0; - } - entry.1 += 1; + state.auth_failures.record_failure(ip); } respond(&mut stream, 401, &json!({"error": "unauthorized"})); return; @@ -777,14 +720,6 @@ fn route(state: &State, method: &str, path: &str, body: &str) -> (u16, serde_jso mod tests { use super::*; - #[test] - fn ct_eq_matches_only_equal_slices() { - assert!(ct_eq(b"secret", b"secret")); - assert!(!ct_eq(b"secret", b"secreT")); - assert!(!ct_eq(b"secret", b"secre")); - assert!(ct_eq(b"", b"")); - } - #[test] fn reason_phrases_cover_every_status_the_routes_emit() { assert_eq!(reason_phrase(202), "Accepted"); diff --git a/crates/tole-cli/tests/mcp_http_hardening.rs b/crates/tole-cli/tests/mcp_http_hardening.rs new file mode 100644 index 0000000..d77b2f3 --- /dev/null +++ b/crates/tole-cli/tests/mcp_http_hardening.rs @@ -0,0 +1,168 @@ +//! Hardening of `tole serve --transport mcp` over the REAL binary, at parity +//! with the REST transport: a failed-auth limiter per source IP, a request +//! body ceiling (413, declared or streamed) and a working happy path. + +// Spawns the real `tole serve --transport mcp`, which needs both features. +#![cfg(all(feature = "shell-tools", feature = "mcp-http"))] + +use std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::process::{Child, Command, Stdio}; +use std::time::Duration; + +const TOKEN: &str = "t0ken"; +const INIT: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}"#; + +struct Server { + child: Child, + port: u16, +} + +impl Server { + fn spawn() -> Self { + let port = TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port(); + let child = Command::new(env!("CARGO_BIN_EXE_tole")) + .args([ + "serve", + "--transport", + "mcp", + "--port", + &port.to_string(), + "--token", + TOKEN, + ]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn tole serve --transport mcp"); + let me = Self { child, port }; + // Ready = it answers TCP. (A wrong-token probe would burn the + // limiter budget the tests below measure.) + for _ in 0..100 { + if TcpStream::connect(("127.0.0.1", port)).is_ok() { + return me; + } + std::thread::sleep(Duration::from_millis(100)); + } + panic!("mcp server never listened"); + } + + fn connect(&self) -> TcpStream { + let s = TcpStream::connect(("127.0.0.1", self.port)).unwrap(); + s.set_read_timeout(Some(Duration::from_secs(15))).unwrap(); + s + } + + /// POST `body` with `auth` (the whole header value, or none) and return + /// the status code. + fn post(&self, auth: Option<&str>, body: &str) -> u16 { + let mut s = self.connect(); + let auth = auth.map_or(String::new(), |a| format!("Authorization: {a}\r\n")); + let req = format!( + "POST /mcp HTTP/1.1\r\nHost: 127.0.0.1:{}\r\n{auth}Content-Type: application/json\r\nAccept: application/json, text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + self.port, + body.len() + ); + s.write_all(req.as_bytes()).unwrap(); + status_of(&mut s) + } +} + +impl Drop for Server { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +/// Read until the status line is complete and parse the code. +fn status_of(s: &mut TcpStream) -> u16 { + let mut buf = Vec::new(); + let mut chunk = [0u8; 4096]; + loop { + match s.read(&mut chunk) { + Ok(0) | Err(_) => break, + Ok(n) => { + buf.extend_from_slice(&chunk[..n]); + if buf.windows(2).any(|w| w == b"\r\n") { + break; + } + } + } + } + let line = String::from_utf8_lossy(&buf); + line.split_whitespace() + .nth(1) + .and_then(|c| c.parse().ok()) + .unwrap_or_else(|| panic!("no status line in {line:?}")) +} + +#[test] +fn a_valid_request_still_works_and_the_scheme_is_case_insensitive() { + let srv = Server::spawn(); + assert_eq!(srv.post(Some("Bearer t0ken"), INIT), 200); + assert_eq!(srv.post(Some("bearer t0ken"), INIT), 200); +} + +#[test] +fn wrong_tokens_are_401_then_the_ip_is_limited_even_with_the_right_token() { + let srv = Server::spawn(); + assert_eq!(srv.post(None, INIT), 401, "missing header"); + for i in 1..10 { + assert_eq!( + srv.post(Some("Bearer nope"), INIT), + 401, + "failure #{}", + i + 1 + ); + } + // The tenth failure used the budget up: from now on 429, and the right + // token does not get through until the window rolls over (REST parity). + assert_eq!(srv.post(Some("Bearer nope"), INIT), 429); + assert_eq!(srv.post(Some("Bearer t0ken"), INIT), 429); +} + +#[test] +fn a_declared_oversized_body_is_413_before_it_is_read() { + let srv = Server::spawn(); + let mut s = srv.connect(); + // Headers only: the server must answer from Content-Length alone. + let req = format!( + "POST /mcp HTTP/1.1\r\nHost: 127.0.0.1:{}\r\nAuthorization: Bearer {TOKEN}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + srv.port, + 2 * 1024 * 1024 + ); + s.write_all(req.as_bytes()).unwrap(); + assert_eq!(status_of(&mut s), 413); +} + +#[test] +fn a_chunked_oversized_body_is_413_while_it_streams() { + let srv = Server::spawn(); + let mut s = srv.connect(); + let head = format!( + "POST /mcp HTTP/1.1\r\nHost: 127.0.0.1:{}\r\nAuthorization: Bearer {TOKEN}\r\nContent-Type: application/json\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n", + srv.port + ); + s.write_all(head.as_bytes()).unwrap(); + let block = vec![b'x'; 64 * 1024]; + // 3 MiB in 64 KiB chunks; the server stops reading at the ceiling, so a + // later write may fail with a reset: that is fine, the answer is what + // counts. + for _ in 0..48 { + let chunk = format!("{:x}\r\n", block.len()); + if s.write_all(chunk.as_bytes()).is_err() + || s.write_all(&block).is_err() + || s.write_all(b"\r\n").is_err() + { + break; + } + } + let _ = s.write_all(b"0\r\n\r\n"); + assert_eq!(status_of(&mut s), 413); +}