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