From 90400701fad890d4e383358825a3fc730e57c56d Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sat, 25 Jul 2026 12:56:23 +0200 Subject: [PATCH] Bound local HTTP resource use --- src/server.rs | 84 ++++++++++++++++++++++++++++++++++++++++++---- src/server/http.rs | 42 +++++++++++++++++++---- 2 files changed, 114 insertions(+), 12 deletions(-) diff --git a/src/server.rs b/src/server.rs index aab41b8..70d0ebd 100644 --- a/src/server.rs +++ b/src/server.rs @@ -34,13 +34,16 @@ use std::fs::File; use std::io::{Read, Write}; use std::net::{TcpListener, TcpStream}; use std::path::PathBuf; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; const MAX_HEADER_BYTES: usize = 64 * 1024; const MAX_BODY_BYTES: usize = 64 * 1024 * 1024; +const MAX_CONNECTIONS: usize = 64; +const HTTP_IO_TIMEOUT: Duration = Duration::from_secs(10); +const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); const PREFILL_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5); const TOOLS_PROMPT: &str = "## Tools\n\n\ You have access to a set of tools to help answer the user question. You can invoke tools by writing a \"<｜DSML｜tool_calls>\" block like the following:\n\n\ @@ -71,6 +74,7 @@ struct State { models_path: PathBuf, cache_path: PathBuf, sequence: AtomicU64, + connections: Arc, tool_memory: Mutex>, metrics: Arc, } @@ -226,6 +230,7 @@ impl ServerHandle { models_path, cache_path, sequence: AtomicU64::new(0), + connections: Arc::new(AtomicUsize::new(0)), tool_memory: Mutex::new(HashMap::new()), metrics: Arc::clone(&metrics), }); @@ -257,10 +262,20 @@ fn serve(listener: TcpListener, state: Arc, stop: Arc) { while !stop.load(Ordering::Relaxed) { match listener.accept() { Ok((stream, _)) => { + let Some(slot) = ConnectionSlot::acquire(&state.connections) else { + continue; + }; let state = Arc::clone(&state); - let _ = thread::Builder::new() - .name("http-request".into()) - .spawn(move || handle(stream, &state)); + if let Err(error) = + thread::Builder::new() + .name("http-request".into()) + .spawn(move || { + let _slot = slot; + handle(stream, &state); + }) + { + eprintln!("DS4Server: endpoint request thread failed: {error}"); + } } Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { thread::sleep(Duration::from_millis(50)); @@ -273,10 +288,29 @@ fn serve(listener: TcpListener, state: Arc, stop: Arc) { } } +struct ConnectionSlot(Arc); + +impl ConnectionSlot { + fn acquire(active: &Arc) -> Option { + active + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { + (count < MAX_CONNECTIONS).then_some(count + 1) + }) + .ok()?; + Some(Self(Arc::clone(active))) + } +} + +impl Drop for ConnectionSlot { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Relaxed); + } +} + fn handle(mut stream: TcpStream, state: &State) { - let _ = stream.set_read_timeout(Some(Duration::from_secs(30))); + let _ = stream.set_write_timeout(Some(HTTP_IO_TIMEOUT)); let started = Instant::now(); - let request = match read_request(&mut stream) { + let request = match read_request(&mut stream, started + HTTP_REQUEST_TIMEOUT) { Ok(request) => request, Err(error) => { state.metrics.http_started("", 0); @@ -518,6 +552,41 @@ fn empty_arguments() -> String { mod tests { use super::*; + #[test] + fn connection_slots_enforce_the_limit_and_release_on_drop() { + let active = Arc::new(AtomicUsize::new(0)); + let mut slots = (0..MAX_CONNECTIONS) + .map(|_| ConnectionSlot::acquire(&active).unwrap()) + .collect::>(); + + assert!(ConnectionSlot::acquire(&active).is_none()); + slots.pop(); + assert!(ConnectionSlot::acquire(&active).is_some()); + } + + #[test] + fn request_deadline_stops_slow_clients() { + let listener = match TcpListener::bind("127.0.0.1:0") { + Ok(listener) => listener, + Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => return, + Err(error) => panic!("could not start test server: {error}"), + }; + let address = listener.local_addr().unwrap(); + let client = thread::spawn(move || { + let mut stream = TcpStream::connect(address).unwrap(); + stream.write_all(b"GET /v1/models HTTP/1.1\r\n").unwrap(); + thread::sleep(Duration::from_millis(200)); + }); + let (mut stream, _) = listener.accept().unwrap(); + + let error = read_request(&mut stream, Instant::now() + Duration::from_millis(50)) + .err() + .unwrap(); + + client.join().unwrap(); + assert_eq!(error, "HTTP request timed out"); + } + #[test] fn tool_prompt_matches_ds4_server_text() { assert_eq!(TOOLS_PROMPT.len(), 1_183); @@ -612,6 +681,7 @@ mod tests { models_path: PathBuf::new(), cache_path: PathBuf::new(), sequence: AtomicU64::new(0), + connections: Arc::new(AtomicUsize::new(0)), tool_memory: Mutex::new(HashMap::new()), metrics, }; @@ -647,6 +717,7 @@ mod tests { models_path: PathBuf::new(), cache_path: PathBuf::new(), sequence: AtomicU64::new(0), + connections: Arc::new(AtomicUsize::new(0)), tool_memory: Mutex::new(HashMap::new()), metrics, }; @@ -673,6 +744,7 @@ mod tests { models_path: PathBuf::new(), cache_path: PathBuf::new(), sequence: AtomicU64::new(0), + connections: Arc::new(AtomicUsize::new(0)), tool_memory: Mutex::new(HashMap::new()), metrics, }; diff --git a/src/server/http.rs b/src/server/http.rs index 0fa630b..a80b3e1 100644 --- a/src/server/http.rs +++ b/src/server/http.rs @@ -72,21 +72,27 @@ pub(super) fn send_response( .map_err(|error| error.to_string()) } -pub(super) fn read_request(stream: &mut TcpStream) -> Result { +pub(super) fn read_request( + stream: &mut TcpStream, + deadline: Instant, +) -> Result { let mut bytes = Vec::new(); let header_end = loop { - if bytes.len() >= MAX_HEADER_BYTES { - return Err("HTTP headers are too large".into()); - } let mut chunk = [0_u8; 4096]; - let read = stream.read(&mut chunk).map_err(|error| error.to_string())?; + let read = read_with_deadline(stream, &mut chunk, deadline)?; if read == 0 { return Err("bad HTTP request".into()); } bytes.extend_from_slice(&chunk[..read]); if let Some(end) = find_header_end(&bytes) { + if end > MAX_HEADER_BYTES { + return Err("HTTP headers are too large".into()); + } break end; } + if bytes.len() >= MAX_HEADER_BYTES { + return Err("HTTP headers are too large".into()); + } }; let header = std::str::from_utf8(&bytes[..header_end]) .map_err(|_| "HTTP headers are not UTF-8".to_owned())?; @@ -115,7 +121,7 @@ pub(super) fn read_request(stream: &mut TcpStream) -> Result Result Result { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err("HTTP request timed out".into()); + } + stream + .set_read_timeout(Some(remaining.min(HTTP_IO_TIMEOUT))) + .map_err(|error| error.to_string())?; + stream.read(buffer).map_err(|error| { + if matches!( + error.kind(), + std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock + ) { + "HTTP request timed out".into() + } else { + error.to_string() + } + }) +} + fn find_header_end(bytes: &[u8]) -> Option { bytes .windows(4)