diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md index 23727d30..be2d7693 100644 --- a/axum/CHANGELOG.md +++ b/axum/CHANGELOG.md @@ -7,7 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 # Unreleased -- None. +- **added:** Add `Router::route_layer` for applying middleware that + will only run on requests that match a route ([#474]) + +[#474]: https://github.com/tokio-rs/axum/pull/474 # 0.3.1 (06. November, 2021) diff --git a/axum/src/docs/routing/route_layer.md b/axum/src/docs/routing/route_layer.md new file mode 100644 index 00000000..58a42cd8 --- /dev/null +++ b/axum/src/docs/routing/route_layer.md @@ -0,0 +1,28 @@ +Apply a [`tower::Layer`] to the router that will only run if the request matches +a route. + +This works similarly to [`Router::layer`] except the middleware will only run if +the request matches a route. This is useful for middleware that return early +(such as authorization) which might otherwise convert a `404 Not Found` into a +`401 Unauthorized`. + +# Example + +```rust +use axum::{ + routing::get, + Router, +}; +use tower_http::auth::RequireAuthorizationLayer; + +let app = Router::new() + .route("/foo", get(|| async {})) + .route_layer(RequireAuthorizationLayer::bearer("password")); + +// `GET /foo` with a valid token will receive `200 OK` +// `GET /foo` with a invalid token will receive `401 Unauthorized` +// `GET /not-found` with a invalid token will receive `404 Not Found` +# async { +# axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap(); +# }; +``` diff --git a/axum/src/extract/extractor_middleware.rs b/axum/src/extract/extractor_middleware.rs index e7aeee79..85015268 100644 --- a/axum/src/extract/extractor_middleware.rs +++ b/axum/src/extract/extractor_middleware.rs @@ -37,38 +37,52 @@ use tower_service::Service; /// /// ```rust /// use axum::{ -/// Router, -/// async_trait, /// extract::{extractor_middleware, FromRequest, RequestParts}, -/// http::StatusCode, /// routing::{get, post}, +/// Router, /// }; -/// use std::convert::Infallible; +/// use http::StatusCode; +/// use async_trait::async_trait; /// -/// struct MyExtractor; +/// // An extractor that performs authorization. +/// struct RequireAuth; /// /// #[async_trait] -/// impl FromRequest for MyExtractor +/// impl FromRequest for RequireAuth /// where /// B: Send, /// { -/// type Rejection = Infallible; +/// type Rejection = StatusCode; /// /// async fn from_request(req: &mut RequestParts) -> Result { -/// # Ok(Self) -/// // ... +/// let auth_header = req +/// .headers() +/// .and_then(|headers| headers.get(http::header::AUTHORIZATION)) +/// .and_then(|value| value.to_str().ok()); +/// +/// if let Some(value) = auth_header { +/// if value == "secret" { +/// return Ok(Self); +/// } +/// } +/// +/// Err(StatusCode::UNAUTHORIZED) /// } /// } /// -/// async fn handler() {} +/// async fn handler() { +/// // If we get here the request has been authorized +/// } /// -/// async fn other_handler() {} +/// async fn other_handler() { +/// // If we get here the request has been authorized +/// } /// /// let app = Router::new() /// .route("/", get(handler)) /// .route("/foo", post(other_handler)) /// // The extractor will run before all routes -/// .layer(extractor_middleware::()); +/// .route_layer(extractor_middleware::()); /// # async { /// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap(); /// # }; diff --git a/axum/src/routing/mod.rs b/axum/src/routing/mod.rs index 68321d4d..b9d8ba06 100644 --- a/axum/src/routing/mod.rs +++ b/axum/src/routing/mod.rs @@ -223,19 +223,16 @@ where } #[doc = include_str!("../docs/routing/layer.md")] - pub fn layer(self, layer: L) -> Router + pub fn layer(self, layer: L) -> Router where L: Layer>, - L::Service: Service< - Request, - Response = Response, - Error = Infallible, - > + Clone + L::Service: Service, Response = Response, Error = Infallible> + + Clone + Send + 'static, - >>::Future: Send + 'static, - LayeredResBody: http_body::Body + Send + 'static, - LayeredResBody::Error: Into, + >>::Future: Send + 'static, + NewResBody: http_body::Body + Send + 'static, + NewResBody::Error: Into, { let layer = ServiceBuilder::new() .layer_fn(Route::new) @@ -249,7 +246,7 @@ where let route = Layer::layer(&layer, route); (id, route) }) - .collect::>>(); + .collect(); let fallback = self.fallback.map(|svc| Layer::layer(&layer, svc)); @@ -260,6 +257,39 @@ where } } + #[doc = include_str!("../docs/routing/route_layer.md")] + pub fn route_layer(self, layer: L) -> Self + where + L: Layer>, + L::Service: Service, Response = Response, Error = Infallible> + + Clone + + Send + + 'static, + >>::Future: Send + 'static, + NewResBody: http_body::Body + Send + 'static, + NewResBody::Error: Into, + { + let layer = ServiceBuilder::new() + .layer_fn(Route::new) + .layer(MapResponseBodyLayer::new(box_body)) + .layer(layer); + + let routes = self + .routes + .into_iter() + .map(|(id, route)| { + let route = Layer::layer(&layer, route); + (id, route) + }) + .collect(); + + Router { + routes, + node: self.node, + fallback: self.fallback, + } + } + #[doc = include_str!("../docs/routing/fallback.md")] pub fn fallback(mut self, svc: T) -> Self where diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs index c674bd26..84bb4565 100644 --- a/axum/src/routing/tests/mod.rs +++ b/axum/src/routing/tests/mod.rs @@ -478,3 +478,31 @@ async fn middleware_still_run_for_unmatched_requests() { async fn routing_to_router_panics() { TestClient::new(Router::new().route("/", Router::new())); } + +#[tokio::test] +async fn route_layer() { + let app = Router::new() + .route("/foo", get(|| async {})) + .route_layer(RequireAuthorizationLayer::bearer("password")); + + let client = TestClient::new(app); + + let res = client + .get("/foo") + .header("authorization", "Bearer password") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let res = client.get("/foo").send().await; + assert_eq!(res.status(), StatusCode::UNAUTHORIZED); + + let res = client.get("/not-found").send().await; + assert_eq!(res.status(), StatusCode::NOT_FOUND); + + // it would be nice if this would return `405 Method Not Allowed` + // but that requires knowing more about which method route we're calling, which we + // don't know currently since its just a generic `Service` + let res = client.post("/foo").send().await; + assert_eq!(res.status(), StatusCode::UNAUTHORIZED); +}