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_COOKIE: &str = "semlith_token";
const WORKERS: usize = 8;
const MAX_HEADERS: usize = 16 * 1024;
const MAX_BODY: usize = 1024 * 1024;
const IO_TIMEOUT: Duration = Duration::from_secs(30);
static WAKE: AtomicBool = AtomicBool::new(false);
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Refusal {
Unauthorized,
ForeignHost,
Malformed,
}
impl Refusal {
pub fn as_str(self) -> &'static str {
match self {
Refusal::Unauthorized => "unauthorized",
Refusal::ForeignHost => "foreign host",
Refusal::Malformed => "malformed",
}
}
}
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> {
Ok(serde_json::from_slice(&self.body)?)
}
}
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)>,
body: Body,
}
impl Response {
pub fn new(status: u16, content_type: &'static str, body: Vec<u8>) -> Self {
Self {
status,
content_type,
extra: Vec::new(),
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 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(),
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_token(self, token: &str) -> Self {
self.header(
"Set-Cookie",
format!("{TOKEN_COOKIE}={token}; Path=/; HttpOnly; SameSite=Strict"),
)
}
}
pub type Handler = Arc<dyn Fn(&Request) -> Response + Send + Sync>;
pub struct Server {
listener: TcpListener,
port: u16,
token: Arc<Mutex<String>>,
}
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())),
})
}
pub fn port(&self) -> u16 {
self.port
}
pub fn token(&self) -> String {
self.token.lock().expect("the token lock").clone()
}
pub fn rotate(&self) -> String {
let fresh = new_token();
*self.token.lock().expect("the token lock") = 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::<TcpStream>();
let rx = Arc::new(Mutex::new(rx));
let refused = Arc::new(refused);
let mut workers = Vec::with_capacity(WORKERS);
for _ in 0..WORKERS {
let rx = Arc::clone(&rx);
let handler = Arc::clone(&handler);
let token = Arc::clone(&self.token);
let refused = Arc::clone(&refused);
workers.push(std::thread::spawn(move || {
loop {
let next = {
let guard = rx.lock().expect("the queue lock");
guard.recv()
};
let Ok(stream) = next else { return };
let want = token.lock().expect("the token lock").clone();
if let Some(class) = answer(stream, &handler, &want) {
refused(class);
}
}
}));
}
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));
if tx.send(stream).is_err() {
break;
}
}
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();
drop(tx);
for worker in workers {
let _ = worker.join();
}
Ok(())
}
}
fn answer(mut stream: TcpStream, handler: &Handler, want: &str) -> Option<Refusal> {
let request = match read_request(&mut stream) {
Ok(Some(r)) => r,
Ok(None) => return None,
Err(_) => {
let _ = write_response(&mut stream, Response::new(400, "text/plain", Vec::new()));
return Some(Refusal::Malformed);
}
};
if !loopback_host(request.header("host")) {
let _ = write_response(&mut stream, Response::new(400, "text/plain", Vec::new()));
return Some(Refusal::ForeignHost);
}
let from_cookie = cookie(request.header("cookie"), TOKEN_COOKIE);
let by_cookie = from_cookie.as_deref().is_some_and(|t| same(t, want));
let by_query = request.query("token").is_some_and(|t| same(t, want));
if !by_cookie && !by_query {
let _ = write_response(&mut stream, Response::new(401, "text/plain", Vec::new()));
return Some(Refusal::Unauthorized);
}
let head_only = request.method == "HEAD";
let request = if head_only {
Request {
method: "GET".to_string(),
..request
}
} else {
request
};
let mut response = handler(&request);
if by_query && !by_cookie {
response = response.with_token(want);
}
if head_only {
response.body = Body::Bytes(Vec::new());
}
let _ = write_response(&mut stream, response);
None
}
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 cookie(header: Option<&str>, name: &str) -> Option<String> {
header?.split(';').find_map(|pair| {
let (key, value) = pair.split_once('=')?;
(key.trim() == name).then(|| value.trim().to_string())
})
}
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
}
fn read_request(stream: &mut TcpStream) -> Result<Option<Request>> {
let mut reader = BufReader::new(stream.try_clone()?);
let mut head = String::new();
let mut read = 0;
loop {
let mut line = String::new();
let n = reader.read_line(&mut line)?;
if n == 0 {
return Ok(head.is_empty().then_some(None).flatten());
}
read += n;
anyhow::ensure!(read <= MAX_HEADERS, "request headers too large");
if line == "\r\n" || line == "\n" {
break;
}
head.push_str(&line);
}
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(':') {
headers.insert(name.trim().to_ascii_lowercase(), value.trim().to_string());
}
}
let (raw_path, raw_query) = match target.split_once('?') {
Some((p, q)) => (p, q),
None => (target.as_str(), ""),
};
let mut body = Vec::new();
if let Some(len) = headers.get("content-length").and_then(|v| v.parse().ok()) {
let len: usize = len;
anyhow::ensure!(len <= MAX_BODY, "request body too large");
body.resize(len, 0);
reader.read_exact(&mut body)?;
}
Ok(Some(Request {
method,
path: percent_decode(raw_path),
query: parse_query(raw_query),
headers,
body,
}))
}
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 Response {
status,
content_type,
extra,
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(
"Content-Security-Policy: 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'\r\n",
);
head.push_str("X-Content-Type-Options: nosniff\r\n");
head.push_str("Referrer-Policy: no-referrer\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
}
}
}
fn reason(status: u16) -> &'static str {
match status {
200 => "OK",
400 => "Bad Request",
401 => "Unauthorized",
404 => "Not Found",
405 => "Method Not Allowed",
409 => "Conflict",
500 => "Internal Server Error",
_ => "OK",
}
}
fn new_token() -> String {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let entropy = Box::new(0u8);
let address = &*entropy as *const u8 as usize;
let mut hasher = blake3::Hasher::new();
hasher.update(&now.to_le_bytes());
hasher.update(&std::process::id().to_le_bytes());
hasher.update(&address.to_le_bytes());
hasher.update(&std::time::Instant::now().elapsed().as_nanos().to_le_bytes());
hasher.finalize().to_hex().to_string()
}
#[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));
}
#[test]
fn a_cookie_is_read_out_of_a_header_with_others_in_it() {
let header = Some("theme=dark; semlith_token=abc123; other=1");
assert_eq!(cookie(header, TOKEN_COOKIE), Some("abc123".to_string()));
assert_eq!(cookie(Some("theme=dark"), TOKEN_COOKIE), None);
assert_eq!(cookie(None, TOKEN_COOKIE), None);
assert_eq!(cookie(Some("semlith_tokens=x"), TOKEN_COOKIE), None);
}
#[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");
}
}