mini_serve/

handler.rs

use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};

use hyper::body::{Bytes, Incoming};
use hyper::{Request, Response};
use http_body_util::combinators::BoxBody;
use http_body_util::{BodyExt, Full};

use crate::error::ServeError;
use crate::state::State;

/// An error that can occur while streaming a response body.
///
/// This error type wraps any error that might occur during the actual transmission
/// of the response body to the client (e.g., I/O errors from a file being streamed).
/// If the body stream yields an error after headers have already been sent to the client,
/// that error will cause the connection to be aborted — it cannot be converted back to
/// an HTTP error response.
#[derive(Debug)]
pub struct BodyError {
	inner: Box<dyn std::error::Error + Send + Sync>,
}

impl BodyError {
	/// Wrap an error in a `BodyError`.
	pub fn new(err: impl std::error::Error + Send + Sync + 'static) -> Self {
		BodyError {
			inner: Box::new(err),
		}
	}
}

impl fmt::Display for BodyError {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		write!(f, "body streaming error: {}", self.inner)
	}
}

impl std::error::Error for BodyError {
	fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
		Some(&*self.inner)
	}
}

impl From<Box<dyn std::error::Error + Send + Sync>> for BodyError {
	fn from(err: Box<dyn std::error::Error + Send + Sync>) -> Self {
		BodyError { inner: err }
	}
}

/// The response body type used by all handlers.
///
/// The error type is `BodyError`, which can represent errors that occur during body
/// streaming (e.g., disk I/O failures while reading a large file). These errors will
/// abort the connection rather than being converted to an HTTP error response.
pub type ResponseBody = BoxBody<Bytes, BodyError>;

/// Create a response body from raw bytes.
///
/// Since the bytes come from an infallible source (an in-memory `Full`), the body
/// stream can never actually fail. The error type is converted from `Infallible` to
/// `BodyError` to satisfy the `ResponseBody` type.
pub fn body(bytes: Bytes) -> ResponseBody {
	BoxBody::new(Full::new(bytes).map_err(|never: std::convert::Infallible| match never {}))
}

/// The upgraded connection handed to an [`OnUpgrade`] callback.
///
/// Wrapped so it implements tokio's `AsyncRead`/`AsyncWrite`, which is what a protocol
/// crate wants; the raw `hyper::upgrade::Upgraded` is reachable through it if needed.
pub type UpgradedIo = hyper_util::rt::TokioIo<hyper::upgrade::Upgraded>;

/// Take over a connection once the response has been written.
///
/// Attach one to a `101 Switching Protocols` response and the server hands you the raw
/// stream after the response goes out. Everything past that point speaks whatever protocol
/// you like — this crate stops interpreting the bytes.
///
/// The callback runs **inside the connection's own task**, which is deliberate and is the
/// reason this type exists rather than callers using `hyper::upgrade::on` directly. That
/// task holds the connection's semaphore permit and is the one shutdown aborts, so an
/// upgraded connection still counts against [`RouteBuilder::with_max_connections`] and is
/// still ended by the shutdown drain. Servicing the stream from a detached `tokio::spawn`
/// — the usual hyper pattern — escapes both.
///
/// Requires [`RouteBuilder::with_upgrades`]; without it the response is sent and the
/// callback never runs.
///
/// ```no_run
/// # use hyper::{Response, StatusCode};
/// # use mini_serve::{OnUpgrade, ResponseBody, ServeError};
/// # fn example() -> Result<Response<ResponseBody>, ServeError> {
/// let mut response = Response::builder()
///     .status(StatusCode::SWITCHING_PROTOCOLS)
///     .body(mini_serve::body(hyper::body::Bytes::new()))
///     .unwrap();
/// response.extensions_mut().insert(OnUpgrade::new(|_io| async move {
///     // speak your protocol here
/// }));
/// Ok(response)
/// # }
/// ```
///
/// The callback is held behind `Arc<Mutex<Option<..>>>` rather than directly, because
/// `http::Extensions` requires `Clone + Send + Sync` and a `FnOnce` is none of those. The
/// `Option` is what makes it callable once: the connection takes it, leaving `None`.
/// The boxed callback inside an [`OnUpgrade`]. Named because the nested type is otherwise
/// unreadable at every use site.
type UpgradeCallback =
	Box<dyn FnOnce(UpgradedIo) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send>;

#[derive(Clone)]
pub struct OnUpgrade(Arc<Mutex<Option<UpgradeCallback>>>);

impl OnUpgrade {
	/// Wrap a callback to run once the connection has been upgraded.
	pub fn new<F, Fut>(f: F) -> Self
	where
		F: FnOnce(UpgradedIo) -> Fut + Send + 'static,
		Fut: Future<Output = ()> + Send + 'static,
	{
		OnUpgrade(Arc::new(Mutex::new(Some(Box::new(move |io| {
			Box::pin(f(io)) as Pin<Box<dyn Future<Output = ()> + Send>>
		})))))
	}

	/// Run the callback, if it has not already been taken.
	///
	/// A poisoned lock or a second call are both no-ops rather than panics: failing to
	/// upgrade a connection is not worth taking a server down for.
	pub(crate) async fn run(self, io: UpgradedIo) {
		let taken = self.0.lock().ok().and_then(|mut slot| slot.take());
		if let Some(callback) = taken {
			callback(io).await;
		}
	}
}

impl fmt::Debug for OnUpgrade {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.write_str("OnUpgrade")
	}
}

/// A request handler that processes an HTTP request and returns a response or error.
///
/// Handlers receive the full request (method, path, headers, body) and the app state,
/// and return either a response or a `ServeError` (which is converted to an HTTP error response).
pub type Handler<S> = Arc<
	dyn Fn(Request<Incoming>, State<S>)
			-> Pin<Box<dyn Future<Output = Result<Response<ResponseBody>, ServeError>> + Send>>
		+ Send
		+ Sync,
>;

/// Wrap an async function to create a handler.
///
/// # Example
///
/// ```ignore
/// use mini_serve::handler;
/// use hyper::StatusCode;
/// use hyper::body::Bytes;
///
/// let h = handler(|req, state| async move {
///     Ok::<_, mini_serve::ServeError>(
///         mini_serve::json(StatusCode::OK, &serde_json::json!({"status": "ok"}))
///     )
/// });
/// ```
pub fn handler<S, F, Fut>(f: F) -> Handler<S>
where
	S: Send + Sync + 'static,
	F: Fn(Request<Incoming>, State<S>) -> Fut + Send + Sync + 'static,
	Fut: Future<Output = Result<Response<ResponseBody>, ServeError>> + Send + 'static,
{
	Arc::new(move |req, state| Box::pin(f(req, state)))
}

/// A middleware transforms a `Handler` into a new `Handler`, typically by
/// running logic before and/or after calling the inner handler — or by
/// short-circuiting and never calling it at all (e.g. to block a request).
///
/// Registered on a [`crate::RouteBuilder`] via `.wrap()`, and applied to every
/// route registered after that call.
pub type Middleware<S> = Arc<dyn Fn(Handler<S>) -> Handler<S> + Send + Sync>;