aro_web/app.rs
1//! Application builder for the Aro web framework.
2
3use axum::Router;
4use axum::extract::DefaultBodyLimit;
5use std::future::Future;
6use std::net::SocketAddr;
7use std::pin::Pin;
8use tokio::net::TcpListener;
9use tower_http::trace::TraceLayer;
10
11use std::sync::Arc;
12
13use crate::fallback::{fallback_404, fallback_405};
14use crate::state::{AroState, AroStateBuilder};
15
16/// A boxed startup hook: an async function that receives the resolved `AroState`
17/// and may fail with any error.
18type StartupHook = Box<
19 dyn FnOnce(
20 Arc<AroState>,
21 ) -> Pin<Box<dyn Future<Output = Result<(), Box<dyn std::error::Error>>>>>
22 + Send,
23>;
24
25struct UserLayer {
26 apply: Box<dyn Fn(Router) -> Router + Send + Sync>,
27}
28
29/// Builder for an Aro web application.
30///
31/// Wraps an Axum `Router` and applies framework defaults (tracing middleware,
32/// JSON error fallbacks) while allowing full customization.
33///
34/// # Example
35///
36/// ```no_run
37/// use aro_web::App;
38/// use axum::Router;
39/// use axum::routing::get;
40///
41/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
42/// App::new()
43/// .routes(Router::new().route("/", get(|| async { "ok" })))
44/// .serve("0.0.0.0:3000".parse()?)
45/// .await?;
46/// # Ok(())
47/// # }
48/// ```
49pub struct App {
50 router: Router,
51 stateful_router: Option<Router<AroState>>,
52 state: Option<AroState>,
53 builder: AroStateBuilder,
54 default_middleware: bool,
55 body_limit: Option<usize>,
56 user_layers: Vec<UserLayer>,
57 startup_hooks: Vec<StartupHook>,
58}
59
60impl App {
61 /// Create a new `App` with default settings.
62 pub fn new() -> Self {
63 Self {
64 router: Router::new(),
65 stateful_router: None,
66 state: None,
67 builder: AroState::builder(),
68 default_middleware: true,
69 body_limit: None,
70 user_layers: Vec::new(),
71 startup_hooks: Vec::new(),
72 }
73 }
74
75 /// Merge the given router's routes into the application.
76 ///
77 /// Can be called multiple times to compose routes from different modules.
78 pub fn routes(mut self, router: Router) -> Self {
79 self.router = self.router.merge(router);
80 self
81 }
82
83 /// Set the application state for dependency injection.
84 ///
85 /// The state is made available to handlers via the
86 /// [`Dep<T>`](crate::dep::Dep) extractor.
87 ///
88 /// Registrations from [`register`](Self::register) take precedence over
89 /// entries with the same type when both are present at [`build`](Self::build) time.
90 pub fn state(mut self, state: AroState) -> Self {
91 self.state = Some(state);
92 self
93 }
94
95 /// Register a dependency for injection.
96 ///
97 /// The type parameter `T` should be a trait object type (e.g. `dyn UserRepo`).
98 /// Registered dependencies are accumulated internally and merged into the
99 /// final [`AroState`] at [`build`](Self::build) time. When combined with
100 /// [`state`](Self::state), registrations from this method take precedence
101 /// over entries with the same type from external state.
102 ///
103 /// # Example
104 ///
105 /// ```ignore
106 /// // Domain types (`UserRepo`, `UserRepoImpl`) come from your app.
107 /// App::new()
108 /// .register::<dyn UserRepo>(Arc::new(UserRepoImpl::new(pool)))
109 /// .routes_with_state(routes)
110 /// .build()
111 /// ```
112 pub fn register<T: ?Sized + 'static>(mut self, service: Arc<T>) -> Self
113 where
114 Arc<T>: Send + Sync + Clone + 'static,
115 {
116 self.builder = self.builder.register::<T>(service);
117 self
118 }
119
120 /// Merge routes that require `AroState` into the application.
121 ///
122 /// Use this for routers containing handlers that use `Dep<T>` extractors.
123 /// The state must be set via [`state`](Self::state) before calling
124 /// [`build`](Self::build).
125 pub fn routes_with_state(mut self, router: Router<AroState>) -> Self {
126 self.stateful_router = Some(match self.stateful_router {
127 Some(existing) => existing.merge(router),
128 None => router,
129 });
130 self
131 }
132
133 /// Add a tower `Layer` to the application.
134 ///
135 /// User layers wrap the default middleware stack (outermost).
136 /// They execute in the order they are added and apply to all routes,
137 /// including those registered via [`routes_with_state`](Self::routes_with_state).
138 /// Each call wraps the stack, so the last layer added is outermost on inbound
139 /// requests.
140 pub fn layer<L>(mut self, layer: L) -> Self
141 where
142 L: tower::Layer<axum::routing::Route> + Clone + Send + Sync + 'static,
143 L::Service:
144 tower::Service<axum::http::Request<axum::body::Body>> + Clone + Send + Sync + 'static,
145 <L::Service as tower::Service<axum::http::Request<axum::body::Body>>>::Response:
146 axum::response::IntoResponse + 'static,
147 <L::Service as tower::Service<axum::http::Request<axum::body::Body>>>::Error:
148 Into<std::convert::Infallible> + 'static,
149 <L::Service as tower::Service<axum::http::Request<axum::body::Body>>>::Future:
150 Send + 'static,
151 {
152 self.user_layers.push(UserLayer {
153 apply: Box::new(move |router| router.layer(layer.clone())),
154 });
155 self
156 }
157
158 /// Set the maximum allowed request body size in bytes.
159 ///
160 /// Applies [`DefaultBodyLimit::max`] as a global layer, overriding
161 /// Axum's default 2 MB limit for extractors like `Json`, `Form`, and
162 /// `Bytes`. Individual routes can still override this with a
163 /// per-route `DefaultBodyLimit` layer.
164 ///
165 /// # Example
166 ///
167 /// ```no_run
168 /// use aro_web::App;
169 /// use axum::Router;
170 ///
171 /// let _app = App::new()
172 /// .body_limit(10 * 1024 * 1024) // 10 MB
173 /// .routes(Router::new())
174 /// .build();
175 /// ```
176 pub fn body_limit(mut self, max_bytes: usize) -> Self {
177 self.body_limit = Some(max_bytes);
178 self
179 }
180
181 /// Load layered configuration and register it as a dependency.
182 ///
183 /// Uses [`load_config`](crate::config::load_config) to build `T`
184 /// from environment variables and config files, then registers the result
185 /// so it can be extracted with [`Dep<T>`](crate::dep::Dep).
186 ///
187 /// # Errors
188 ///
189 /// Returns an error if configuration loading or deserialization fails.
190 pub fn config<T: serde::de::DeserializeOwned + Send + Sync + 'static>(
191 self,
192 ) -> Result<Self, crate::config::ConfigError> {
193 let cfg: T = crate::config::load_config()?;
194 Ok(self.register::<T>(Arc::new(cfg)))
195 }
196
197 /// Register an async startup hook.
198 ///
199 /// Hooks run in the order they are registered during [`serve()`](Self::serve)
200 /// or [`build_with_hooks()`](Self::build_with_hooks), **before** the TCP
201 /// listener is bound. Each hook receives a shared reference to the resolved
202 /// [`AroState`]. If any hook returns an error, the server will not start
203 /// and `serve()` will return the error.
204 ///
205 /// # Example
206 ///
207 /// ```no_run
208 /// use aro_web::App;
209 /// use axum::Router;
210 ///
211 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
212 /// App::new()
213 /// .on_startup(|_state| Box::pin(async move { Ok(()) }))
214 /// .routes(Router::new())
215 /// .serve("127.0.0.1:0".parse()?)
216 /// .await?;
217 /// # Ok(())
218 /// # }
219 /// ```
220 pub fn on_startup<F>(mut self, hook: F) -> Self
221 where
222 F: FnOnce(
223 Arc<AroState>,
224 )
225 -> Pin<Box<dyn Future<Output = Result<(), Box<dyn std::error::Error>>>>>
226 + Send
227 + 'static,
228 {
229 self.startup_hooks.push(Box::new(hook));
230 self
231 }
232
233 /// Apply a request timeout to all routes.
234 ///
235 /// Requests that are not completed within the given duration will receive
236 /// a `408 Request Timeout` response.
237 ///
238 /// This is the canonical builder entry point: the timeout applies to all
239 /// routes, including stateful DI routes added via
240 /// [`routes_with_state`](Self::routes_with_state).
241 ///
242 /// For manual `Router::layer(...)` composition, see
243 /// [`crate::middleware::timeout`].
244 ///
245 /// # Example
246 ///
247 /// ```no_run
248 /// use std::time::Duration;
249 ///
250 /// use aro_web::App;
251 /// use axum::Router;
252 ///
253 /// let _app = App::new()
254 /// .routes(Router::new())
255 /// .timeout(Duration::from_secs(30))
256 /// .build();
257 /// ```
258 pub fn timeout(self, duration: std::time::Duration) -> Self {
259 self.layer(crate::middleware::timeout(duration))
260 }
261
262 /// Disable the default middleware stack (`TraceLayer`).
263 ///
264 /// Call this if you want full control over the middleware pipeline.
265 /// The JSON fallback handler is always applied.
266 pub fn without_defaults(mut self) -> Self {
267 self.default_middleware = false;
268 self
269 }
270
271 /// Build the final `Router`, applying default middleware.
272 ///
273 /// The default middleware stack includes:
274 /// - JSON 404/405 fallback handlers for unmatched routes and unsupported methods
275 /// - `TraceLayer` from tower-http for request/response logging
276 ///
277 /// **Note:** Startup hooks registered via [`on_startup`](Self::on_startup)
278 /// are *not* executed by `build()`. Use [`build_with_hooks()`](Self::build_with_hooks)
279 /// or [`serve()`](Self::serve) to run them.
280 pub fn build(self) -> Router {
281 let (router, _state, _hooks) = self.into_parts();
282 router
283 }
284
285 /// Build the final `Router`, running all startup hooks first.
286 ///
287 /// This is the same as [`build()`](Self::build) but runs all registered
288 /// startup hooks in order before returning the router. Use this in tests
289 /// when your app relies on startup hooks (e.g., database migrations).
290 ///
291 /// # Errors
292 ///
293 /// Returns an error if any startup hook fails.
294 ///
295 /// # Example
296 ///
297 /// ```no_run
298 /// use aro_web::App;
299 /// use axum::Router;
300 ///
301 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
302 /// let _router = App::new()
303 /// .on_startup(|_state| Box::pin(async move { Ok(()) }))
304 /// .routes(Router::new())
305 /// .build_with_hooks()
306 /// .await?;
307 /// # Ok(())
308 /// # }
309 /// ```
310 pub async fn build_with_hooks(self) -> Result<Router, Box<dyn std::error::Error>> {
311 let (router, state, hooks) = self.into_parts();
312
313 let shared_state = Arc::new(state);
314 for hook in hooks {
315 hook(Arc::clone(&shared_state)).await?;
316 }
317
318 Ok(router)
319 }
320
321 /// Build and start serving the application.
322 ///
323 /// Runs all registered startup hooks in order, then binds to the given
324 /// address and serves with graceful shutdown (handles SIGINT/Ctrl+C and
325 /// SIGTERM).
326 ///
327 /// **Note:** Like [`build_with_hooks`](Self::build_with_hooks), this runs
328 /// startup hooks. [`build`](Self::build) alone does not.
329 ///
330 /// If any startup hook returns an error, the server is **not** started
331 /// and the error is returned.
332 pub async fn serve(self, addr: SocketAddr) -> Result<(), Box<dyn std::error::Error>> {
333 let app = self.build_with_hooks().await?;
334
335 let listener = TcpListener::bind(addr).await?;
336 tracing::info!("listening on {}", addr);
337 axum::serve(listener, app)
338 .with_graceful_shutdown(shutdown_signal())
339 .await?;
340 Ok(())
341 }
342
343 /// Consume the App and produce the final router, resolved state, and startup hooks.
344 fn into_parts(self) -> (Router, AroState, Vec<StartupHook>) {
345 let mut router = self.router;
346
347 // Build the final AroState by merging the internal builder with any
348 // externally provided state. Builder registrations (`.register()`)
349 // take precedence over entries from `.state()`.
350 let state = match self.state {
351 Some(external) => external.merge(self.builder.build()),
352 None => self.builder.build(),
353 };
354
355 // Merge stateful routes with the resolved state
356 if let Some(stateful) = self.stateful_router {
357 router = router.merge(stateful.with_state(state.clone()));
358 }
359
360 // Apply JSON fallbacks for unmatched routes and unsupported methods
361 router = router
362 .fallback(fallback_404)
363 .method_not_allowed_fallback(fallback_405);
364
365 // Apply body limit
366 if let Some(max_bytes) = self.body_limit {
367 router = router.layer(DefaultBodyLimit::max(max_bytes));
368 }
369
370 // Apply default middleware
371 if self.default_middleware {
372 router = router.layer(TraceLayer::new_for_http());
373 }
374
375 // Apply user layers outermost so they wrap the default stack
376 for user_layer in self.user_layers {
377 router = (user_layer.apply)(router);
378 }
379
380 (router, state, self.startup_hooks)
381 }
382}
383
384impl Default for App {
385 fn default() -> Self {
386 Self::new()
387 }
388}
389
390/// Wait for a shutdown signal (SIGINT or SIGTERM).
391#[expect(
392 clippy::expect_used,
393 reason = "signal handler installation failure is unrecoverable"
394)]
395async fn shutdown_signal() {
396 let ctrl_c = async {
397 tokio::signal::ctrl_c()
398 .await
399 .expect("failed to install Ctrl+C handler");
400 };
401
402 #[cfg(unix)]
403 let terminate = async {
404 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
405 .expect("failed to install SIGTERM handler")
406 .recv()
407 .await;
408 };
409
410 #[cfg(not(unix))]
411 let terminate = std::future::pending::<()>();
412
413 tokio::select! {
414 () = ctrl_c => {},
415 () = terminate => {},
416 }
417
418 tracing::info!("shutdown signal received, starting graceful shutdown");
419}