pub mod file; use std::collections::HashMap; use std::future::Future; use std::io; use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, Weak}; use std::task::{Context, Poll}; use base64::Engine as _; use bytes::Bytes; use futures::channel::oneshot; use futures::future::BoxFuture; use http_body_util::{BodyExt, Full, StreamBody, combinators::UnsyncBoxBody}; use hyper::body::{Body as HttpBody, Frame, Incoming, SizeHint}; use hyper::server::conn::http1; use hyper::service::service_fn; use hyper::{Request, Response, StatusCode}; use hyper_util::rt::{TokioIo, TokioTimer}; use sha1::Digest; use tokio::sync::{Mutex as AsyncMutex, mpsc}; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; use tokio_util::task::task_tracker::TaskTrackerToken; pub use avril_net::Address; use avril_net::{Endpoint, Listener}; pub type Error = Box; type Body = UnsyncBoxBody; type Requests = Arc>>>; type Upload = mpsc::Sender, Error>>; pub type Authorize = Arc Result<(), Error> + Send + Sync>; pub type Metadata = (u64, String, String, Vec<(String, String)>, bool, Endpoint); pub type RespondArgs = (u64, u16, Vec<(String, String)>, Option, bool); const MAX_ID: u64 = 9_007_199_254_740_991; enum BodyKind { Inline(Bytes), GuestStream, File(Arc), } struct Exchange { server: Weak, method: hyper::Method, version: hyper::Version, headers: hyper::HeaderMap, upgrade: Mutex>, pending_upgrade: Mutex>, body: AsyncMutex, reply: Mutex>>>, upload: Mutex>, done: CancellationToken, disconnected: AtomicBool, } impl Exchange { fn check(&self) -> Result<(), Error> { self.server .upgrade() .ok_or("HTTP server is closed")? .check() } } struct Lifetime { id: u64, exchange: Arc, requests: Requests, complete: bool, } impl Drop for Lifetime { fn drop(&mut self) { self.exchange .disconnected .store(!self.complete, Ordering::Relaxed); self.exchange.done.cancel(); self.exchange.upload.lock().unwrap().take(); self.exchange.pending_upgrade.lock().unwrap().take(); self.requests.lock().unwrap().remove(&self.id); } } struct PendingUpgrade { headers: hyper::HeaderMap, ready: oneshot::Sender, } pub struct Upgrade { on_upgrade: hyper::upgrade::OnUpgrade, ready: oneshot::Receiver, pub canceled: CancellationToken, token: TaskTrackerToken, } pub struct Lease { _lifetime: Lifetime, _token: TaskTrackerToken, } impl Upgrade { pub async fn accept(self) -> Result<(hyper::upgrade::Upgraded, Lease), Error> { let lifetime = tokio::select! { _ = self.canceled.cancelled() => return Err("HTTP server aborted".into()), ready = self.ready => ready.map_err(|_| "WebSocket upgrade response was not returned")?, }; let stream = tokio::select! { _ = self.canceled.cancelled() => return Err("HTTP server aborted".into()), stream = self.on_upgrade => stream?, }; Ok(( stream, Lease { _lifetime: lifetime, _token: self.token, }, )) } } struct ResponseBody { inner: Body, lifetime: Lifetime, remaining: Option, failed: bool, } impl HttpBody for ResponseBody { type Data = Bytes; type Error = Error; fn poll_frame( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll, Error>>> { let result = Pin::new(&mut self.inner).poll_frame(cx); if let Poll::Ready(Some(frame)) = &result { match frame { Ok(frame) => { if let (Some(remaining), Some(bytes)) = (self.remaining, frame.data_ref()) { self.remaining = remaining.checked_sub(bytes.len() as u64); if self.remaining.is_none() { self.failed = true; } } } Err(_) => self.failed = true, } } if (matches!(result, Poll::Ready(None)) || self.inner.is_end_stream()) && self.remaining.is_some_and(|remaining| remaining != 0) { self.failed = true; } if !self.failed && (self.remaining == Some(0) || matches!(result, Poll::Ready(None)) || self.inner.is_end_stream()) { self.lifetime.complete = true; } result } fn is_end_stream(&self) -> bool { self.inner.is_end_stream() } fn size_hint(&self) -> SizeHint { self.inner.size_hint() } } impl Drop for ResponseBody { fn drop(&mut self) { // Hyper can stop polling at Content-Length without ever observing EOF. if !self.failed && (self.remaining == Some(0) || self.inner.is_end_stream()) { self.lifetime.complete = true; } } } struct FileResponseBody { file: Arc, request: Arc, pending: Option>>, remaining: u64, ended: bool, } impl HttpBody for FileResponseBody { type Data = Bytes; type Error = Error; fn poll_frame( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll, Error>>> { if self.ended { return Poll::Ready(None); } if self.pending.is_none() { let file = self.file.clone(); let request = self.request.clone(); self.pending = Some(Box::pin(async move { request.check()?; let result = tokio::select! { biased; _ = request.done.cancelled() => return Err("HTTP request canceled".into()), result = file.read() => result?, }; request.check()?; Ok((Bytes::from(result.0), result.1)) })); } match self.pending.as_mut().unwrap().as_mut().poll(cx) { Poll::Pending => Poll::Pending, Poll::Ready(result) => { self.pending = None; match result { Ok((bytes, done)) => { self.remaining -= bytes.len() as u64; self.ended = done; Poll::Ready(Some(Ok(Frame::data(bytes)))) } Err(error) => { self.ended = true; self.file.close(); Poll::Ready(Some(Err(error))) } } } } } fn is_end_stream(&self) -> bool { self.ended } fn size_hint(&self) -> SizeHint { SizeHint::with_exact(self.remaining) } } impl Drop for FileResponseBody { fn drop(&mut self) { self.file.close(); } } struct Server { addr: Address, authorize: Authorize, upgrades: TaskTracker, queue: mpsc::Sender, incoming: AsyncMutex>, stop: CancellationToken, abort: CancellationToken, done: CancellationToken, error: Mutex>, requests: Requests, next: Arc, } impl Server { fn check(&self) -> Result<(), Error> { (self.authorize)(&self.addr) } async fn request( self: Arc, mut request: Request, remote: Endpoint, ) -> Result, Error> { self.check()?; if request.method() == hyper::Method::CONNECT { return Ok(Response::builder() .status(StatusCode::NOT_IMPLEMENTED) .body( Full::new(Bytes::new()) .map_err(|never| match never {}) .boxed_unsync(), )?); } let authority = request .headers() .get("host") .map(|value| value.to_str()) .transpose()? .map(str::to_owned) .unwrap_or_else(|| match &self.addr { Address::Tcp { hostname, port } => { if hostname.contains(':') { format!("[{hostname}]:{port}") } else { format!("{hostname}:{port}") } } Address::Unix { .. } => "localhost".to_owned(), }); authority.parse::()?; let path = request .uri() .path_and_query() .map_or("/", |path| path.as_str()); let path = if path == "*" { "/" } else { path }; let url = url::Url::parse(&format!("http://{authority}{path}"))?; if !url.username().is_empty() || url.password().is_some() || url.host_str().is_none() { return Err("Invalid HTTP request authority".into()); } let id = self .next .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| { (id <= MAX_ID).then_some(id + 1) }) .map_err(|_| "HTTP request IDs exhausted")?; let headers = request .headers() .iter() .map(|(name, value)| { ( name.to_string(), value .as_bytes() .iter() .map(|byte| char::from(*byte)) .collect(), ) }) .collect(); let has_body = !request.body().is_end_stream(); let upgrade = hyper::upgrade::on(&mut request); let original_headers = request.headers().clone(); let version = request.version(); let method = request.method().clone(); let (reply, receive) = oneshot::channel(); let exchange = Arc::new(Exchange { server: Arc::downgrade(&self), method: method.clone(), version, headers: original_headers, upgrade: Mutex::new(Some(upgrade)), pending_upgrade: Mutex::new(None), body: AsyncMutex::new(request.into_body()), reply: Mutex::new(Some(reply)), upload: Mutex::new(None), done: CancellationToken::new(), disconnected: AtomicBool::new(false), }); self.requests.lock().unwrap().insert(id, exchange.clone()); let lifetime = Lifetime { id, exchange: exchange.clone(), requests: self.requests.clone(), complete: false, }; tokio::select! { biased; _ = self.stop.cancelled() => return Err("HTTP server is shutting down".into()), result = self.queue.send((id, method.to_string(), url.into(), headers, has_body, remote)) => { result.map_err(|_| "HTTP dispatcher stopped")?; } } let response = tokio::select! { _ = exchange.done.cancelled() => return Err("HTTP request canceled".into()), response = receive => response.map_err(|_| "HTTP handler stopped")?, }; if response.status() == StatusCode::SWITCHING_PROTOCOLS { let pending = exchange .pending_upgrade .lock() .unwrap() .take() .ok_or("WebSocket upgrade was not prepared")?; pending .ready .send(lifetime) .map_err(|_| "WebSocket upgrade was canceled")?; return Ok(response); } exchange.pending_upgrade.lock().unwrap().take(); let suppressed = exchange.method == hyper::Method::HEAD || matches!(response.status().as_u16(), 204 | 205 | 304); let remaining = if suppressed { Some(0) } else { response .headers() .get("content-length") .and_then(|value| value.to_str().ok()?.parse().ok()) }; let (parts, inner) = response.into_parts(); Ok(Response::from_parts( parts, ResponseBody { inner, lifetime, remaining, failed: false, } .boxed_unsync(), )) } async fn serve(self: Arc, listener: Listener) -> Result<(), Error> { let mut connections = JoinSet::new(); loop { tokio::select! { biased; _ = self.stop.cancelled() => break, Some(_) = connections.join_next(), if !connections.is_empty() => {}, accepted = listener.accept() => { let avril_net::Accepted { socket, remote, .. } = match accepted { Ok(accepted) => accepted, Err(error) => { if matches!(error.kind(), io::ErrorKind::ConnectionAborted | io::ErrorKind::ConnectionReset | io::ErrorKind::Interrupted) { continue; } let exhausted = error.kind() == io::ErrorKind::OutOfMemory; #[cfg(unix)] let exhausted = exhausted || matches!(error.raw_os_error(), Some(libc::EMFILE | libc::ENFILE | libc::ENOBUFS | libc::ENOMEM)); #[cfg(windows)] let exhausted = exhausted || matches!(error.raw_os_error(), Some(10024 | 10055)); if !exhausted { return Err(error.into()); } tokio::select! { _ = self.stop.cancelled() => break, _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {}, } continue; } }; self.check()?; let server = self.clone(); connections.spawn(async move { let service_server = server.clone(); let connection = http1::Builder::new().timer(TokioTimer::new()).serve_connection(TokioIo::new(socket), service_fn(move |request| { service_server.clone().request(request, remote.clone()) })).with_upgrades(); tokio::pin!(connection); tokio::select! { biased; _ = server.abort.cancelled() => {}, _ = server.stop.cancelled() => { connection.as_mut().graceful_shutdown(); tokio::select! { _ = server.abort.cancelled() => {}, _ = &mut connection => {}, } }, _ = &mut connection => {}, } }); } } } drop(listener); while connections.join_next().await.is_some() {} self.upgrades.close(); self.upgrades.wait().await; Ok(()) } } pub struct Engine { runtime: tokio::runtime::Handle, servers: Mutex>>, requests: Requests, next_server: AtomicU64, next_request: Arc, } impl Drop for Engine { fn drop(&mut self) { for server in self.servers.get_mut().unwrap().values() { server.stop.cancel(); server.abort.cancel(); } } } impl Engine { fn server(&self, id: u64) -> Result, Error> { self.servers .lock() .unwrap() .get(&id) .cloned() .ok_or_else(|| "HTTP server is closed".into()) } fn request(&self, id: u64) -> Result, Error> { self.requests .lock() .unwrap() .get(&id) .cloned() .ok_or_else(|| "HTTP request is closed".into()) } fn run> + Send + 'static>( &self, future: F, ) -> impl Future> + use { let handle = self.runtime.spawn(future); async move { // Dropping the handle detaches the task rather than aborting it. handle .await .map_err(|_| io::Error::other("HTTP worker stopped"))? .map_err(io::Error::other) } } pub fn upgrade( &self, id: u64, protocol: Option<&str>, ) -> Result<(Upgrade, Vec<(String, String)>), Error> { let request = self.request(id)?; request.check()?; let token = |name, token: &str| { request.headers.get_all(name).iter().any(|value| { value.to_str().is_ok_and(|value| { value .split(',') .any(|value| value.trim().eq_ignore_ascii_case(token)) }) }) }; let single = |name| { let mut values = request.headers.get_all(name).iter(); let value = values .next() .ok_or("Missing WebSocket handshake header")? .to_str()?; if values.next().is_some() { return Err("Duplicate WebSocket handshake header".into()); } Ok::<_, Error>(value) }; if request.method != hyper::Method::GET || request.version != hyper::Version::HTTP_11 || !token("connection", "upgrade") || !token("upgrade", "websocket") || single("sec-websocket-version")? != "13" { return Err("Invalid WebSocket upgrade request".into()); } let key = single("sec-websocket-key")?; if base64::engine::general_purpose::STANDARD.decode(key)?.len() != 16 { return Err("Invalid WebSocket key".into()); } if let Some(protocol) = protocol && (protocol.is_empty() || !protocol .bytes() .all(|byte| byte.is_ascii_alphanumeric() || b"!#$%&'*+-.^_`|~".contains(&byte)) || !request .headers .get_all("sec-websocket-protocol") .iter() .any(|value| { value.to_str().is_ok_and(|value| { value.split(',').any(|value| value.trim() == protocol) }) })) { return Err("WebSocket protocol was not offered by the client".into()); } let mut hash = sha1::Sha1::new(); hash.update(key.as_bytes()); hash.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); let mut response = Response::builder() .status(101) .header("connection", "Upgrade") .header("upgrade", "websocket") .header( "sec-websocket-accept", base64::engine::general_purpose::STANDARD.encode(hash.finalize()), ); if let Some(protocol) = protocol { response = response.header("sec-websocket-protocol", protocol); } let response = response.body(())?; let headers = response .headers() .iter() .map(|(name, value)| (name.to_string(), value.to_str().unwrap().to_owned())) .collect(); let server = request.server.upgrade().ok_or("HTTP server is closed")?; let on_upgrade = request .upgrade .lock() .unwrap() .take() .ok_or("WebSocket upgrade already prepared")?; let (ready, receive) = oneshot::channel(); *request.pending_upgrade.lock().unwrap() = Some(PendingUpgrade { headers: response.into_parts().0.headers, ready, }); Ok(( Upgrade { on_upgrade, ready: receive, canceled: server.abort.clone(), token: server.upgrades.token(), }, headers, )) } pub fn next( &self, id: u64, ) -> impl Future, io::Error>> + use<> { let server = self.server(id); self.run(async move { let server = server?; tokio::select! { biased; _ = server.stop.cancelled() => Ok(server.incoming.lock().await.try_recv().ok()), item = async { server.incoming.lock().await.recv().await } => Ok(item), } }) } pub fn finished( self: &Arc, id: u64, ) -> impl Future> + use<> { let server = self.server(id); let pending = self.run(async move { let server = server?; server.done.cancelled().await; let error = server.error.lock().unwrap().clone(); if let Some(error) = error { Err(error.into()) } else { Ok(()) } }); let engine = self.clone(); async move { let result = pending.await; engine.servers.lock().unwrap().remove(&id); result } } pub fn stop_all(&self, force: bool) { for server in self.servers.lock().unwrap().values() { server.stop.cancel(); if force { server.abort.cancel(); } } } pub fn stop(&self, id: u64, force: bool) { if let Ok(server) = self.server(id) { server.stop.cancel(); if force { server.abort.cancel(); } } } pub fn read(&self, id: u64) -> impl Future, io::Error>> + use<> { let request = self.request(id); self.run(async move { let request = request?; request.check()?; tokio::select! { _ = request.done.cancelled() => Err("HTTP request canceled".into()), result = async { let mut body = request.body.lock().await; while let Some(frame) = body.frame().await { if let Ok(data) = frame?.into_data() { return Ok(Some(data)); } } Ok(None) } => result, } }) } pub fn write( &self, id: u64, bytes: Option, ) -> impl Future> + use<> { let request = self.request(id); self.run(async move { let Ok(request) = request else { return Ok(false); }; request.check()?; if let Some(bytes) = bytes { let Some(upload) = request.upload.lock().unwrap().clone() else { return Ok(false); }; tokio::select! { _ = request.done.cancelled() => Ok(false), result = upload.send(Ok(Frame::data(bytes))) => Ok(result.is_ok()), } } else { Ok(request.upload.lock().unwrap().take().is_some()) } }) } pub fn wait(&self, id: u64) -> impl Future> + use<> { let request = self.request(id); self.run(async move { let request = request?; request.done.cancelled().await; Ok(request.disconnected.load(Ordering::Relaxed)) }) } pub fn cancel(&self, id: u64) { if let Ok(request) = self.request(id) { request.disconnected.store(true, Ordering::Relaxed); request.done.cancel(); request.pending_upgrade.lock().unwrap().take(); } } pub fn new(runtime: tokio::runtime::Handle) -> Self { Self { runtime, servers: Mutex::new(HashMap::new()), requests: Arc::new(Mutex::new(HashMap::new())), next_server: AtomicU64::new(1), next_request: Arc::new(AtomicU64::new(1)), } } pub fn listen(&self, address: Address, authorize: Authorize) -> Result<(u64, Address), Error> { authorize(&address)?; let (listener, addr) = Listener::bind(&self.runtime, &address)?; let addr = Address::from(addr); authorize(&addr)?; let id = self .next_server .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| { (id <= MAX_ID).then_some(id + 1) }) .map_err(|_| "HTTP server IDs exhausted")?; let (queue, incoming) = mpsc::channel(64); let server = Arc::new(Server { addr: addr.clone(), authorize, upgrades: TaskTracker::new(), queue, incoming: AsyncMutex::new(incoming), stop: CancellationToken::new(), abort: CancellationToken::new(), done: CancellationToken::new(), error: Mutex::new(None), requests: self.requests.clone(), next: self.next_request.clone(), }); self.servers.lock().unwrap().insert(id, server.clone()); self.runtime.spawn(async move { if let Err(error) = server.clone().serve(listener).await { *server.error.lock().unwrap() = Some(error.to_string()); server.abort.cancel(); server.upgrades.close(); server.upgrades.wait().await; } server.stop.cancel(); server.done.cancel(); }); Ok((id, addr)) } pub fn respond(&self, args: RespondArgs) -> Result { let (id, status, headers, bytes, streaming) = args; let body = if streaming { BodyKind::GuestStream } else { BodyKind::Inline(bytes.unwrap_or_default()) }; self.respond_inner(id, status, headers, body) } pub fn respond_file( &self, id: u64, status: u16, headers: Vec<(String, String)>, file: Arc, ) -> Result { self.respond_inner(id, status, headers, BodyKind::File(file)) } fn respond_inner( &self, id: u64, status: u16, headers: Vec<(String, String)>, body: BodyKind, ) -> Result { let request = self.request(id)?; request.check()?; if status == 101 { let mut response = Response::new( Full::new(Bytes::new()) .map_err(|never| match never {}) .boxed_unsync(), ); *response.status_mut() = StatusCode::SWITCHING_PROTOCOLS; *response.headers_mut() = request .pending_upgrade .lock() .unwrap() .as_ref() .ok_or("WebSocket upgrade was not prepared")? .headers .clone(); request .reply .lock() .unwrap() .take() .ok_or("HTTP request already answered")? .send(response) .map_err(|_| "HTTP client disconnected")?; return Ok(false); } if !(200..=599).contains(&status) { return Err("Invalid HTTP response status".into()); } let mut response = Response::builder().status(status); for (name, value) in headers { if [ "connection", "transfer-encoding", "upgrade", "trailer", "keep-alive", "proxy-connection", ] .contains(&name.to_ascii_lowercase().as_str()) { return Err(format!("Unsupported HTTP response header: {name}").into()); } let value = value .chars() .map(|c| u8::try_from(c as u32)) .collect::, _>>()?; response = response.header(name, value); } let suppressed = request.method == hyper::Method::HEAD || matches!(status, 204 | 205 | 304); if !matches!(body, BodyKind::GuestStream) && !matches!(status, 204 | 304) && !response .headers_ref() .is_some_and(|headers| headers.contains_key("content-length")) { response = response.header( "content-length", if status == 205 { 0 } else { match &body { BodyKind::Inline(bytes) => bytes.len() as u64, BodyKind::File(file) => file.len(), BodyKind::GuestStream => 0, } }, ); } if let Some(length) = response .headers_ref() .and_then(|headers| headers.get("content-length")) { let length = length.to_str()?; if length.is_empty() || !length.bytes().all(|byte| byte.is_ascii_digit()) { return Err("Invalid response Content-Length".into()); } let length = length.parse::()?; let mismatch = match &body { BodyKind::Inline(bytes) => !suppressed && length != bytes.len() as u64, BodyKind::File(file) => !suppressed && length != file.len(), BodyKind::GuestStream => false, }; if status == 204 || (status == 205 && length != 0) || mismatch { return Err("Response Content-Length does not match the body".into()); } } let guest_streaming = matches!(body, BodyKind::GuestStream) && !suppressed; let body = match body { BodyKind::File(file) if !suppressed => FileResponseBody { remaining: file.len(), file, request: request.clone(), pending: None, ended: false, } .boxed_unsync(), BodyKind::GuestStream if !suppressed => { let (send, receive) = mpsc::channel(1); *request.upload.lock().unwrap() = Some(send); let done = request.done.clone(); let stream = futures::stream::unfold( (receive, done, false), |(mut receive, done, ended)| async move { if ended { return None; } tokio::select! { biased; _ = done.cancelled() => Some((Err("HTTP request canceled".into()), (receive, done, true))), frame = receive.recv() => frame.map(|frame| (frame, (receive, done, false))), } }, ); StreamBody::new(stream).boxed_unsync() } BodyKind::Inline(bytes) => Full::new(if suppressed { Bytes::new() } else { bytes }) .map_err(|never| match never {}) .boxed_unsync(), BodyKind::File(_) | BodyKind::GuestStream => Full::new(Bytes::new()) .map_err(|never| match never {}) .boxed_unsync(), }; let response = response.body(body)?; request .reply .lock() .unwrap() .take() .ok_or("HTTP request already answered")? .send(response) .map_err(|_| "HTTP client disconnected")?; Ok(guest_streaming) } }