Merge pull request #646 from swlody/master

Add Hooks::before_routes to give user control over initial axum::Router construction

docs-site/content/docs/resources/faq.md

@@ -54,6 +54,8 @@ Loco is compatible with Axum middlewares. Simply implement `FromRequestParts` in
 <details>
 <summary>Injecting Custom State or Layers in Loco?</summary>
 Yes, you can achieve this by implementing `Hooks::after_routes`. This hook receive Axum routers that Loco has already built, allowing you to seamlessly add any available Axum functions that suit your needs.
+
+If you need your routes or (404) fallback handler to be affected by loco's middleware, you can add them in `Hooks::before_routes` which is called before the middleware is installed.
 </details>
 
 <br/>

src/app.rs

@@ -39,7 +39,7 @@ pub struct AppContext {
     /// The environment in which the application is running.
     pub environment: Environment,
     #[cfg(feature = "with-db")]
-    /// A database connection used by the application.    
+    /// A database connection used by the application.
     pub db: DatabaseConnection,
     /// An optional connection pool for Queue, for worker tasks
     pub queue: Option<Pool<RedisConnectionManager>>,
@@ -131,6 +131,16 @@ pub trait Hooks {
         Ok(false)
     }
 
+    /// Returns the initial Axum router for the application, allowing the user
+    /// to control the construction of the Axum router. This is where a fallback
+    /// handler can be installed before middleware or other routes are added.
+    ///
+    /// # Errors
+    /// Return an [`Result`] when the router could not be created
+    async fn before_routes(_ctx: &AppContext) -> Result<AxumRouter<AppContext>> {
+        Ok(AxumRouter::new())
+    }
+
     /// Invoke this function after the Loco routers have been constructed. This
     /// function enables you to configure custom Axum logics, such as layers,
     /// that are compatible with Axum.
@@ -178,11 +188,11 @@ pub trait Hooks {
     /// function. The truncate controlled from the [`crate::config::Database`]
     /// by changing dangerously_truncate to true (default false).
     /// Truncate can be useful when you want to truncate the database before any
-    /// test.        
+    /// test.
     #[cfg(feature = "with-db")]
     async fn truncate(db: &DatabaseConnection) -> Result<()>;
 
-    /// Seeds the database with initial data.    
+    /// Seeds the database with initial data.
     #[cfg(feature = "with-db")]
     async fn seed(db: &DatabaseConnection, path: &Path) -> Result<()>;
 }

src/boot.rs

@@ -265,7 +265,8 @@ pub async fn run_app<H: Hooks>(mode: &StartMode, app_context: AppContext) -> Res
     }
     match mode {
         StartMode::ServerOnly => {
-            let app = H::routes(&app_context).to_router(app_context.clone())?;
+            let app = H::before_routes(&app_context).await?;
+            let app = H::routes(&app_context).to_router(app_context.clone(), app)?;
             let mut router = H::after_routes(app, &app_context).await?;
             for initializer in &initializers {
                 router = initializer.after_routes(router, &app_context).await?;
@@ -279,7 +280,8 @@ pub async fn run_app<H: Hooks>(mode: &StartMode, app_context: AppContext) -> Res
         }
         StartMode::ServerAndWorker => {
             let processor = create_processor::<H>(&app_context)?;
-            let app = H::routes(&app_context).to_router(app_context.clone())?;
+            let app = H::before_routes(&app_context).await?;
+            let app = H::routes(&app_context).to_router(app_context.clone(), app)?;
             let mut router = H::after_routes(app, &app_context).await?;
             for initializer in &initializers {
                 router = initializer.after_routes(router, &app_context).await?;

src/controller/app_routes.rs

@@ -171,16 +171,14 @@ impl AppRoutes {
         self
     }
 
-    /// Convert the routes to an Axum Router, and set a list of middlewares that
-    /// configure in the [`config::Config`]
+    /// Add the routes to an existing Axum Router, and set a list of middlewares
+    /// that configure in the [`config::Config`]
     ///
     /// # Errors
     /// Return an [`Result`] when could not convert the router setup to
     /// [`axum::Router`].
     #[allow(clippy::cognitive_complexity)]
-    pub fn to_router(&self, ctx: AppContext) -> Result<AXRouter> {
-        let mut app = AXRouter::new();
-
+    pub fn to_router(&self, ctx: AppContext, mut app: AXRouter<AppContext>) -> Result<AXRouter> {
         for router in self.collect() {
             tracing::info!("{}", router.to_string());