use anyhow::{Context, Result};
use std::collections::BTreeMap;
use std::io::{BufRead, BufReader, Read, Write};
use std::net::{Ipv4Addr, SocketAddr, TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::time::Duration;
pub const DEFAULT_PORT: u16 = 7365;
pub const PORT_ENV: &str = "SEMLITH_PORT";
pub const TOKEN_HEADER: &str = "Semlith-Token";
const TOKEN_HEADER_KEY: &str = "semlith-token";
pub const MCP_PATH: &str = "/mcp";
const WORKERS: usize = 8;
const MAX_HEADERS: usize = 16 * 1024;
const IDLE_PER_PEER: usize = 8;
const IDLE_GRACE: Duration = Duration::from_millis(250);
const MAX_WAITING: usize = 48;
const HEAD_POLL: Duration = Duration::from_secs(1);
const MAX_BODY: usize = 1024 * 1024;
const IO_TIMEOUT: Duration = Duration::from_secs(30);
const REQUEST_DEADLINE: Duration = Duration::from_secs(10);
const MAX_CONNECTIONS: usize = 32;
static WAKE: AtomicBool = AtomicBool::new(false);
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Refusal {
Unauthorized,
ForeignHost,
Malformed,
CrossOrigin,
BadContentType,
OutsideBoundary,
DeniedPath,
PrivateAddress,
}
impl Refusal {
pub fn as_str(self) -> &'static str {
match self {
Refusal::Unauthorized => "unauthorized",
Refusal::ForeignHost => "foreign host",
Refusal::Malformed => "malformed",
Refusal::CrossOrigin => "cross-origin",
Refusal::BadContentType => "bad-content-type",
Refusal::OutsideBoundary => "outside-boundary",
Refusal::DeniedPath => "denied-path",
Refusal::PrivateAddress => "private-address",
}
}
}
#[derive(Clone)]
pub struct Request {
pub method: String,
pub path: String,
pub query: BTreeMap<String, String>,
pub headers: BTreeMap<String, String>,
pub body: Vec<u8>,
}
impl Request {
pub fn query(&self, key: &str) -> Option<&str> {
self.query.get(key).map(String::as_str)
}
pub fn query_all(&self, key: &str) -> Vec<String> {
match self.query.get(key) {
Some(v) => v.split('\u{1}').map(str::to_string).collect(),
None => Vec::new(),
}
}
pub fn header(&self, name: &str) -> Option<&str> {
self.headers.get(name).map(String::as_str)
}
pub fn json(&self) -> Result<serde_json::Value> {
anyhow::ensure!(
json_body(self.header("content-type")),
"this route reads JSON; send Content-Type: application/json"
);
Ok(serde_json::from_slice(&self.body)?)
}
}
fn json_body(content_type: Option<&str>) -> bool {
content_type
.map(|value| value.split(';').next().unwrap_or("").trim())
.is_some_and(|media| media.eq_ignore_ascii_case("application/json"))
}
pub struct Chunks<'a> {
out: &'a mut dyn Write,
}
impl Chunks<'_> {
pub fn send(&mut self, line: &str) -> std::io::Result<()> {
write!(self.out, "{:x}\r\n", line.len())?;
self.out.write_all(line.as_bytes())?;
self.out.write_all(b"\r\n")?;
self.out.flush()
}
}
type Streamer = Box<dyn FnOnce(&mut Chunks) -> std::io::Result<()> + Send>;
enum Body {
Bytes(Vec<u8>),
Stream(Streamer),
}
pub struct Response {
status: u16,
content_type: &'static str,
extra: Vec<(String, String)>,
csp: Option<&'static str>,
body: Body,
}
const PAGE_CSP: &str = "default-src 'self'; \
img-src 'self' data:; \
style-src 'self'; \
script-src 'self'; \
font-src 'self'; \
connect-src 'self'; \
form-action 'none'; \
frame-ancestors 'none'; \
base-uri 'none'";
impl Response {
pub fn new(status: u16, content_type: &'static str, body: Vec<u8>) -> Self {
Self {
status,
content_type,
extra: Vec::new(),
csp: None,
body: Body::Bytes(body),
}
}
pub fn json(value: &serde_json::Value) -> Self {
Self::new(
200,
"application/json; charset=utf-8",
serde_json::to_vec(value).unwrap_or_else(|_| b"{}".to_vec()),
)
}
pub fn status(&self) -> u16 {
self.status
}
pub fn json_value(&self) -> Option<serde_json::Value> {
match &self.body {
Body::Bytes(bytes) => serde_json::from_slice(bytes).ok(),
Body::Stream(_) => None,
}
}
pub fn error(status: u16, message: &str) -> Self {
Self::new(
status,
"application/json; charset=utf-8",
serde_json::to_vec(&serde_json::json!({ "error": message }))
.unwrap_or_else(|_| b"{}".to_vec()),
)
}
pub fn asset(content_type: &'static str, bytes: &'static [u8]) -> Self {
Self::new(200, content_type, bytes.to_vec())
}
pub fn stream(f: impl FnOnce(&mut Chunks) -> std::io::Result<()> + Send + 'static) -> Self {
Self {
status: 200,
content_type: "application/x-ndjson; charset=utf-8",
extra: Vec::new(),
csp: None,
body: Body::Stream(Box::new(f)),
}
}
pub fn header(mut self, name: &str, value: impl Into<String>) -> Self {
self.extra.push((name.to_string(), value.into()));
self
}
pub fn with_csp(mut self, csp: &'static str) -> Self {
self.csp = Some(csp);
self
}
pub fn csp(&self) -> &'static str {
self.csp.unwrap_or(PAGE_CSP)
}
}
pub type Handler = Arc<dyn Fn(&Request) -> Response + Send + Sync>;
struct Late {
stream: TcpStream,
due: std::time::Instant,
_in_flight: InFlight,
}
struct InFlight(
Arc<std::sync::atomic::AtomicUsize>,
#[allow(dead_code)] crate::priority::Embedding,
);
impl Drop for InFlight {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::Relaxed);
}
}
impl InFlight {
fn taken(count: &Arc<std::sync::atomic::AtomicUsize>) -> Self {
count.fetch_add(1, Ordering::Relaxed);
Self(Arc::clone(count), crate::priority::request())
}
fn move_to(&mut self, count: &Arc<std::sync::atomic::AtomicUsize>) {
count.fetch_add(1, Ordering::Relaxed);
std::mem::replace(&mut self.0, Arc::clone(count)).fetch_sub(1, Ordering::Relaxed);
}
}
struct Ready {
stream: TcpStream,
reader: BufReader<TcpStream>,
request: Request,
in_flight: InFlight,
accepted: std::time::Instant,
}
#[derive(Default)]
struct HeadState {
heard: AtomicBool,
evicted: AtomicBool,
}
struct Waiting {
state: Arc<HeadState>,
peer: std::net::IpAddr,
since: std::time::Instant,
handle: TcpStream,
slot: InFlight,
}
#[derive(Default)]
struct Pending(std::collections::VecDeque<Waiting>);
impl Pending {
fn admit(
&mut self,
peer: std::net::IpAddr,
handle: TcpStream,
slot: InFlight,
) -> Arc<HeadState> {
let state = Arc::new(HeadState::default());
self.0.push_back(Waiting {
state: Arc::clone(&state),
peer,
since: std::time::Instant::now(),
handle,
slot,
});
let idle = |w: &Waiting| {
w.peer == peer
&& !w.state.heard.load(Ordering::Relaxed)
&& w.since.elapsed() >= IDLE_GRACE
};
while self.0.iter().filter(|w| idle(w)).count() > IDLE_PER_PEER {
let Some(at) = self.0.iter().position(idle) else {
break;
};
self.evict(at);
}
while self.0.len() > MAX_WAITING {
let older = self.0.len() - 1;
let at = self
.0
.iter()
.take(older)
.position(|w| !w.state.heard.load(Ordering::Relaxed))
.unwrap_or(0);
self.evict(at);
}
state
}
fn evict(&mut self, at: usize) {
if let Some(gone) = self.0.remove(at) {
gone.state.evicted.store(true, Ordering::Relaxed);
trace(&gone.handle, || {
"closed while waiting, to make room".to_string()
});
let _ = gone.handle.shutdown(std::net::Shutdown::Both);
}
}
fn release(&mut self, state: &Arc<HeadState>) -> Option<InFlight> {
let at = self.0.iter().position(|w| Arc::ptr_eq(&w.state, state))?;
self.0.remove(at).map(|w| w.slot)
}
}
#[derive(Debug)]
struct HeadRefused(u16);
impl std::fmt::Display for HeadRefused {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "request head refused with {}", self.0)
}
}
impl std::error::Error for HeadRefused {}
#[derive(Default)]
struct Agent {
current: String,
previous: Option<(String, std::time::Instant)>,
}
const KEY_GRACE: Duration = Duration::from_secs(15 * 60);
pub struct Server {
listener: TcpListener,
port: u16,
token: Arc<Mutex<String>>,
agent: Arc<Mutex<Agent>>,
mcp_open: Arc<AtomicBool>,
penalty: Arc<Mutex<Penalty>>,
}
impl Server {
pub fn bind(port: u16) -> Result<Self> {
let addr = SocketAddr::from((Ipv4Addr::LOCALHOST, port));
let listener = TcpListener::bind(addr).with_context(|| {
format!(
"port {port} is already in use — stop whatever is listening on it, \
or choose another with --port or {PORT_ENV}"
)
})?;
let port = listener.local_addr().map(|a| a.port()).unwrap_or(port);
Ok(Self {
listener,
port,
token: Arc::new(Mutex::new(new_token())),
agent: Arc::new(Mutex::new(Agent::default())),
mcp_open: Arc::new(AtomicBool::new(true)),
penalty: Arc::new(Mutex::new(Penalty::default())),
})
}
pub fn set_agent_key(&self, key: &str) {
self.agent.lock().unwrap_or_else(|e| e.into_inner()).current = key.to_string();
}
pub fn agent_key(&self) -> String {
self.agent
.lock()
.unwrap_or_else(|e| e.into_inner())
.current
.clone()
}
pub fn rotate_agent(&self, key: &str, now: bool) {
let mut agent = self.agent.lock().unwrap_or_else(|e| e.into_inner());
let was = std::mem::replace(&mut agent.current, key.to_string());
agent.previous = if now || was.is_empty() {
None
} else {
Some((was, std::time::Instant::now() + KEY_GRACE))
};
}
pub fn key_grace(&self) -> Option<Duration> {
let agent = self.agent.lock().unwrap_or_else(|e| e.into_inner());
agent
.previous
.as_ref()
.map(|(_, until)| until.saturating_duration_since(std::time::Instant::now()))
.filter(|left| !left.is_zero())
}
pub fn mcp_open(&self) -> bool {
self.mcp_open.load(Ordering::Relaxed)
}
pub fn set_mcp_open(&self, open: bool) {
self.mcp_open.store(open, Ordering::Relaxed);
}
pub fn port(&self) -> u16 {
self.port
}
pub fn token(&self) -> String {
self.token.lock().unwrap_or_else(|e| e.into_inner()).clone()
}
pub fn rotate(&self) -> String {
let fresh = new_token();
*self.token.lock().unwrap_or_else(|e| e.into_inner()) = fresh.clone();
fresh
}
pub fn url(&self) -> String {
format!("http://127.0.0.1:{}/?token={}", self.port, self.token())
}
pub fn serve(
&self,
handler: Handler,
stop: &'static AtomicBool,
refused: impl Fn(Refusal) + Send + Sync + 'static,
) -> Result<()> {
let (tx, rx) = mpsc::channel::<Ready>();
let pool = Arc::new(Mutex::new(Some(tx)));
let rx = Arc::new(Mutex::new(rx));
let refused = Arc::new(refused);
let live = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let pending = Arc::new(Mutex::new(Pending::default()));
let waiting = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let (late_tx, late_rx) = mpsc::channel::<Late>();
let waiting_room = std::thread::spawn(move || {
while let Ok(late) = late_rx.recv() {
let now = std::time::Instant::now();
if late.due > now {
std::thread::sleep(late.due - now);
}
let mut stream = late.stream;
let _ = write_response(&mut stream, Response::new(401, "text/plain", Vec::new()));
}
});
let port = self.port;
let spawn_worker = {
let rx = Arc::clone(&rx);
let late_tx = late_tx.clone();
let penalty = Arc::clone(&self.penalty);
let handler = Arc::clone(&handler);
let token = Arc::clone(&self.token);
let agent = Arc::clone(&self.agent);
let mcp_open = Arc::clone(&self.mcp_open);
let refused = Arc::clone(&refused);
move || {
let rx = Arc::clone(&rx);
let late_tx = late_tx.clone();
let penalty = Arc::clone(&penalty);
let handler = Arc::clone(&handler);
let token = Arc::clone(&token);
let agent = Arc::clone(&agent);
let mcp_open = Arc::clone(&mcp_open);
let refused = Arc::clone(&refused);
std::thread::spawn(move || {
loop {
let next = {
let guard = rx.lock().unwrap_or_else(|e| e.into_inner());
guard.recv()
};
let Ok(ready) = next else {
return;
};
let hold = |stream, in_flight| {
let delay = penalty.lock().unwrap_or_else(|e| e.into_inner()).charge();
let _ = late_tx.send(Late {
stream,
due: std::time::Instant::now() + delay,
_in_flight: in_flight,
});
};
let auth = {
let agent = agent.lock().unwrap_or_else(|e| e.into_inner());
Auth {
token: token.lock().unwrap_or_else(|e| e.into_inner()).clone(),
agent: agent.current.clone(),
previous: agent.previous.clone(),
mcp_open: mcp_open.load(Ordering::Relaxed),
port,
}
};
if let Some(class) = answer(ready, &handler, &auth, &hold) {
refused(class);
}
}
})
}
};
let mut workers: Vec<_> = (0..WORKERS).map(|_| spawn_worker()).collect();
let waker = {
let port = self.port;
std::thread::spawn(move || {
while !stop.load(Ordering::Relaxed) && !WAKE.load(Ordering::Relaxed) {
std::thread::sleep(Duration::from_millis(250));
}
let _ = TcpStream::connect(("127.0.0.1", port));
})
};
while !stop.load(Ordering::Relaxed) {
match self.listener.accept() {
Ok((stream, _)) => {
if stop.load(Ordering::Relaxed) {
break;
}
let _ = stream.set_read_timeout(Some(IO_TIMEOUT));
let _ = stream.set_write_timeout(Some(IO_TIMEOUT));
let busy = live.load(Ordering::Relaxed);
trace(&stream, || format!("accepted, {busy} live"));
if busy >= MAX_CONNECTIONS {
let in_flight = InFlight::taken(&live);
std::thread::spawn(move || {
refuse_busy(stream);
drop(in_flight);
});
continue;
}
let accepted = std::time::Instant::now();
let (Ok(peer), Ok(handle)) = (stream.peer_addr(), stream.try_clone()) else {
continue;
};
let state = pending.lock().unwrap_or_else(|e| e.into_inner()).admit(
peer.ip(),
handle,
InFlight::taken(&waiting),
);
let pool = Arc::clone(&pool);
let pending = Arc::clone(&pending);
let refused = Arc::clone(&refused);
let live = Arc::clone(&live);
std::thread::spawn(move || {
let mut reader = BufReader::new(stream);
let head = read_head(&mut reader, accepted, &state);
let Some(mut in_flight) = pending
.lock()
.unwrap_or_else(|e| e.into_inner())
.release(&state)
else {
return;
};
match head {
Ok(Some(request)) => {
in_flight.move_to(&live);
let _ = reader.get_ref().set_read_timeout(Some(IO_TIMEOUT));
if let (Ok(stream), Some(tx)) = (
reader.get_ref().try_clone(),
pool.lock().unwrap_or_else(|e| e.into_inner()).as_ref(),
) {
let _ = tx.send(Ready {
stream,
reader,
request,
in_flight,
accepted,
});
}
}
Ok(None) => {
trace(reader.get_ref(), || "closed before a request".to_string())
}
Err(e) => {
trace(reader.get_ref(), || format!("bad request head: {e:#}"));
let status = e.downcast_ref::<HeadRefused>().map_or(400, |r| r.0);
let _ = write_response(
reader.get_mut(),
Response::new(status, "text/plain", Vec::new()),
);
drop(in_flight);
refused(Refusal::Malformed);
}
}
});
for worker in &mut workers {
if worker.is_finished() {
*worker = spawn_worker();
}
}
}
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => {
WAKE.store(true, Ordering::Relaxed);
let _ = waker.join();
return Err(e).context("accepting a connection");
}
}
}
WAKE.store(true, Ordering::Relaxed);
let _ = waker.join();
pool.lock().unwrap_or_else(|e| e.into_inner()).take();
for worker in workers {
let _ = worker.join();
}
drop(spawn_worker);
drop(late_tx);
let _ = waiting_room.join();
Ok(())
}
}
struct Auth {
token: String,
agent: String,
previous: Option<(String, std::time::Instant)>,
mcp_open: bool,
port: u16,
}
impl Auth {
fn is_agent(&self, bearer: &str) -> bool {
if bearer.is_empty() {
return false;
}
if !self.agent.is_empty() && same(bearer, &self.agent) {
return true;
}
self.previous.as_ref().is_some_and(|(previous, until)| {
!previous.is_empty() && *until > std::time::Instant::now() && same(bearer, previous)
})
}
}
fn tracing() -> bool {
static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
*ON.get_or_init(|| std::env::var("SEMLITH_HTTP_TRACE").is_ok_and(|v| v == "1"))
}
fn trace(stream: &TcpStream, what: impl FnOnce() -> String) {
if tracing() {
let port = stream.peer_addr().map(|a| a.port()).unwrap_or(0);
eprintln!("semlith http: {port} {}", what());
}
}
const BUSY_READ: Duration = Duration::from_millis(50);
const BUSY_DRAIN: usize = 64 * 1024;
fn refuse_busy(mut stream: TcpStream) {
let started = std::time::Instant::now();
let _ = stream.set_read_timeout(Some(BUSY_READ));
let mut seen: Vec<u8> = Vec::new();
let mut want: Option<usize> = None;
let mut buf = [0u8; 8 * 1024];
let mut first = [0u8; 1];
if !matches!(stream.peek(&mut first), Ok(n) if n > 0) {
trace(&stream, || "over the limit, closed: no request".to_string());
return;
}
trace(&stream, || "over the limit, answered 503".to_string());
while seen.len() < BUSY_DRAIN && started.elapsed() < BUSY_READ {
if want.is_some_and(|total| seen.len() >= total) {
break;
}
match stream.read(&mut buf) {
Ok(0) | Err(_) => break,
Ok(n) => seen.extend_from_slice(&buf[..n]),
}
if want.is_none()
&& let Some(end) = seen.windows(4).position(|w| w == b"\r\n\r\n")
{
let head = String::from_utf8_lossy(&seen[..end]).to_ascii_lowercase();
let body = head
.lines()
.find_map(|l| l.strip_prefix("content-length:"))
.and_then(|v| v.trim().parse::<usize>().ok())
.unwrap_or(0);
want = Some(end + 4 + body);
}
}
let _ = write_response(&mut stream, Response::new(503, "text/plain", Vec::new()));
}
fn answer(
ready: Ready,
handler: &Handler,
auth: &Auth,
hold: &impl Fn(TcpStream, InFlight),
) -> Option<Refusal> {
let Ready {
mut stream,
mut reader,
mut request,
in_flight,
accepted: started,
} = ready;
if !loopback_host(request.header("host")) {
drain_body(&mut reader, &request, started);
let _ = write_response(&mut stream, Response::new(400, "text/plain", Vec::new()));
return Some(Refusal::ForeignHost);
}
let want = auth.token.as_str();
let for_mcp = request.path == MCP_PATH;
if for_mcp && !auth.mcp_open {
drain_body(&mut reader, &request, started);
let _ = write_response(&mut stream, Response::error(404, "no such route"));
return None;
}
let head_only = request.method == "HEAD";
let reading = head_only || request.method == "GET";
if !reading {
if !same_origin(&request, auth.port) {
drain_body(&mut reader, &request, started);
let _ = write_response(&mut stream, Response::new(403, "text/plain", Vec::new()));
return Some(Refusal::CrossOrigin);
}
if !json_body(request.header("content-type")) {
drain_body(&mut reader, &request, started);
let _ = write_response(&mut stream, Response::new(403, "text/plain", Vec::new()));
return Some(Refusal::BadContentType);
}
}
if !public(&request, reading) {
let by_agent = for_mcp
&& request
.header("authorization")
.and_then(|value| value.strip_prefix("Bearer "))
.is_some_and(|bearer| auth.is_agent(bearer.trim()));
let by_header = request
.header(TOKEN_HEADER_KEY)
.is_some_and(|t| same(t.trim(), want));
let agent_elsewhere = !for_mcp
&& !by_header
&& request
.header("authorization")
.and_then(|value| value.strip_prefix("Bearer "))
.is_some_and(|bearer| auth.is_agent(bearer.trim()));
if agent_elsewhere {
drain_body(&mut reader, &request, started);
let _ = write_response(&mut stream, Response::new(403, "text/plain", Vec::new()));
return Some(Refusal::Unauthorized);
}
if !by_header && !by_agent {
drain_body(&mut reader, &request, started);
hold(stream, in_flight);
return Some(Refusal::Unauthorized);
}
}
if read_body(&mut reader, &mut request).is_err() {
let _ = write_response(&mut stream, Response::new(400, "text/plain", Vec::new()));
return Some(Refusal::Malformed);
}
if started.elapsed() > REQUEST_DEADLINE {
let _ = write_response(&mut stream, Response::new(408, "text/plain", Vec::new()));
return Some(Refusal::Malformed);
}
let request = if head_only {
Request {
method: "GET".to_string(),
..request
}
} else {
request
};
let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| handler(&request)));
let mut response = match caught {
Ok(response) => response,
Err(_) => {
eprintln!(
"semlith: {} {} panicked; answered 500 and kept serving",
request.method, request.path
);
Response::new(500, "text/plain", Vec::new())
}
};
if head_only {
response.body = Body::Bytes(Vec::new());
}
if matches!(response.body, Body::Stream(_)) {
std::thread::spawn(move || {
let _in_flight = in_flight;
let _ = write_response(&mut stream, response);
});
return None;
}
let status = response.status;
let wrote = write_response(&mut stream, response);
trace(&stream, || {
format!(
"{} {} -> {status} in {} ms, write {:?}",
request.method,
request.path,
started.elapsed().as_millis(),
wrote.as_ref().map(|_| "ok")
)
});
None
}
fn public(request: &Request, reading: bool) -> bool {
if !reading {
return false;
}
if request.path == "/" || crate::portal::asset(&request.path).is_some() {
return true;
}
!request.path.starts_with("/api") && request.path != MCP_PATH
}
fn same_origin(request: &Request, port: u16) -> bool {
if let Some(site) = request.header("sec-fetch-site") {
return site.trim().eq_ignore_ascii_case("same-origin");
}
match request.header("origin") {
Some(origin) => own_origin(origin.trim(), port),
None => true,
}
}
fn own_origin(origin: &str, port: u16) -> bool {
let Some(authority) = origin.strip_prefix("http://") else {
return false;
};
let Some((host, stated)) = authority.rsplit_once(':') else {
return false;
};
let host = host.trim_start_matches('[').trim_end_matches(']');
let loopback = host.eq_ignore_ascii_case("localhost") || host == "127.0.0.1" || host == "::1";
loopback && stated.parse::<u16>() == Ok(port)
}
fn loopback_host(host: Option<&str>) -> bool {
let Some(host) = host else { return false };
let name = host.rsplit_once(':').map_or(host, |(n, _)| n);
let name = name.trim_start_matches('[').trim_end_matches(']');
name.eq_ignore_ascii_case("localhost") || name == "127.0.0.1" || name == "::1"
}
fn same(a: &str, b: &str) -> bool {
if a.len() != b.len() {
return false;
}
let mut diff = 0u8;
for (x, y) in a.bytes().zip(b.bytes()) {
diff |= x ^ y;
}
diff == 0
}
#[doc(hidden)]
pub fn fuzz_head(bytes: &[u8]) {
use std::io::Write;
static LISTENER: std::sync::OnceLock<std::net::TcpListener> = std::sync::OnceLock::new();
let listener = LISTENER
.get_or_init(|| std::net::TcpListener::bind("127.0.0.1:0").expect("a loopback port"));
let Ok(mut peer) = TcpStream::connect(listener.local_addr().expect("a bound port")) else {
return;
};
let _ = peer.write_all(bytes);
let _ = peer.shutdown(std::net::Shutdown::Write);
let Ok((accepted, _)) = listener.accept() else {
return;
};
let state = HeadState {
heard: AtomicBool::new(false),
evicted: AtomicBool::new(false),
};
let _ = read_head(
&mut BufReader::new(accepted),
std::time::Instant::now(),
&state,
);
}
fn read_head(
reader: &mut BufReader<TcpStream>,
accepted: std::time::Instant,
state: &HeadState,
) -> Result<Option<Request>> {
let mut head = String::new();
let mut line: Vec<u8> = Vec::new();
let mut read = 0;
loop {
let nothing_yet = read == 0;
let left = REQUEST_DEADLINE.saturating_sub(accepted.elapsed());
if left.is_zero() || state.evicted.load(Ordering::Relaxed) {
if nothing_yet {
return Ok(None);
}
return Err(HeadRefused(408).into());
}
reader
.get_ref()
.set_read_timeout(Some(left.min(HEAD_POLL)))?;
let buf = match reader.fill_buf() {
Ok(buf) => buf,
Err(e)
if matches!(
e.kind(),
std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
) =>
{
continue;
}
Err(_) if nothing_yet => return Ok(None),
Err(e) => return Err(e.into()),
};
if buf.is_empty() {
anyhow::ensure!(nothing_yet, "the connection closed mid-head");
return Ok(None);
}
state.heard.store(true, Ordering::Relaxed);
let (taken, complete) = match buf.iter().position(|b| *b == b'\n') {
Some(at) => (at + 1, true),
None => (buf.len(), false),
};
read += taken;
if read > MAX_HEADERS {
return Err(HeadRefused(431).into());
}
line.extend_from_slice(&buf[..taken]);
reader.consume(taken);
if !complete {
continue;
}
let text = String::from_utf8(std::mem::take(&mut line)).context("head is not UTF-8")?;
if text == "\r\n" || text == "\n" {
break;
}
head.push_str(&text);
}
let mut lines = head.lines();
let start = lines.next().context("empty request")?;
let mut parts = start.split_whitespace();
let method = parts.next().context("no method")?.to_string();
let target = parts.next().context("no target")?.to_string();
let mut headers = BTreeMap::new();
for line in lines {
if let Some((name, value)) = line.split_once(':') {
let name = name.trim().to_ascii_lowercase();
if SINGLE.contains(&name.as_str()) && headers.contains_key(&name) {
return Err(HeadRefused(400).into());
}
headers.insert(name, value.trim().to_string());
}
}
let (raw_path, raw_query) = match target.split_once('?') {
Some((p, q)) => (p, q),
None => (target.as_str(), ""),
};
Ok(Some(Request {
method,
path: percent_decode(raw_path),
query: parse_query(raw_query),
headers,
body: Vec::new(),
}))
}
fn drain_body(reader: &mut BufReader<TcpStream>, request: &Request, started: std::time::Instant) {
let Some(len) = request
.headers
.get("content-length")
.and_then(|v| v.parse::<usize>().ok())
else {
return;
};
let mut left = len.min(MAX_BODY);
let mut sink = [0u8; 8 * 1024];
while left > 0 && started.elapsed() <= REQUEST_DEADLINE {
let want = left.min(sink.len());
match reader.read(&mut sink[..want]) {
Ok(0) | Err(_) => return,
Ok(n) => left -= n,
}
}
}
fn read_body(reader: &mut BufReader<TcpStream>, request: &mut Request) -> Result<()> {
let Some(len) = request
.headers
.get("content-length")
.and_then(|v| v.parse::<usize>().ok())
else {
return Ok(());
};
anyhow::ensure!(len <= MAX_BODY, "request body too large");
request.body.resize(len, 0);
reader.read_exact(&mut request.body)?;
Ok(())
}
fn parse_query(raw: &str) -> BTreeMap<String, String> {
let mut out: BTreeMap<String, String> = BTreeMap::new();
for pair in raw.split('&').filter(|p| !p.is_empty()) {
let (key, value) = pair.split_once('=').unwrap_or((pair, ""));
let key = percent_decode(&key.replace('+', " "));
let value = percent_decode(&value.replace('+', " "));
out.entry(key)
.and_modify(|existing| {
existing.push('\u{1}');
existing.push_str(&value);
})
.or_insert(value);
}
out
}
fn percent_decode(raw: &str) -> String {
let bytes = raw.as_bytes();
let mut out = Vec::with_capacity(bytes.len());
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'%' && i + 2 < bytes.len() {
let hex = std::str::from_utf8(&bytes[i + 1..i + 3]).unwrap_or("");
if let Ok(byte) = u8::from_str_radix(hex, 16) {
out.push(byte);
i += 3;
continue;
}
}
out.push(bytes[i]);
i += 1;
}
String::from_utf8_lossy(&out).into_owned()
}
fn write_response(stream: &mut TcpStream, response: Response) -> std::io::Result<()> {
let written = write_whole(stream, response);
linger(stream);
written
}
const LINGER: Duration = Duration::from_millis(500);
fn linger(stream: &mut TcpStream) {
if stream.shutdown(std::net::Shutdown::Write).is_err() {
return;
}
let started = std::time::Instant::now();
let _ = stream.set_read_timeout(Some(LINGER));
let mut sink = [0u8; 4096];
while started.elapsed() < LINGER {
match stream.read(&mut sink) {
Ok(0) | Err(_) => return,
Ok(_) => {}
}
}
}
fn write_whole(stream: &mut TcpStream, response: Response) -> std::io::Result<()> {
let Response {
status,
content_type,
extra,
csp,
body,
} = response;
let mut head = format!("HTTP/1.1 {status} {}\r\n", reason(status));
head.push_str(&format!("Content-Type: {content_type}\r\n"));
head.push_str(&format!(
"Content-Security-Policy: {}\r\n",
csp.unwrap_or(PAGE_CSP)
));
head.push_str("X-Content-Type-Options: nosniff\r\n");
head.push_str("Referrer-Policy: no-referrer\r\n");
head.push_str("Permissions-Policy: camera=(), microphone=(), geolocation=()\r\n");
head.push_str("Cross-Origin-Opener-Policy: same-origin\r\n");
head.push_str("Cross-Origin-Resource-Policy: same-origin\r\n");
head.push_str("Cache-Control: no-store\r\n");
head.push_str("Connection: close\r\n");
for (name, value) in &extra {
head.push_str(&format!("{name}: {value}\r\n"));
}
match body {
Body::Bytes(bytes) => {
head.push_str(&format!("Content-Length: {}\r\n\r\n", bytes.len()));
stream.write_all(head.as_bytes())?;
stream.write_all(&bytes)?;
stream.flush()
}
Body::Stream(f) => {
head.push_str("Transfer-Encoding: chunked\r\n\r\n");
stream.write_all(head.as_bytes())?;
stream.flush()?;
let result = {
let mut chunks = Chunks { out: stream };
f(&mut chunks)
};
stream.write_all(b"0\r\n\r\n")?;
stream.flush()?;
result
}
}
}
const SINGLE: &[&str] = &[
"host",
"origin",
"sec-fetch-site",
"content-type",
"content-length",
"semlith-token",
"authorization",
];
fn reason(status: u16) -> &'static str {
match status {
200 => "OK",
400 => "Bad Request",
401 => "Unauthorized",
403 => "Forbidden",
404 => "Not Found",
405 => "Method Not Allowed",
408 => "Request Timeout",
409 => "Conflict",
431 => "Request Header Fields Too Large",
413 => "Content Too Large",
500 => "Internal Server Error",
502 => "Bad Gateway",
503 => "Service Unavailable",
_ => "Unknown",
}
}
fn new_token() -> String {
let mut bytes = [0u8; 32];
let filled = getrandom::fill(&mut bytes).map(|()| bytes);
token_from(filled)
}
fn token_from(filled: Result<[u8; 32], getrandom::Error>) -> String {
let bytes = filled
.expect("the OS random source refused; semlith will not mint a session token without it");
let mut out = String::with_capacity(64);
for byte in bytes {
use std::fmt::Write;
let _ = write!(out, "{byte:02x}");
}
out
}
#[derive(Default)]
struct Penalty {
recent: std::collections::VecDeque<std::time::Instant>,
}
const REFUSAL_DELAY: Duration = Duration::from_millis(250);
const MAX_REFUSAL_DELAY: Duration = Duration::from_secs(4);
const REFUSALS_PER_STEP: usize = 20;
impl Penalty {
fn charge(&mut self) -> Duration {
let now = std::time::Instant::now();
while self
.recent
.front()
.is_some_and(|t| now.duration_since(*t) > Duration::from_secs(60))
{
self.recent.pop_front();
}
self.recent.push_back(now);
let steps = (self.recent.len().saturating_sub(1)) / REFUSALS_PER_STEP;
let delay = REFUSAL_DELAY * 2u32.saturating_pow(steps.min(8) as u32);
delay.min(MAX_REFUSAL_DELAY)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn only_loopback_hosts_are_answered() {
assert!(loopback_host(Some("localhost")));
assert!(loopback_host(Some("localhost:7365")));
assert!(loopback_host(Some("127.0.0.1:7365")));
assert!(loopback_host(Some("LOCALHOST:7365")));
assert!(loopback_host(Some("[::1]:7365")));
assert!(!loopback_host(Some("evil.example")));
assert!(!loopback_host(Some("evil.example:7365")));
assert!(!loopback_host(Some("notlocalhost")));
assert!(!loopback_host(Some("localhost.evil.example")));
assert!(!loopback_host(Some("127.1")));
assert!(!loopback_host(None));
}
fn write(headers: &[(&str, &str)]) -> Request {
Request {
method: "POST".into(),
path: "/api/index".into(),
query: BTreeMap::new(),
headers: headers
.iter()
.map(|(k, v)| (k.to_ascii_lowercase(), v.to_string()))
.collect(),
body: Vec::new(),
}
}
#[test]
fn only_a_same_origin_write_is_accepted() {
assert!(same_origin(
&write(&[("Sec-Fetch-Site", "same-origin")]),
7365
));
assert!(!same_origin(
&write(&[("Sec-Fetch-Site", "same-site")]),
7365
));
assert!(!same_origin(
&write(&[("Sec-Fetch-Site", "cross-site")]),
7365
));
assert!(!same_origin(&write(&[("Sec-Fetch-Site", "none")]), 7365));
assert!(!same_origin(
&write(&[
("Sec-Fetch-Site", "same-site"),
("Origin", "http://127.0.0.1:7365"),
]),
7365
));
assert!(same_origin(&write(&[]), 7365));
assert!(same_origin(
&write(&[("Origin", "http://localhost:7365")]),
7365
));
assert!(!same_origin(
&write(&[("Origin", "http://127.0.0.1:9999")]),
7365
));
assert!(!same_origin(&write(&[("Origin", "null")]), 7365));
assert!(!same_origin(
&write(&[("Origin", "https://127.0.0.1:7365")]),
7365
));
}
#[test]
fn only_a_json_body_is_read() {
assert!(json_body(Some("application/json")));
assert!(json_body(Some("application/json; charset=utf-8")));
assert!(json_body(Some("Application/JSON")));
assert!(!json_body(Some("text/plain")));
assert!(!json_body(Some("application/x-www-form-urlencoded")));
assert!(!json_body(Some("multipart/form-data; boundary=x")));
assert!(!json_body(Some("application/json-patch+json")));
assert!(!json_body(None));
}
#[test]
fn only_the_page_and_its_assets_are_public() {
let get = |path: &str| Request {
method: "GET".into(),
path: path.into(),
query: BTreeMap::new(),
headers: BTreeMap::new(),
body: Vec::new(),
};
assert!(public(&get("/"), true));
assert!(public(&get("/app.js"), true));
assert!(public(&get("/style.css"), true));
assert!(public(&get("/nope.png"), true));
assert!(public(&get("/icons/nope.png"), true));
assert!(!public(&get("/api/stores"), true));
assert!(!public(&get("/api/nope"), true));
assert!(!public(&get("/api/image"), true));
assert!(!public(&get(MCP_PATH), true));
assert!(!public(&get("/"), false));
}
#[test]
fn tokens_compare_by_value_and_by_length() {
assert!(same("abc", "abc"));
assert!(!same("abc", "abd"));
assert!(!same("abc", "abcd"));
assert!(!same("", "a"));
}
#[test]
fn a_repeated_query_key_keeps_every_value() {
let q = parse_query("lang=rust&lang=go&k=5");
let request = Request {
method: "GET".into(),
path: "/".into(),
query: q,
headers: BTreeMap::new(),
body: Vec::new(),
};
assert_eq!(request.query_all("lang"), vec!["rust", "go"]);
assert_eq!(request.query("k"), Some("5"));
assert!(request.query_all("ext").is_empty());
}
#[test]
fn percent_and_plus_both_decode() {
assert_eq!(percent_decode("a%20b"), "a b");
assert_eq!(percent_decode("%2Fsrc%2Flib.rs"), "/src/lib.rs");
assert_eq!(percent_decode("100%"), "100%");
let q = parse_query("query=Fleet%3A%3Awritable+opens");
assert_eq!(q.get("query").unwrap(), "Fleet::writable opens");
}
#[test]
fn every_token_is_new() {
let a = new_token();
let b = new_token();
assert_ne!(a, b);
assert_eq!(a.len(), 64, "256 bits, hex encoded");
assert!(
a.bytes()
.all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()),
"not lowercase hex: {a}"
);
}
#[test]
fn a_token_is_never_minted_without_the_os_random_source() {
let refused = std::panic::catch_unwind(|| token_from(Err(getrandom::Error::UNSUPPORTED)))
.expect_err("a failing random source minted a token anyway");
let message = refused
.downcast_ref::<String>()
.map(String::as_str)
.or_else(|| refused.downcast_ref::<&str>().copied())
.unwrap_or("");
assert!(
message.contains("the OS random source refused"),
"the panic does not say what refused: {message}"
);
assert_eq!(token_from(Ok([0u8; 32])), "0".repeat(64));
assert_eq!(&token_from(Ok([0xab; 32]))[..4], "abab");
}
#[test]
fn a_connection_past_the_limit_is_answered_not_reset() {
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
let port = listener.local_addr().unwrap().port();
let client = std::thread::spawn(move || {
let mut c = TcpStream::connect((Ipv4Addr::LOCALHOST, port)).unwrap();
let body = vec![b'x'; 8 * 1024];
let head = format!(
"POST /api/x HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Length: {}\r\n\r\n",
body.len()
);
c.write_all(head.as_bytes()).unwrap();
c.write_all(&body).unwrap();
c.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
let mut answer = Vec::new();
c.read_to_end(&mut answer).map(|_| answer)
});
let (stream, _) = listener.accept().unwrap();
std::thread::sleep(Duration::from_millis(100));
refuse_busy(stream);
let answer = client.join().unwrap().expect("the client saw a reset");
assert!(
String::from_utf8_lossy(&answer).starts_with("HTTP/1.1 503"),
"{}",
String::from_utf8_lossy(&answer)
);
}
#[test]
fn a_run_of_refusals_gets_slower() {
let mut penalty = Penalty::default();
assert_eq!(penalty.charge(), REFUSAL_DELAY, "the first refusal");
for _ in 1..REFUSALS_PER_STEP {
assert_eq!(penalty.charge(), REFUSAL_DELAY);
}
assert_eq!(penalty.charge(), REFUSAL_DELAY * 2);
for _ in 0..REFUSALS_PER_STEP {
penalty.charge();
}
assert_eq!(penalty.charge(), REFUSAL_DELAY * 4);
for _ in 0..200 {
penalty.charge();
}
assert_eq!(
penalty.charge(),
MAX_REFUSAL_DELAY,
"the delay is capped, because holding a connection forever is the \
other kind of denial of service"
);
let old = std::time::Instant::now() - Duration::from_secs(61);
penalty.recent = std::iter::repeat_n(old, 500).collect();
assert_eq!(penalty.charge(), REFUSAL_DELAY);
}
}