1use crate::error::{Error, Result};
4use crate::handler::{BoxFuture, DynHandler, Handler, HandlerWithParams};
5use crate::middleware::{Layered, Middleware, Next};
6use crate::openapi::{self, OpenApiInfo, Operation, RouteKind};
7use crate::request::{Params, Request, StateMap};
8use crate::response::Response;
9use crate::router::Router;
10use crate::shutdown::{Context, Shutdown};
11use http_body_util::BodyExt;
12use hyper::server::conn::http1;
13use hyper::service::service_fn;
14use hyper_util::rt::TokioIo;
15use std::collections::HashMap;
16use std::convert::Infallible;
17use std::net::SocketAddr;
18use std::sync::Arc;
19use std::time::Duration;
20use tokio::net::TcpListener;
21
22pub struct App {
24 router: Router,
25 middlewares: Vec<Arc<dyn Middleware>>,
26 max_body_size: Option<usize>, docs: HashMap<(String, String), Operation>,
28 openapi: Option<OpenApiInfo>,
29 state: StateMap,
30 background: Vec<BackgroundTask>,
31 shutdown_grace: Duration,
32 body_limits: Vec<(String, Option<usize>)>,
33}
34
35type BackgroundTask = Box<dyn FnOnce(Context) -> BoxFuture<'static, ()> + Send>;
36
37struct Shared {
39 router: Router,
40 middlewares: Vec<Arc<dyn Middleware>>,
41 max_body_size: Option<usize>,
42 body_limits: Vec<(String, Option<usize>)>,
43 state: Arc<StateMap>,
44 shutdown: Shutdown,
45}
46
47pub struct RouteGroup {
49 prefix: String,
50 routes: Vec<(String, String, DynHandler)>, middlewares: Vec<Arc<dyn Middleware>>,
52 docs: Vec<(String, String, Operation)>, body_limit: Option<Option<usize>>,
54}
55
56impl App {
57 pub fn new() -> Self {
58 Self {
59 router: Router::new(),
60 middlewares: Vec::new(),
61 max_body_size: Some(10 * 1024 * 1024), docs: HashMap::new(),
63 openapi: None,
64 state: StateMap::default(),
65 background: Vec::new(),
66 shutdown_grace: Duration::from_secs(10),
67 body_limits: Vec::new(),
68 }
69 }
70
71 pub fn max_body_size(mut self, size: Option<usize>) -> Self {
73 self.max_body_size = size;
74 self
75 }
76
77 pub fn body_limit(mut self, prefix: &str, limit: Option<usize>) -> Self {
84 self.body_limits
85 .push((prefix.trim_end_matches('/').to_string(), limit));
86 self
87 }
88
89 pub fn state<T: Send + Sync + 'static>(mut self, value: T) -> Self {
111 self.state.insert(value);
112 self
113 }
114
115 pub fn background<F, Fut>(mut self, task: F) -> Self
136 where
137 F: FnOnce(Context) -> Fut + Send + 'static,
138 Fut: std::future::Future<Output = ()> + Send + 'static,
139 {
140 self.background
141 .push(Box::new(move |ctx| Box::pin(task(ctx))));
142 self
143 }
144
145 pub fn shutdown_grace(mut self, grace: Duration) -> Self {
148 self.shutdown_grace = grace;
149 self
150 }
151
152 pub fn group(mut self, prefix: &str, configure: impl FnOnce(RouteGroup) -> RouteGroup) -> Self {
154 let group = RouteGroup::new(prefix);
155 let group = configure(group);
156
157 if let Some(limit) = group.body_limit {
158 self.body_limits.push((group.prefix.clone(), limit));
159 }
160 for (method, path, op) in group.docs {
161 let full_path = format!("{}{}", group.prefix, path);
162 self.docs.insert((method, full_path), op);
163 }
164
165 let middlewares = group.middlewares;
167 for (method, path, handler) in group.routes {
168 let full_path = format!("{}{}", group.prefix, path);
169 let handler: DynHandler = if middlewares.is_empty() {
170 handler
171 } else {
172 Arc::new(Layered::new(middlewares.clone(), handler))
173 };
174 if let Err(e) = self.router.insert(&method, &full_path, handler) {
175 panic!("Failed to register {} route {}: {}", method, full_path, e);
176 }
177 }
178
179 self
180 }
181
182 pub fn doc(mut self, method: &str, path: &str, operation: Operation) -> Self {
186 self.docs
187 .insert((method.to_ascii_uppercase(), path.to_string()), operation);
188 self
189 }
190
191 pub fn openapi(mut self, info: OpenApiInfo) -> Self {
195 self.openapi = Some(info);
196 self
197 }
198
199 pub fn get<H>(mut self, path: &str, handler: H) -> Self
201 where
202 H: Handler,
203 {
204 if let Err(e) = self
205 .router
206 .insert("GET", path, Arc::new(handler) as DynHandler)
207 {
208 panic!("Failed to register GET route {}: {}", path, e);
209 }
210 self
211 }
212
213 pub fn get_with_params<F, Fut>(mut self, path: &str, handler: F) -> Self
215 where
216 F: Fn(Params) -> Fut + Send + Sync + 'static,
217 Fut: std::future::Future<Output = Result<Response>> + Send + 'static,
218 {
219 if let Err(e) = self.router.insert(
220 "GET",
221 path,
222 Arc::new(HandlerWithParams::new(handler)) as DynHandler,
223 ) {
224 panic!("Failed to register GET route {}: {}", path, e);
225 }
226 self
227 }
228
229 pub fn post<H>(mut self, path: &str, handler: H) -> Self
231 where
232 H: Handler,
233 {
234 if let Err(e) = self
235 .router
236 .insert("POST", path, Arc::new(handler) as DynHandler)
237 {
238 panic!("Failed to register POST route {}: {}", path, e);
239 }
240 self
241 }
242
243 pub fn put<H>(mut self, path: &str, handler: H) -> Self
245 where
246 H: Handler,
247 {
248 if let Err(e) = self
249 .router
250 .insert("PUT", path, Arc::new(handler) as DynHandler)
251 {
252 panic!("Failed to register PUT route {}: {}", path, e);
253 }
254 self
255 }
256
257 pub fn delete<H>(mut self, path: &str, handler: H) -> Self
259 where
260 H: Handler,
261 {
262 if let Err(e) = self
263 .router
264 .insert("DELETE", path, Arc::new(handler) as DynHandler)
265 {
266 panic!("Failed to register DELETE route {}: {}", path, e);
267 }
268 self
269 }
270
271 pub fn patch<H>(mut self, path: &str, handler: H) -> Self
273 where
274 H: Handler,
275 {
276 if let Err(e) = self
277 .router
278 .insert("PATCH", path, Arc::new(handler) as DynHandler)
279 {
280 panic!("Failed to register PATCH route {}: {}", path, e);
281 }
282 self
283 }
284
285 #[cfg(feature = "websocket")]
291 pub fn ws<F, Fut>(self, path: &str, handler: F) -> Self
292 where
293 F: Fn(Request, crate::websocket::WebSocket) -> Fut + Send + Sync + 'static,
294 Fut: std::future::Future<Output = ()> + Send + 'static,
295 {
296 self.ws_with(path, crate::websocket::WsOptions::default(), handler)
297 }
298
299 #[cfg(feature = "websocket")]
301 pub fn ws_with<F, Fut>(
302 mut self,
303 path: &str,
304 options: crate::websocket::WsOptions,
305 handler: F,
306 ) -> Self
307 where
308 F: Fn(Request, crate::websocket::WebSocket) -> Fut + Send + Sync + 'static,
309 Fut: std::future::Future<Output = ()> + Send + 'static,
310 {
311 let route = crate::websocket::WsRoute::new(handler, options);
312 if let Err(e) = self.router.insert_as(
313 RouteKind::WebSocket,
314 "GET",
315 path,
316 Arc::new(route) as DynHandler,
317 ) {
318 panic!("Failed to register WebSocket route {}: {}", path, e);
319 }
320 self
321 }
322
323 pub fn middleware<M>(mut self, middleware: M) -> Self
325 where
326 M: Middleware,
327 {
328 self.middlewares.push(Arc::new(middleware));
329 self
330 }
331
332 pub fn serve_static(mut self, route_prefix: &str, dir_path: &str) -> Self {
334 use crate::static_files::StaticFileHandler;
335 let handler = StaticFileHandler::new(dir_path);
336 let prefix = route_prefix.trim_end_matches('/');
338 let handler: DynHandler = Arc::new(handler);
339 let route = format!("{}/{{*filepath}}", prefix);
340 if let Err(e) = self
341 .router
342 .insert_as(RouteKind::Hidden, "GET", &route, handler.clone())
343 {
344 panic!("Failed to register static file route {}: {}", route, e);
345 }
346 let root = if prefix.is_empty() { "/" } else { prefix };
349 let _ = self
350 .router
351 .insert_as(RouteKind::Hidden, "GET", root, handler.clone());
352 if !prefix.is_empty() {
353 let _ =
354 self.router
355 .insert_as(RouteKind::Hidden, "GET", &format!("{}/", prefix), handler);
356 }
357 self
358 }
359
360 pub async fn run(self, addr: &str) -> Result<()> {
365 let addr: SocketAddr = addr
366 .parse()
367 .map_err(|e| Error::internal(format!("Invalid address: {}", e)))?;
368
369 let listener = TcpListener::bind(addr)
370 .await
371 .map_err(|e| Error::internal(format!("Failed to bind: {}", e)))?;
372
373 println!("š ullun-api listening on http://{}", addr);
374
375 let result = self
376 .serve_with_shutdown(listener, async {
377 termination_signal().await;
378 println!("\nš Shutting down gracefully...");
379 })
380 .await;
381 println!("ā
Server stopped");
382 result
383 }
384
385 pub async fn serve(self, listener: TcpListener) -> Result<()> {
392 self.serve_with_shutdown(listener, std::future::pending())
393 .await
394 }
395
396 pub async fn serve_with_shutdown(
403 mut self,
404 listener: TcpListener,
405 signal: impl std::future::Future<Output = ()>,
406 ) -> Result<()> {
407 self = self.finalize();
408 let grace = self.shutdown_grace;
409 let background = std::mem::take(&mut self.background);
410 let (trigger, rx) = tokio::sync::watch::channel(false);
411 let shutdown = Shutdown::new(rx);
412 let shared = Arc::new(self.into_shared(shutdown.clone()));
413
414 let tasks: Vec<_> = background
415 .into_iter()
416 .map(|task| {
417 let ctx = Context::new(shared.state.clone(), shutdown.clone());
418 tokio::spawn(task(ctx))
419 })
420 .collect();
421
422 let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::<()>(1);
424 tokio::pin!(signal);
425 loop {
426 tokio::select! {
427 _ = &mut signal => break,
428 accepted = listener.accept() => match accepted {
429 Ok((stream, peer)) => {
430 tokio::spawn(serve_connection(
431 stream,
432 peer,
433 shared.clone(),
434 shutdown.clone(),
435 done_tx.clone(),
436 ));
437 }
438 Err(e) => {
439 eprintln!("Failed to accept connection: {}", e);
441 tokio::time::sleep(Duration::from_millis(100)).await;
442 }
443 },
444 }
445 }
446
447 let _ = trigger.send(true);
448 drop(done_tx);
449 drop(listener);
450 let drain = async {
451 let _ = done_rx.recv().await;
452 for task in tasks {
453 let _ = task.await;
454 }
455 };
456 if tokio::time::timeout(grace, drain).await.is_err() {
457 eprintln!(
458 "Shutdown grace period of {:?} elapsed, stopping anyway",
459 grace
460 );
461 }
462 Ok(())
463 }
464
465 fn into_shared(self, shutdown: Shutdown) -> Shared {
466 Shared {
467 router: self.router,
468 middlewares: self.middlewares,
469 max_body_size: self.max_body_size,
470 body_limits: self.body_limits,
471 state: Arc::new(self.state),
472 shutdown,
473 }
474 }
475
476 fn finalize(mut self) -> Self {
479 let Some(info) = self.openapi.take() else {
480 return self;
481 };
482 let spec = openapi::build(&info, self.router.registered(), &self.docs);
483 let spec = bytes::Bytes::from(serde_json::to_vec(&spec).unwrap_or_default());
484
485 let spec_handler = {
486 let spec = spec.clone();
487 move |_req: Request| {
488 let spec = spec.clone();
489 async move {
490 Ok(Response::ok()
491 .with_header("content-type", "application/json")
492 .with_body(spec))
493 }
494 }
495 };
496 self.insert_hidden("GET", &info.spec_path, Arc::new(spec_handler));
497
498 if let Some(docs_path) = &info.docs_path {
499 let page = openapi::swagger_ui_html(&info.title, &info.spec_path);
500 let page_handler = move |_req: Request| {
501 let page = page.clone();
502 async move { Ok(Response::html(page)) }
503 };
504 self.insert_hidden("GET", docs_path, Arc::new(page_handler));
505 }
506 self
507 }
508
509 fn insert_hidden(&mut self, method: &str, path: &str, handler: DynHandler) {
510 if let Err(e) = self
511 .router
512 .insert_as(RouteKind::Hidden, method, path, handler)
513 {
514 panic!("Failed to register {} route {}: {}", method, path, e);
515 }
516 }
517}
518
519async fn serve_connection(
521 stream: tokio::net::TcpStream,
522 peer: SocketAddr,
523 shared: Arc<Shared>,
524 shutdown: Shutdown,
525 _done: tokio::sync::mpsc::Sender<()>,
526) {
527 let io = TokioIo::new(stream);
528 let service = service_fn(|req| {
529 let shared = shared.clone();
530 async move { Ok::<_, Infallible>(handle_request(req, peer, shared).await) }
531 });
532 let conn = http1::Builder::new()
533 .serve_connection(io, service)
534 .with_upgrades();
535 tokio::pin!(conn);
536
537 let mut draining = false;
538 loop {
539 tokio::select! {
540 result = conn.as_mut() => {
541 if let Err(err) = result {
542 eprintln!("Error serving connection: {:?}", err);
543 }
544 return;
545 }
546 _ = shutdown.wait(), if !draining => {
547 conn.as_mut().graceful_shutdown();
549 draining = true;
550 }
551 }
552 }
553}
554
555async fn termination_signal() {
557 #[cfg(unix)]
558 {
559 use tokio::signal::unix::{signal, SignalKind};
560 match signal(SignalKind::terminate()) {
561 Ok(mut term) => {
562 tokio::select! {
563 _ = tokio::signal::ctrl_c() => {}
564 _ = term.recv() => {}
565 }
566 return;
567 }
568 Err(e) => eprintln!("Cannot listen for SIGTERM: {}", e),
569 }
570 }
571 let _ = tokio::signal::ctrl_c().await;
572}
573
574impl Default for App {
575 fn default() -> Self {
576 Self::new()
577 }
578}
579
580impl RouteGroup {
581 fn new(prefix: &str) -> Self {
582 let prefix = prefix.trim_end_matches('/');
583 Self {
584 prefix: prefix.to_string(),
585 routes: Vec::new(),
586 middlewares: Vec::new(),
587 docs: Vec::new(),
588 body_limit: None,
589 }
590 }
591
592 pub fn max_body_size(mut self, limit: Option<usize>) -> Self {
595 self.body_limit = Some(limit);
596 self
597 }
598
599 pub fn doc(mut self, method: &str, path: &str, operation: Operation) -> Self {
602 self.docs
603 .push((method.to_ascii_uppercase(), path.to_string(), operation));
604 self
605 }
606
607 pub fn get<H>(mut self, path: &str, handler: H) -> Self
608 where
609 H: Handler,
610 {
611 self.routes.push((
612 "GET".to_string(),
613 path.to_string(),
614 Arc::new(handler) as DynHandler,
615 ));
616 self
617 }
618
619 pub fn post<H>(mut self, path: &str, handler: H) -> Self
620 where
621 H: Handler,
622 {
623 self.routes.push((
624 "POST".to_string(),
625 path.to_string(),
626 Arc::new(handler) as DynHandler,
627 ));
628 self
629 }
630
631 pub fn put<H>(mut self, path: &str, handler: H) -> Self
632 where
633 H: Handler,
634 {
635 self.routes.push((
636 "PUT".to_string(),
637 path.to_string(),
638 Arc::new(handler) as DynHandler,
639 ));
640 self
641 }
642
643 pub fn delete<H>(mut self, path: &str, handler: H) -> Self
644 where
645 H: Handler,
646 {
647 self.routes.push((
648 "DELETE".to_string(),
649 path.to_string(),
650 Arc::new(handler) as DynHandler,
651 ));
652 self
653 }
654
655 pub fn patch<H>(mut self, path: &str, handler: H) -> Self
656 where
657 H: Handler,
658 {
659 self.routes.push((
660 "PATCH".to_string(),
661 path.to_string(),
662 Arc::new(handler) as DynHandler,
663 ));
664 self
665 }
666
667 pub fn middleware<M>(mut self, middleware: M) -> Self
672 where
673 M: Middleware,
674 {
675 self.middlewares.push(Arc::new(middleware));
676 self
677 }
678}
679
680async fn handle_request(
681 req: hyper::Request<hyper::body::Incoming>,
682 peer: SocketAddr,
683 shared: Arc<Shared>,
684) -> hyper::Response<crate::response::HyperBody> {
685 let Shared {
686 router,
687 middlewares,
688 max_body_size,
689 body_limits,
690 state,
691 shutdown,
692 } = &*shared;
693 let method = req.method().to_string();
694 let is_head = req.method() == http::Method::HEAD;
695 let path = req.uri().path().to_string();
696 let max_body_size = body_limit_for(body_limits, &path, *max_body_size);
697
698 let (parts, body) = req.into_parts();
700
701 let body_bytes = match max_body_size {
704 Some(max_size) => match read_limited(body, max_size).await {
705 Ok(bytes) => bytes,
706 Err(too_large) => return too_large.into_hyper(max_size),
707 },
708 None => match body.collect().await {
709 Ok(collected) => collected.to_bytes(),
710 Err(_) => bytes::Bytes::new(),
711 },
712 };
713
714 let mut request = Request::new(parts.method, parts.uri, parts.headers, body_bytes);
715 request.extensions = parts.extensions;
716 request.remote_addr = Some(peer);
717 request.state = Some(state.clone());
718 request.shutdown = Some(shutdown.clone());
719
720 let route = match router.match_route(&method, &path) {
723 Ok((handler, params)) => {
724 request.params = params;
725 Route::Found(handler)
726 }
727 Err(_) => Route::Miss {
728 allow: router.allowed_methods(&path),
729 },
730 };
731
732 let result = if middlewares.is_empty() {
733 dispatch(route, request).await
734 } else {
735 let mut chain = middlewares.to_vec();
737 let final_handler = Arc::new(move |req: Request, _next: Next| {
738 let route = route.clone();
739 async move { dispatch(route, req).await }
740 });
741 chain.push(final_handler as Arc<dyn Middleware>);
742
743 Next::new(Arc::new(chain)).run(request).await
744 };
745
746 let response = match result {
748 Ok(resp) => resp,
749 Err(err) => Response::from_error(err),
750 };
751
752 let response = if is_head {
754 response.into_head()
755 } else {
756 response
757 };
758
759 response.into_hyper()
760}
761
762const DRAIN_CAP: u64 = 8 * 1024 * 1024;
766const DRAIN_TIME: Duration = Duration::from_secs(3);
767
768struct BodyTooLarge;
769
770impl BodyTooLarge {
771 fn into_hyper(self, max_size: usize) -> hyper::Response<crate::response::HyperBody> {
772 Response::from_error(Error::custom(
773 413,
774 format!("Request body too large (max {} bytes)", max_size),
775 ))
776 .with_header("connection", "close")
778 .into_hyper()
779 }
780}
781
782async fn read_limited(
784 mut body: hyper::body::Incoming,
785 max: usize,
786) -> std::result::Result<bytes::Bytes, BodyTooLarge> {
787 use bytes::BufMut;
788 use hyper::body::Body;
789
790 if body.size_hint().lower() > max as u64 {
792 drain(&mut body, 0).await;
793 return Err(BodyTooLarge);
794 }
795
796 let mut buf = bytes::BytesMut::new();
797 while let Some(frame) = body.frame().await {
798 let Ok(frame) = frame else {
799 return Ok(bytes::Bytes::new());
800 };
801 if let Ok(data) = frame.into_data() {
802 if buf.len() + data.len() > max {
803 drain(&mut body, (buf.len() + data.len()) as u64).await;
804 return Err(BodyTooLarge);
805 }
806 buf.put(data);
807 }
808 }
809 Ok(buf.freeze())
810}
811
812async fn drain(body: &mut hyper::body::Incoming, already: u64) {
814 let _ = tokio::time::timeout(DRAIN_TIME, async {
815 let mut seen = already;
816 while seen < DRAIN_CAP {
817 match body.frame().await {
818 Some(Ok(frame)) => {
819 if let Ok(data) = frame.into_data() {
820 seen += data.len() as u64;
821 }
822 }
823 _ => return,
824 }
825 }
826 })
827 .await;
828}
829
830fn body_limit_for(
832 rules: &[(String, Option<usize>)],
833 path: &str,
834 default: Option<usize>,
835) -> Option<usize> {
836 rules
837 .iter()
838 .filter(|(prefix, _)| {
839 prefix.is_empty()
840 || path == prefix
841 || path
842 .strip_prefix(prefix.as_str())
843 .is_some_and(|rest| rest.starts_with('/'))
844 })
845 .max_by_key(|(prefix, _)| prefix.len())
846 .map(|(_, limit)| *limit)
847 .unwrap_or(default)
848}
849
850#[derive(Clone)]
852enum Route {
853 Found(DynHandler),
854 Miss {
856 allow: Vec<&'static str>,
857 },
858}
859
860async fn dispatch(route: Route, req: Request) -> Result<Response> {
861 match route {
862 Route::Found(handler) => handler.call(req).await,
863 Route::Miss { allow } if allow.is_empty() => Err(Error::not_found("Route not found")),
864 Route::Miss { allow } => {
865 let allow = allow.join(", ");
866 if req.method == http::Method::OPTIONS {
867 Ok(Response::new(http::StatusCode::NO_CONTENT).with_header("allow", allow))
870 } else {
871 Ok(
872 Response::from_error(Error::custom(405, "Method not allowed"))
873 .with_header("allow", allow),
874 )
875 }
876 }
877 }
878}