use crate::error::{Error, Result};
use crate::handler::{BoxFuture, DynHandler, Handler, HandlerWithParams};
use crate::middleware::{Layered, Middleware, Next};
use crate::openapi::{self, OpenApiInfo, Operation, RouteKind};
use crate::request::{Params, Request, StateMap};
use crate::response::Response;
use crate::router::Router;
use crate::shutdown::{Context, Shutdown};
use http_body_util::BodyExt;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper_util::rt::TokioIo;
use std::collections::HashMap;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;
pub struct App {
router: Router,
middlewares: Vec<Arc<dyn Middleware>>,
max_body_size: Option<usize>, docs: HashMap<(String, String), Operation>,
openapi: Option<OpenApiInfo>,
state: StateMap,
background: Vec<BackgroundTask>,
shutdown_grace: Duration,
body_limits: Vec<(String, Option<usize>)>,
#[cfg(feature = "tls")]
tls: Option<crate::tls::TlsConfig>,
#[cfg(feature = "tls")]
tls_reload: Option<Duration>,
}
struct ConnInfo {
peer: SocketAddr,
secure: bool,
client_cert: Option<Vec<u8>>,
}
type BackgroundTask = Box<dyn FnOnce(Context) -> BoxFuture<'static, ()> + Send>;
struct Shared {
router: Router,
middlewares: Vec<Arc<dyn Middleware>>,
max_body_size: Option<usize>,
body_limits: Vec<(String, Option<usize>)>,
state: Arc<StateMap>,
shutdown: Shutdown,
}
pub struct RouteGroup {
prefix: String,
routes: Vec<(String, String, DynHandler)>, middlewares: Vec<Arc<dyn Middleware>>,
docs: Vec<(String, String, Operation)>, body_limit: Option<Option<usize>>,
ws_paths: Vec<String>,
}
impl App {
pub fn new() -> Self {
Self {
router: Router::new(),
middlewares: Vec::new(),
max_body_size: Some(10 * 1024 * 1024), docs: HashMap::new(),
openapi: None,
state: StateMap::default(),
background: Vec::new(),
shutdown_grace: Duration::from_secs(10),
body_limits: Vec::new(),
#[cfg(feature = "tls")]
tls: None,
#[cfg(feature = "tls")]
tls_reload: None,
}
}
pub fn max_body_size(mut self, size: Option<usize>) -> Self {
self.max_body_size = size;
self
}
pub fn body_limit(mut self, prefix: &str, limit: Option<usize>) -> Self {
self.body_limits
.push((prefix.trim_end_matches('/').to_string(), limit));
self
}
pub fn state<T: Send + Sync + 'static>(mut self, value: T) -> Self {
self.state.insert(value);
self
}
pub fn background<F, Fut>(mut self, task: F) -> Self
where
F: FnOnce(Context) -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + Send + 'static,
{
self.background
.push(Box::new(move |ctx| Box::pin(task(ctx))));
self
}
#[cfg(feature = "tls")]
pub fn tls(mut self, config: crate::tls::TlsConfig) -> Self {
self.tls = Some(config);
self
}
#[cfg(feature = "tls")]
pub fn tls_reload_every(mut self, interval: Duration) -> Self {
self.tls_reload = Some(interval);
self
}
pub fn shutdown_grace(mut self, grace: Duration) -> Self {
self.shutdown_grace = grace;
self
}
pub fn group(mut self, prefix: &str, configure: impl FnOnce(RouteGroup) -> RouteGroup) -> Self {
let group = RouteGroup::new(prefix);
let group = configure(group);
if let Some(limit) = group.body_limit {
self.body_limits.push((group.prefix.clone(), limit));
}
for (method, path, op) in group.docs {
let full_path = format!("{}{}", group.prefix, path);
self.docs.insert((method, full_path), op);
}
let middlewares = group.middlewares;
for (method, path, handler) in group.routes {
let full_path = format!("{}{}", group.prefix, path);
let handler: DynHandler = if middlewares.is_empty() {
handler
} else {
Arc::new(Layered::new(middlewares.clone(), handler))
};
let kind = if method == "GET" && group.ws_paths.contains(&path) {
RouteKind::WebSocket
} else {
RouteKind::Http
};
if let Err(e) = self.router.insert_as(kind, &method, &full_path, handler) {
panic!("Failed to register {} route {}: {}", method, full_path, e);
}
}
self
}
pub fn doc(mut self, method: &str, path: &str, operation: Operation) -> Self {
self.docs
.insert((method.to_ascii_uppercase(), path.to_string()), operation);
self
}
pub fn openapi(mut self, info: OpenApiInfo) -> Self {
self.openapi = Some(info);
self
}
pub fn get<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
if let Err(e) = self
.router
.insert("GET", path, Arc::new(handler) as DynHandler)
{
panic!("Failed to register GET route {}: {}", path, e);
}
self
}
pub fn get_with_params<F, Fut>(mut self, path: &str, handler: F) -> Self
where
F: Fn(Params) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = Result<Response>> + Send + 'static,
{
if let Err(e) = self.router.insert(
"GET",
path,
Arc::new(HandlerWithParams::new(handler)) as DynHandler,
) {
panic!("Failed to register GET route {}: {}", path, e);
}
self
}
pub fn post<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
if let Err(e) = self
.router
.insert("POST", path, Arc::new(handler) as DynHandler)
{
panic!("Failed to register POST route {}: {}", path, e);
}
self
}
pub fn put<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
if let Err(e) = self
.router
.insert("PUT", path, Arc::new(handler) as DynHandler)
{
panic!("Failed to register PUT route {}: {}", path, e);
}
self
}
pub fn delete<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
if let Err(e) = self
.router
.insert("DELETE", path, Arc::new(handler) as DynHandler)
{
panic!("Failed to register DELETE route {}: {}", path, e);
}
self
}
pub fn patch<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
if let Err(e) = self
.router
.insert("PATCH", path, Arc::new(handler) as DynHandler)
{
panic!("Failed to register PATCH route {}: {}", path, e);
}
self
}
#[cfg(feature = "websocket")]
pub fn ws<F, Fut>(self, path: &str, handler: F) -> Self
where
F: Fn(Request, crate::websocket::WebSocket) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = ()> + Send + 'static,
{
self.ws_with(path, crate::websocket::WsOptions::default(), handler)
}
#[cfg(feature = "websocket")]
pub fn ws_with<F, Fut>(
mut self,
path: &str,
options: crate::websocket::WsOptions,
handler: F,
) -> Self
where
F: Fn(Request, crate::websocket::WebSocket) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = ()> + Send + 'static,
{
let route = crate::websocket::WsRoute::new(handler, options);
if let Err(e) = self.router.insert_as(
RouteKind::WebSocket,
"GET",
path,
Arc::new(route) as DynHandler,
) {
panic!("Failed to register WebSocket route {}: {}", path, e);
}
self
}
pub fn middleware<M>(mut self, middleware: M) -> Self
where
M: Middleware,
{
self.middlewares.push(Arc::new(middleware));
self
}
pub fn serve_static(mut self, route_prefix: &str, dir_path: &str) -> Self {
use crate::static_files::StaticFileHandler;
let handler = StaticFileHandler::new(dir_path);
let prefix = route_prefix.trim_end_matches('/');
let handler: DynHandler = Arc::new(handler);
let route = format!("{}/{{*filepath}}", prefix);
if let Err(e) = self
.router
.insert_as(RouteKind::Hidden, "GET", &route, handler.clone())
{
panic!("Failed to register static file route {}: {}", route, e);
}
let root = if prefix.is_empty() { "/" } else { prefix };
let _ = self
.router
.insert_as(RouteKind::Hidden, "GET", root, handler.clone());
if !prefix.is_empty() {
let _ =
self.router
.insert_as(RouteKind::Hidden, "GET", &format!("{}/", prefix), handler);
}
self
}
pub async fn run(self, addr: &str) -> Result<()> {
let addr: SocketAddr = addr
.parse()
.map_err(|e| Error::internal(format!("Invalid address: {}", e)))?;
let listener = TcpListener::bind(addr)
.await
.map_err(|e| Error::internal(format!("Failed to bind: {}", e)))?;
#[cfg(feature = "tls")]
let scheme = if self.tls.is_some() { "https" } else { "http" };
#[cfg(not(feature = "tls"))]
let scheme = "http";
println!("🚀 ullun-api listening on {}://{}", scheme, addr);
let result = self
.serve_with_shutdown(listener, async {
termination_signal().await;
println!("\n🛑 Shutting down gracefully...");
})
.await;
println!("✅ Server stopped");
result
}
pub async fn serve(self, listener: TcpListener) -> Result<()> {
self.serve_with_shutdown(listener, std::future::pending())
.await
}
pub async fn serve_with_shutdown(
mut self,
listener: TcpListener,
signal: impl std::future::Future<Output = ()>,
) -> Result<()> {
self = self.finalize();
let grace = self.shutdown_grace;
let background = std::mem::take(&mut self.background);
#[cfg(feature = "tls")]
let tls = self.tls.take();
#[cfg(feature = "tls")]
let tls_reload = self.tls_reload;
let (trigger, rx) = tokio::sync::watch::channel(false);
let shutdown = Shutdown::new(rx);
let shared = Arc::new(self.into_shared(shutdown.clone()));
#[allow(unused_mut)]
let mut tasks: Vec<_> = background
.into_iter()
.map(|task| {
let ctx = Context::new(shared.state.clone(), shutdown.clone());
tokio::spawn(task(ctx))
})
.collect();
#[cfg(feature = "tls")]
if let (Some(config), Some(every)) = (tls.clone(), tls_reload) {
let shutdown = shutdown.clone();
tasks.push(tokio::spawn(async move {
while shutdown.sleep(every).await {
match config.reload() {
Ok(true) => println!("🔐 TLS certificate reloaded"),
Ok(false) => {}
Err(e) => eprintln!(
"TLS certificate reload failed, keeping the previous one: {}",
e.message()
),
}
}
}));
}
let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::<()>(1);
tokio::pin!(signal);
loop {
tokio::select! {
_ = &mut signal => break,
accepted = listener.accept() => match accepted {
Ok((stream, peer)) => {
#[cfg(feature = "tls")]
if let Some(config) = tls.clone() {
let (shared, shutdown, done) = (shared.clone(), shutdown.clone(), done_tx.clone());
tokio::spawn(async move {
if let Ok((tls_stream, client_cert)) = config.accept(stream).await {
let info = ConnInfo { peer, secure: true, client_cert };
serve_connection(tls_stream, info, shared, shutdown, done).await;
}
});
continue;
}
let info = ConnInfo { peer, secure: false, client_cert: None };
tokio::spawn(serve_connection(
stream,
info,
shared.clone(),
shutdown.clone(),
done_tx.clone(),
));
}
Err(e) => {
eprintln!("Failed to accept connection: {}", e);
tokio::time::sleep(Duration::from_millis(100)).await;
}
},
}
}
let _ = trigger.send(true);
drop(done_tx);
drop(listener);
let drain = async {
let _ = done_rx.recv().await;
for task in tasks {
let _ = task.await;
}
};
if tokio::time::timeout(grace, drain).await.is_err() {
eprintln!(
"Shutdown grace period of {:?} elapsed, stopping anyway",
grace
);
}
Ok(())
}
fn into_shared(self, shutdown: Shutdown) -> Shared {
Shared {
router: self.router,
middlewares: self.middlewares,
max_body_size: self.max_body_size,
body_limits: self.body_limits,
state: Arc::new(self.state),
shutdown,
}
}
fn finalize(mut self) -> Self {
let Some(info) = self.openapi.take() else {
return self;
};
let spec = openapi::build(&info, self.router.registered(), &self.docs);
let spec = bytes::Bytes::from(serde_json::to_vec(&spec).unwrap_or_default());
let spec_handler = {
let spec = spec.clone();
move |_req: Request| {
let spec = spec.clone();
async move {
Ok(Response::ok()
.with_header("content-type", "application/json")
.with_body(spec))
}
}
};
self.insert_hidden("GET", &info.spec_path, Arc::new(spec_handler));
if let Some(docs_path) = &info.docs_path {
let page = openapi::swagger_ui_html(&info.title, &info.spec_path);
let page_handler = move |_req: Request| {
let page = page.clone();
async move { Ok(Response::html(page)) }
};
self.insert_hidden("GET", docs_path, Arc::new(page_handler));
}
self
}
fn insert_hidden(&mut self, method: &str, path: &str, handler: DynHandler) {
if let Err(e) = self
.router
.insert_as(RouteKind::Hidden, method, path, handler)
{
panic!("Failed to register {} route {}: {}", method, path, e);
}
}
}
async fn serve_connection<S>(
stream: S,
info: ConnInfo,
shared: Arc<Shared>,
shutdown: Shutdown,
_done: tokio::sync::mpsc::Sender<()>,
) where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
let io = TokioIo::new(stream);
let info = Arc::new(info);
let service = service_fn(|req| {
let (shared, info) = (shared.clone(), info.clone());
async move { Ok::<_, Infallible>(handle_request(req, &info, shared).await) }
});
let conn = http1::Builder::new()
.serve_connection(io, service)
.with_upgrades();
tokio::pin!(conn);
let mut draining = false;
loop {
tokio::select! {
result = conn.as_mut() => {
if let Err(err) = result {
eprintln!("Error serving connection: {:?}", err);
}
return;
}
_ = shutdown.wait(), if !draining => {
conn.as_mut().graceful_shutdown();
draining = true;
}
}
}
}
async fn termination_signal() {
#[cfg(unix)]
{
use tokio::signal::unix::{signal, SignalKind};
match signal(SignalKind::terminate()) {
Ok(mut term) => {
tokio::select! {
_ = tokio::signal::ctrl_c() => {}
_ = term.recv() => {}
}
return;
}
Err(e) => eprintln!("Cannot listen for SIGTERM: {}", e),
}
}
let _ = tokio::signal::ctrl_c().await;
}
impl Default for App {
fn default() -> Self {
Self::new()
}
}
impl RouteGroup {
fn new(prefix: &str) -> Self {
let prefix = prefix.trim_end_matches('/');
Self {
prefix: prefix.to_string(),
routes: Vec::new(),
middlewares: Vec::new(),
docs: Vec::new(),
body_limit: None,
ws_paths: Vec::new(),
}
}
#[cfg(feature = "websocket")]
pub fn ws_with<F, Fut>(
mut self,
path: &str,
options: crate::websocket::WsOptions,
handler: F,
) -> Self
where
F: Fn(Request, crate::websocket::WebSocket) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = ()> + Send + 'static,
{
let route = crate::websocket::WsRoute::new(handler, options);
self.routes.push((
"GET".to_string(),
path.to_string(),
Arc::new(route) as DynHandler,
));
self.ws_paths.push(path.to_string());
self
}
pub fn max_body_size(mut self, limit: Option<usize>) -> Self {
self.body_limit = Some(limit);
self
}
pub fn doc(mut self, method: &str, path: &str, operation: Operation) -> Self {
self.docs
.push((method.to_ascii_uppercase(), path.to_string(), operation));
self
}
pub fn get<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
self.routes.push((
"GET".to_string(),
path.to_string(),
Arc::new(handler) as DynHandler,
));
self
}
pub fn post<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
self.routes.push((
"POST".to_string(),
path.to_string(),
Arc::new(handler) as DynHandler,
));
self
}
pub fn put<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
self.routes.push((
"PUT".to_string(),
path.to_string(),
Arc::new(handler) as DynHandler,
));
self
}
pub fn delete<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
self.routes.push((
"DELETE".to_string(),
path.to_string(),
Arc::new(handler) as DynHandler,
));
self
}
pub fn patch<H>(mut self, path: &str, handler: H) -> Self
where
H: Handler,
{
self.routes.push((
"PATCH".to_string(),
path.to_string(),
Arc::new(handler) as DynHandler,
));
self
}
pub fn middleware<M>(mut self, middleware: M) -> Self
where
M: Middleware,
{
self.middlewares.push(Arc::new(middleware));
self
}
}
async fn handle_request(
req: hyper::Request<hyper::body::Incoming>,
conn: &ConnInfo,
shared: Arc<Shared>,
) -> hyper::Response<crate::response::HyperBody> {
let Shared {
router,
middlewares,
max_body_size,
body_limits,
state,
shutdown,
} = &*shared;
let method = req.method().to_string();
let is_head = req.method() == http::Method::HEAD;
let path = req.uri().path().to_string();
let max_body_size = body_limit_for(body_limits, &path, *max_body_size);
let (parts, body) = req.into_parts();
let body_bytes = match max_body_size {
Some(max_size) => match read_limited(body, max_size).await {
Ok(bytes) => bytes,
Err(too_large) => return too_large.into_hyper(max_size),
},
None => match body.collect().await {
Ok(collected) => collected.to_bytes(),
Err(_) => bytes::Bytes::new(),
},
};
let mut request = Request::new(parts.method, parts.uri, parts.headers, body_bytes);
request.extensions = parts.extensions;
request.remote_addr = Some(conn.peer);
request.secure = conn.secure;
if let Some(cert) = &conn.client_cert {
request
.extensions
.insert(crate::request::ClientCertificate(cert.clone()));
}
request.state = Some(state.clone());
request.shutdown = Some(shutdown.clone());
let route = match router.match_route_with_pattern(&method, &path) {
Ok((handler, params, pattern)) => {
request.params = params;
request.matched_path = Some(pattern);
Route::Found(handler)
}
Err(_) => Route::Miss {
allow: router.allowed_methods(&path),
},
};
let result = if middlewares.is_empty() {
dispatch(route, request).await
} else {
let mut chain = middlewares.to_vec();
let final_handler = Arc::new(move |req: Request, _next: Next| {
let route = route.clone();
async move { dispatch(route, req).await }
});
chain.push(final_handler as Arc<dyn Middleware>);
Next::new(Arc::new(chain)).run(request).await
};
let response = match result {
Ok(resp) => resp,
Err(err) => Response::from_error(err),
};
let response = if is_head {
response.into_head()
} else {
response
};
response.into_hyper()
}
const DRAIN_CAP: u64 = 8 * 1024 * 1024;
const DRAIN_TIME: Duration = Duration::from_secs(3);
struct BodyTooLarge;
impl BodyTooLarge {
fn into_hyper(self, max_size: usize) -> hyper::Response<crate::response::HyperBody> {
Response::from_error(Error::custom(
413,
format!("Request body too large (max {} bytes)", max_size),
))
.with_header("connection", "close")
.into_hyper()
}
}
async fn read_limited(
mut body: hyper::body::Incoming,
max: usize,
) -> std::result::Result<bytes::Bytes, BodyTooLarge> {
use bytes::BufMut;
use hyper::body::Body;
if body.size_hint().lower() > max as u64 {
drain(&mut body, 0).await;
return Err(BodyTooLarge);
}
let mut buf = bytes::BytesMut::new();
while let Some(frame) = body.frame().await {
let Ok(frame) = frame else {
return Ok(bytes::Bytes::new());
};
if let Ok(data) = frame.into_data() {
if buf.len() + data.len() > max {
drain(&mut body, (buf.len() + data.len()) as u64).await;
return Err(BodyTooLarge);
}
buf.put(data);
}
}
Ok(buf.freeze())
}
async fn drain(body: &mut hyper::body::Incoming, already: u64) {
let _ = tokio::time::timeout(DRAIN_TIME, async {
let mut seen = already;
while seen < DRAIN_CAP {
match body.frame().await {
Some(Ok(frame)) => {
if let Ok(data) = frame.into_data() {
seen += data.len() as u64;
}
}
_ => return,
}
}
})
.await;
}
fn body_limit_for(
rules: &[(String, Option<usize>)],
path: &str,
default: Option<usize>,
) -> Option<usize> {
rules
.iter()
.filter(|(prefix, _)| {
prefix.is_empty()
|| path == prefix
|| path
.strip_prefix(prefix.as_str())
.is_some_and(|rest| rest.starts_with('/'))
})
.max_by_key(|(prefix, _)| prefix.len())
.map(|(_, limit)| *limit)
.unwrap_or(default)
}
#[derive(Clone)]
enum Route {
Found(DynHandler),
Miss {
allow: Vec<&'static str>,
},
}
async fn dispatch(route: Route, req: Request) -> Result<Response> {
match route {
Route::Found(handler) => handler.call(req).await,
Route::Miss { allow } if allow.is_empty() => Err(Error::not_found("Route not found")),
Route::Miss { allow } => {
let allow = allow.join(", ");
if req.method == http::Method::OPTIONS {
Ok(Response::new(http::StatusCode::NO_CONTENT).with_header("allow", allow))
} else {
Ok(
Response::from_error(Error::custom(405, "Method not allowed"))
.with_header("allow", allow),
)
}
}
}
}