use crate::server::ServerCore;
use std::io::{BufRead, BufReader, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
pub const MAX_BODY: usize = 1 << 20;
pub const BODY_CHUNK: usize = 64 * 1024;
pub const MAX_HEADERS: usize = 64;
pub const MAX_LINE: usize = 16 * 1024;
pub const READ_TIMEOUT: Duration = Duration::from_secs(30);
pub const WRITE_TIMEOUT: Duration = Duration::from_secs(10);
pub const REQUEST_DEADLINE: Duration = Duration::from_secs(60);
pub const CONNECTION_TTL: Duration = Duration::from_secs(300);
pub const MAX_REFUSALS_IN_FLIGHT: usize = 64;
pub const MAX_CONNECTIONS: usize = 256;
pub const HEALTH_ROUTE: &str = "/health";
pub const SSE_ANTI_BUFFERING: &str =
"Cache-Control: no-cache, no-transform\r\nX-Accel-Buffering: no\r\n";
pub struct BoundServer {
pub(crate) listener: TcpListener,
pub(crate) core: Arc<ServerCore>,
pub(crate) route: Arc<String>,
pub(crate) allow_origins: Arc<Vec<String>>,
pub(crate) read_timeout: Duration,
pub(crate) max_connections: usize,
pub(crate) connection_ttl: Duration,
pub(crate) live_connections: Arc<AtomicUsize>,
pub(crate) refusals_in_flight: Arc<AtomicUsize>,
pub(crate) shutdown: Arc<std::sync::atomic::AtomicBool>,
}
impl BoundServer {
pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
self.listener.local_addr()
}
#[must_use]
pub fn handle(&self) -> ServerHandle {
let addr = self
.listener
.local_addr()
.unwrap_or_else(|_| SocketAddr::from(([127, 0, 0, 1], 0)));
ServerHandle {
shutdown: Arc::clone(&self.shutdown),
addr,
}
}
pub fn serve(self) -> std::io::Result<()> {
for stream in self.listener.incoming() {
if self.shutdown.load(Ordering::Relaxed) {
return Ok(());
}
let Ok(mut stream) = stream else {
std::thread::sleep(Duration::from_millis(10));
continue;
};
let _ = stream.set_read_timeout(Some(self.read_timeout));
let _ = stream.set_write_timeout(Some(WRITE_TIMEOUT));
let _ = stream.set_nodelay(true);
let prev = self.live_connections.fetch_add(1, Ordering::Relaxed);
if prev >= self.max_connections {
self.live_connections.fetch_sub(1, Ordering::Relaxed);
let pending = Arc::clone(&self.refusals_in_flight);
let slot = pending.fetch_add(1, Ordering::Relaxed);
if slot < MAX_REFUSALS_IN_FLIGHT {
match stream.try_clone() {
Ok(mut worker) => {
let spawned = std::thread::Builder::new()
.name("litemcp-refus".into())
.spawn(move || {
let _guard = RefusalGuard(pending);
let _ = respond(
&mut worker,
"503 Service Unavailable",
"",
"application/json",
br#"{"jsonrpc":"2.0","id":null,"error":{"code":-32000,"message":"trop de connexions"}}"#,
false,
false,
);
shutdown_gracefully(&mut worker);
});
if spawned.is_ok() {
drop(stream);
continue;
}
self.refusals_in_flight.fetch_sub(1, Ordering::Relaxed);
}
Err(_) => {
self.refusals_in_flight.fetch_sub(1, Ordering::Relaxed);
}
}
} else {
pending.fetch_sub(1, Ordering::Relaxed);
}
let _ = respond(
&mut &stream,
"503 Service Unavailable",
"",
"application/json",
br#"{"jsonrpc":"2.0","id":null,"error":{"code":-32000,"message":"moteur sature"}}"#,
false,
false,
);
shutdown_gracefully(&mut stream);
continue;
}
let core = Arc::clone(&self.core);
let route = Arc::clone(&self.route);
let allow_origins = Arc::clone(&self.allow_origins);
let live = Arc::clone(&self.live_connections);
let shutdown = Arc::clone(&self.shutdown);
let Ok(spawn_stream) = stream.try_clone() else {
self.live_connections.fetch_sub(1, Ordering::Relaxed);
continue;
};
let spawned = std::thread::Builder::new()
.name("litemcp-conn".into())
.spawn(move || {
struct LiveGuard(Arc<AtomicUsize>);
impl Drop for LiveGuard {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::Relaxed);
}
}
let _guard = LiveGuard(live);
let _ = handle_connection(
spawn_stream,
&core,
&route,
&allow_origins,
self.connection_ttl,
&shutdown,
);
});
if spawned.is_err() {
self.live_connections.fetch_sub(1, Ordering::Relaxed);
let _ = respond(
&mut &stream,
"503 Service Unavailable",
"",
"application/json",
br#"{"jsonrpc":"2.0","id":null,"error":{"code":-32000,"message":"serveur sature"}}"#,
false,
false,
);
shutdown_gracefully(&mut stream);
}
}
Ok(())
}
}
#[derive(Clone)]
pub struct ServerHandle {
shutdown: Arc<std::sync::atomic::AtomicBool>,
addr: SocketAddr,
}
impl ServerHandle {
pub fn shutdown(&self) {
if !self.shutdown.swap(true, Ordering::Relaxed) {
let _ = std::net::TcpStream::connect_timeout(&self.addr, Duration::from_millis(250));
}
}
}
struct RefusalGuard(Arc<AtomicUsize>);
impl Drop for RefusalGuard {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::Relaxed);
}
}
fn shutdown_gracefully(stream: &mut TcpStream) {
let _ = stream.shutdown(std::net::Shutdown::Write);
let _ = stream.set_read_timeout(Some(Duration::from_millis(50)));
let mut sink = [0u8; 8192];
let mut drained = 0usize;
loop {
match stream.read(&mut sink) {
Ok(0) => break,
Ok(n) => {
drained += n;
if drained >= 64 * 1024 {
break;
}
}
Err(_) => break,
}
}
}
enum LineFail {
TooLong,
Slow,
Closed,
}
fn read_line_capped(
reader: &mut BufReader<TcpStream>,
out: &mut String,
scratch: &mut Vec<u8>,
cap: usize,
deadline: Instant,
) -> Result<(), LineFail> {
out.clear();
scratch.clear();
loop {
if Instant::now() >= deadline {
return Err(LineFail::Slow);
}
let available = match reader.fill_buf() {
Ok(bytes) => bytes,
Err(error)
if matches!(
error.kind(),
std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
) =>
{
return Err(LineFail::Slow);
}
Err(_) => return Err(LineFail::Closed),
};
if available.is_empty() {
return Err(LineFail::Closed);
}
if let Some(pos) = available.iter().position(|&byte| byte == b'\n') {
let chunk = &available[..=pos];
if scratch.len() + chunk.len() > cap {
return Err(LineFail::TooLong);
}
scratch.extend_from_slice(chunk);
reader.consume(pos + 1);
return match std::str::from_utf8(scratch) {
Ok(text) => {
if !text.ends_with("\r\n") {
return Err(LineFail::Closed);
}
out.push_str(text);
Ok(())
}
Err(_) => Err(LineFail::Closed), };
}
if scratch.len() + available.len() > cap {
return Err(LineFail::TooLong);
}
scratch.extend_from_slice(available);
let filled = available.len();
reader.consume(filled);
}
}
fn handle_connection(
stream: TcpStream,
core: &ServerCore,
route: &str,
allow_origins: &[String],
connection_ttl: Duration,
shutdown: &std::sync::atomic::AtomicBool,
) -> std::io::Result<()> {
let mut reader = BufReader::with_capacity(4 * 1024, stream.try_clone()?);
let mut stream = stream;
let mut line = String::new();
let mut scratch = Vec::new();
let connection_started = Instant::now();
loop {
if shutdown.load(Ordering::Relaxed) {
break;
}
if !connection_ttl.is_zero() && connection_started.elapsed() >= connection_ttl {
break;
}
let deadline = Instant::now() + REQUEST_DEADLINE;
match read_line_capped(&mut reader, &mut line, &mut scratch, MAX_LINE, deadline) {
Ok(()) => {}
Err(LineFail::TooLong) => {
respond(
&mut stream,
"431 Request Header Fields Too Large",
"",
"application/json",
br#"{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"ligne trop longue"}}"#,
false,
false,
)?;
return Ok(());
}
Err(_) => return Ok(()),
}
let request_line = line.strip_suffix("\r\n").unwrap_or(&line);
if request_line.contains(['\r', '\n']) {
respond_error(
&mut stream,
"400 Bad Request",
"fin de ligne invalide dans la ligne de requête",
"",
false,
)?;
return Ok(());
}
let mut parts = request_line.split(' ').filter(|part| !part.is_empty());
let verb = parts.next().unwrap_or("");
let full_path = parts.next().unwrap_or("");
let version = parts.next().unwrap_or("");
if verb.is_empty() || full_path.is_empty() || version.is_empty() || parts.next().is_some() {
respond_error(
&mut stream,
"400 Bad Request",
"ligne de requête invalide",
"",
verb == "HEAD",
)?;
return Ok(());
}
if verb.contains('\t')
|| full_path.contains('\t')
|| version.contains('\t')
|| verb.chars().any(|c| c.is_control())
|| full_path.chars().any(|c| c.is_control())
|| version.chars().any(|c| c.is_control())
{
respond_error(
&mut stream,
"400 Bad Request",
"caractère de contrôle dans la ligne de requête",
"",
verb == "HEAD",
)?;
return Ok(());
}
let http10 = version == "HTTP/1.0";
if !http10 && version != "HTTP/1.1" {
respond_error(
&mut stream,
"505 HTTP Version Not Supported",
"version HTTP non supportée",
"",
verb == "HEAD",
)?;
return Ok(());
}
let path = full_path.split('?').next().unwrap_or(full_path);
let is_post = verb == "POST";
let is_get = verb == "GET";
let is_head = verb == "HEAD";
let route_matched = path == route;
let health_probe = path == HEALTH_ROUTE;
let mut content_length: Option<usize> = None;
let mut saw_content_length = false;
let mut malformed_length = false;
let mut duplicate_length = false;
let mut transfer_encoding = false;
let mut origin_seen = false;
let mut origin_duplicate = false;
let mut origin_allowed = false;
let mut host_seen = false;
let mut host_duplicate = false;
let mut accept_json = false;
let mut accept_sse = false;
let mut expect_continue = false;
let mut content_type: Option<String> = None;
let mut content_type_duplicate = false;
let mut mcp_method: Option<String> = None;
let mut mcp_name: Option<String> = None;
let mut mcp_method_duplicate = false;
let mut mcp_name_duplicate = false;
let mut connection_close_seen = false;
let mut close = http10; let mut header_count = 0usize;
loop {
match read_line_capped(&mut reader, &mut line, &mut scratch, MAX_LINE, deadline) {
Ok(()) => {}
Err(LineFail::TooLong) => {
respond(
&mut stream,
"431 Request Header Fields Too Large",
"",
"application/json",
r#"{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"ligne trop longue"}}"#.as_bytes(),
false,
is_head,
)?;
return Ok(());
}
Err(_) => return Ok(()),
}
let header = line.trim_end();
if header.is_empty() {
break;
}
if header.starts_with(' ') || header.starts_with('\t') {
respond_error(
&mut stream,
"400 Bad Request",
"ligne d'en-tête obsolète (obs-fold)",
"",
is_head,
)?;
return Ok(());
}
header_count += 1;
if header_count > MAX_HEADERS {
respond(
&mut stream,
"431 Request Header Fields Too Large",
"",
"application/json",
r#"{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"trop d'en-têtes"}}"#.as_bytes(),
false,
is_head,
)?;
return Ok(());
}
if let Some((name_raw, value)) = header.split_once(':') {
if name_raw.ends_with(' ') || name_raw.ends_with('\t') {
respond_error(
&mut stream,
"400 Bad Request",
"espace avant deux-points",
"",
is_head,
)?;
return Ok(());
}
let name = name_raw.trim();
let value = value.trim();
if name.is_empty() {
respond_error(
&mut stream,
"400 Bad Request",
"nom d'en-tête vide",
"",
is_head,
)?;
return Ok(());
}
if name.eq_ignore_ascii_case("content-length") {
if saw_content_length {
duplicate_length = true;
} else {
saw_content_length = true;
if !value.is_empty() && value.bytes().all(|b| b.is_ascii_digit()) {
content_length = value.parse().ok();
if content_length.is_none() {
malformed_length = true;
}
} else {
malformed_length = true;
}
}
} else if name.eq_ignore_ascii_case("transfer-encoding") {
transfer_encoding = true; } else if name.eq_ignore_ascii_case("origin") {
if origin_seen {
origin_duplicate = true;
} else {
origin_seen = true;
origin_allowed =
allow_origins.iter().any(|o| o.eq_ignore_ascii_case(value));
}
} else if name.eq_ignore_ascii_case("host") {
if host_seen {
host_duplicate = true;
} else {
host_seen = true;
}
} else if name.eq_ignore_ascii_case("accept") {
accept_json |= header_has_media_type(value, "application/json");
accept_sse |= header_has_media_type(value, "text/event-stream");
} else if name.eq_ignore_ascii_case("content-type") {
if content_type.is_some() {
content_type_duplicate = true;
} else {
content_type = Some(value.to_string());
}
} else if name.eq_ignore_ascii_case("mcp-method") {
if mcp_method.is_some() {
mcp_method_duplicate = true;
} else {
mcp_method = Some(value.to_string());
}
} else if name.eq_ignore_ascii_case("mcp-name") {
if mcp_name.is_some() {
mcp_name_duplicate = true;
} else {
mcp_name = Some(value.to_string());
}
} else if name.eq_ignore_ascii_case("connection") {
let mut has_close = false;
let mut has_keep_alive = false;
for token in value.split(',') {
let token = token.trim();
if token.eq_ignore_ascii_case("close") {
has_close = true;
} else if token.eq_ignore_ascii_case("keep-alive") {
has_keep_alive = true;
}
}
if has_close {
close = true;
connection_close_seen = true;
}
if http10 && !connection_close_seen {
if has_keep_alive {
close = false;
}
}
} else if name.eq_ignore_ascii_case("expect")
&& header_has_token(value, "100-continue")
{
expect_continue = true;
}
} else {
respond_error(
&mut stream,
"400 Bad Request",
"en-tête malformé",
"",
is_head,
)?;
return Ok(());
}
}
if origin_duplicate {
respond_error(
&mut stream,
"400 Bad Request",
"en-tete Origin duplique",
"",
is_head,
)?;
return Ok(());
}
if content_type_duplicate || mcp_method_duplicate || mcp_name_duplicate {
respond_error(
&mut stream,
"400 Bad Request",
"en-tête à valeur unique dupliqué",
"",
is_head,
)?;
return Ok(());
}
if origin_seen && !origin_allowed {
respond_error(&mut stream, "403 Forbidden", "origin refusée", "", is_head)?;
return Ok(());
}
if !http10 {
if !host_seen {
respond_error(
&mut stream,
"400 Bad Request",
"en-tête Host requis",
"",
is_head,
)?;
return Ok(());
}
if host_duplicate {
respond_error(
&mut stream,
"400 Bad Request",
"en-tête Host dupliqué",
"",
is_head,
)?;
return Ok(());
}
}
if transfer_encoding || duplicate_length {
respond_error(
&mut stream,
"400 Bad Request",
"encodage de requête ambigu",
"",
is_head,
)?;
return Ok(());
}
let len = match content_length {
Some(len) => len,
None if malformed_length => {
respond_error(
&mut stream,
"400 Bad Request",
"Content-Length invalide",
"",
is_head,
)?;
return Ok(());
}
None => {
if !is_post {
0
} else {
respond_error(
&mut stream,
"411 Length Required",
"Content-Length requis",
"",
is_head,
)?;
return Ok(());
}
}
};
if len > MAX_BODY {
respond_error(
&mut stream,
"413 Payload Too Large",
"corps trop volumineux",
"",
is_head,
)?;
return Ok(());
}
if health_probe {
if is_get || is_head {
if len != 0 {
respond_error(
&mut stream,
"400 Bad Request",
"corps non admis sur la sonde de vie",
"",
is_head,
)?;
return Ok(());
}
respond(
&mut stream,
"200 OK",
"",
"text/plain",
b"ok",
!close,
is_head,
)?;
if close {
break;
}
continue;
}
respond_error(
&mut stream,
"405 Method Not Allowed",
"GET ou HEAD uniquement",
"Allow: GET, HEAD\r\n",
is_head,
)?;
return Ok(());
}
if !route_matched {
respond_error(&mut stream, "404 Not Found", "route inconnue", "", is_head)?;
return Ok(());
}
if !is_post {
respond_error(
&mut stream,
"405 Method Not Allowed",
"POST uniquement",
"Allow: POST\r\n",
is_head,
)?;
return Ok(());
}
if let Some(value) = &content_type {
let media_type = value
.split(';')
.next()
.unwrap_or("")
.trim()
.to_ascii_lowercase();
if media_type != "application/json" {
respond_error(
&mut stream,
"415 Unsupported Media Type",
"Content-Type attendu : application/json",
"",
is_head,
)?;
return Ok(());
}
}
if expect_continue && len > 0 {
stream.write_all(b"HTTP/1.1 100 Continue\r\n\r\n")?;
stream.flush()?;
}
let body = match read_body(&mut reader, len, deadline) {
Ok(body) => body,
Err(_) => {
let _ = respond_error(
&mut stream,
"400 Bad Request",
"corps tronque ou trop lent",
"",
false,
);
return Ok(());
}
};
let dispatch_result = std::panic::catch_unwind(AssertUnwindSafe(|| {
let headers = crate::server::McpHeaders {
method: mcp_method.clone(),
name: mcp_name.clone(),
};
core.dispatch_with_headers(&body, Some(&headers))
}));
let reply = match dispatch_result {
Ok(Some(reply)) => reply,
Ok(None) => {
respond(
&mut stream,
"202 Accepted",
"",
"application/json",
b"",
!close,
false,
)?;
if close {
break;
}
continue;
}
Err(_) => {
let request_id = extract_request_id(&body);
let payload = serde_json::json!({
"jsonrpc": "2.0",
"id": request_id,
"error": { "code": -32603, "message": "erreur interne" }
});
respond(
&mut stream,
"200 OK",
"",
"application/json",
payload.to_string().as_bytes(),
false,
false,
)?;
return Ok(());
}
};
if accept_sse && !accept_json {
let mut event = Vec::with_capacity(reply.len() + 8);
event.extend_from_slice(b"data: ");
event.extend_from_slice(&reply);
event.extend_from_slice(b"\n\n");
respond(
&mut stream,
"200 OK",
SSE_ANTI_BUFFERING,
"text/event-stream",
&event,
false,
false,
)?;
break; }
respond(
&mut stream,
"200 OK",
"",
"application/json",
&reply,
!close,
false,
)?;
if close {
break;
}
}
Ok(())
}
fn respond_error(
stream: &mut TcpStream,
status: &str,
message: &str,
extra_headers: &str,
head: bool,
) -> std::io::Result<()> {
let payload: &[u8] = &error_payload(message);
respond(
stream,
status,
extra_headers,
"application/json",
payload,
false,
head,
)?;
shutdown_gracefully(stream);
Ok(())
}
fn error_payload(message: &str) -> Vec<u8> {
serde_json::json!({
"jsonrpc": "2.0",
"id": null,
"error": { "code": -32600, "message": message }
})
.to_string()
.into_bytes()
}
fn read_body(
reader: &mut BufReader<TcpStream>,
len: usize,
deadline: Instant,
) -> std::io::Result<Vec<u8>> {
if len == 0 {
return Ok(Vec::new());
}
let mut body = Vec::with_capacity(len.min(BODY_CHUNK));
let mut chunk = vec![0u8; BODY_CHUNK.min(len)];
let mut remaining = len;
while remaining > 0 {
let wanted = remaining.min(chunk.len());
if Instant::now() >= deadline {
return Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"deadline dépassée",
));
}
let n = reader.read(&mut chunk[..wanted])?;
if n == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"corps tronqué",
));
}
body.extend_from_slice(&chunk[..n]);
remaining -= n;
}
Ok(body)
}
fn respond(
stream: &mut impl Write,
status: &str,
extra_headers: &str,
content_type: &str,
body: &[u8],
keep_alive: bool,
head: bool,
) -> std::io::Result<()> {
let connection = if keep_alive { "keep-alive" } else { "close" };
let mut buf = Vec::with_capacity(
96 + status.len() + extra_headers.len() + content_type.len() + body.len(),
);
write!(
buf,
"HTTP/1.1 {status}\r\n{extra_headers}Content-Type: {content_type}\r\nContent-Length: {}\r\nConnection: {connection}\r\n\r\n",
body.len()
)?;
if !head {
buf.extend_from_slice(body);
}
stream.write_all(&buf)?;
stream.flush()
}
fn extract_request_id(body: &[u8]) -> serde_json::Value {
serde_json::from_slice::<serde_json::Value>(body)
.ok()
.and_then(|value| value.get("id").cloned())
.unwrap_or(serde_json::Value::Null)
}
fn header_has_token(value: &str, token: &str) -> bool {
value
.split(',')
.map(|part| part.trim())
.any(|part| part.eq_ignore_ascii_case(token))
}
fn header_has_media_type(value: &str, media_type: &str) -> bool {
value.split(',').map(|part| part.trim()).any(|part| {
part.split(';')
.next()
.unwrap_or("")
.trim()
.eq_ignore_ascii_case(media_type)
})
}