Merge pull request #5804 from Turbo87/axum-404

Handle 404 errors in `axum`

Author: Turbo87

src/lib.rs

@@ -31,7 +31,6 @@ use std::sync::Arc;
 
 use crate::app::AppState;
 use crate::router::build_axum_router;
-use conduit_axum::ConduitFallback;
 use tikv_jemallocator::Jemalloc;
 
 #[global_allocator]
@@ -84,11 +83,8 @@ pub enum Env {
 pub fn build_handler(app: Arc<App>) -> axum::Router {
     let state = AppState(app);
 
-    let endpoints = router::build_router();
-    let conduit_handler = middleware::build_middleware(state.clone(), endpoints);
-
     let axum_router = build_axum_router(state.clone());
-    middleware::apply_axum_middleware(state, axum_router.conduit_fallback(conduit_handler))
+    middleware::apply_axum_middleware(state, axum_router)
 }
 
 /// Convenience function requiring that an environment variable is set.

src/middleware.rs

@@ -5,16 +5,12 @@ mod prelude {
     pub use http::{header, Response, StatusCode};
 }
 
-use self::app::AppMiddleware;
-use self::known_error_to_json::KnownErrorToJson;
-
 pub mod app;
 mod balance_capacity;
 mod block_traffic;
 mod debug;
 mod ember_html;
 mod head;
-mod known_error_to_json;
 pub mod log_request;
 pub mod normalize_path;
 mod require_user_agent;
@@ -23,9 +19,6 @@ pub mod session;
 mod static_or_continue;
 mod update_metrics;
 
-use conduit_middleware::MiddlewareBuilder;
-use conduit_router::RouteBuilder;
-
 use ::sentry::integrations::tower as sentry_tower;
 use axum::error_handling::HandleErrorLayer;
 use axum::middleware::{from_fn, from_fn_with_state};
@@ -107,12 +100,3 @@ pub fn apply_axum_middleware(state: AppState, router: Router) -> Router {
 async fn dummy_error_handler(_err: axum::BoxError) -> http::StatusCode {
     http::StatusCode::INTERNAL_SERVER_ERROR
 }
-
-pub fn build_middleware(app: AppState, endpoints: RouteBuilder) -> MiddlewareBuilder {
-    let mut m = MiddlewareBuilder::new(endpoints);
-
-    m.add(AppMiddleware::new(app));
-    m.add(KnownErrorToJson);
-
-    m
-}

src/middleware/app.rs

@@ -6,24 +6,6 @@ use http::Request;
 
 use crate::app::AppState;
 
-/// Middleware that injects the `App` instance into the `Request` extensions
-pub struct AppMiddleware {
-    app: AppState,
-}
-
-impl AppMiddleware {
-    pub fn new(app: AppState) -> AppMiddleware {
-        AppMiddleware { app }
-    }
-}
-
-impl Middleware for AppMiddleware {
-    fn before(&self, req: &mut dyn RequestExt) -> BeforeResult {
-        req.mut_extensions().insert(self.app.clone());
-        Ok(())
-    }
-}
-
 /// `axum` middleware that injects the `AppState` instance into the `Request` extensions.
 pub async fn add_app_state_extension<B>(
     State(app_state): State<AppState>,

src/middleware/known_error_to_json.rs

@@ -1,15 +1,16 @@
 use axum::handler::Handler as AxumHandler;
 use axum::middleware::from_fn_with_state;
+use axum::response::IntoResponse;
 use axum::routing::{delete, get, post, put};
 use axum::Router;
 use conduit::{Handler, HandlerResult, RequestExt};
 use conduit_axum::{CauseField, ConduitAxumHandler};
-use conduit_router::{RouteBuilder, RoutePattern};
+use conduit_router::RoutePattern;
 
 use crate::app::AppState;
 use crate::controllers::*;
 use crate::middleware::app::{add_app_state_extension, RequestApp};
-use crate::util::errors::{AppError, RouteBlocked};
+use crate::util::errors::{not_found, AppError, RouteBlocked};
 use crate::util::EndpointResult;
 use crate::Env;
 
@@ -197,11 +198,9 @@ pub fn build_axum_router(state: AppState) -> Router {
         );
     }
 
-    router.with_state(state)
-}
-
-pub fn build_router() -> RouteBuilder {
-    RouteBuilder::new()
+    router
+        .fallback(|| async { not_found().into_response() })
+        .with_state(state)
 }
 
 struct C(pub fn(&mut dyn RequestExt) -> EndpointResult);