OpenSource/axum
OpenSource/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2025-01-01 08:56:15 +01:00
Code Issues Projects Releases Packages Wiki Activity

Add map_request and friends (#1408)

Browse source 
* Add `map_request` and friends

* finish it

* changelog ref

* Apply suggestions from code review

Co-authored-by: Jonas Platte <jplatte+git@posteo.de>

* address review feedback

* Apply suggestions from code review

Co-authored-by: Jonas Platte <jplatte+git@posteo.de>

Co-authored-by: Jonas Platte <jplatte+git@posteo.de>
This commit is contained in:
David Pedersen 2022-09-25 13:44:10 +02:00 committed by GitHub
parent 896ffc5fba
commit 2077d50021
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 483 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
axum/CHANGELOG.md
View file

@ -18,6 +18,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- **added:** Add `middleware::from_extractor_with_state` and
`middleware::from_extractor_with_state_arc` ([#1396])
- **added:** Add `DefaultBodyLimit::max` for changing the default body limit ([#1397])
- **added:** Add `map_request`, `map_request_with_state`, and
`map_request_with_state_arc` for transforming the request with an async
function ([#1408])
- **breaking:** `ContentLengthLimit` has been removed. `Use DefaultBodyLimit` instead ([#1400])
[#1371]: https://github.com/tokio-rs/axum/pull/1371
@ -26,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
[#1396]: https://github.com/tokio-rs/axum/pull/1396
[#1397]: https://github.com/tokio-rs/axum/pull/1397
[#1400]: https://github.com/tokio-rs/axum/pull/1400
[#1408]: https://github.com/tokio-rs/axum/pull/1408
# 0.6.0-rc.2 (10. September, 2022)

11
axum/src/middleware/from_fn.rs
View file

@ -77,6 +77,9 @@ use tower_service::Service;
/// async fn my_middleware<B>(
/// TypedHeader(auth): TypedHeader<Authorization<Bearer>>,
/// Query(query_params): Query<HashMap<String, String>>,
/// // you can add more extractors here but the last
/// // extractor must implement `FromRequest` which
/// // `Request` does
/// req: Request<B>,
/// next: Next<B>,
/// ) -> Response {
@ -117,7 +120,9 @@ pub fn from_fn<F, T>(f: F) -> FromFnLayer<F, (), T> {
///
/// async fn my_middleware<B>(
/// State(state): State<AppState>,
/// // you can add more extractors here...
/// // you can add more extractors here but the last
/// // extractor must implement `FromRequest` which
/// // `Request` does
/// req: Request<B>,
/// next: Next<B>,
/// ) -> Response {
@ -140,6 +145,8 @@ pub fn from_fn_with_state<F, S, T>(state: S, f: F) -> FromFnLayer<F, S, T> {
/// Create a middleware from an async function with the given [`Arc`]'ed state.
///
/// See [`from_fn_with_state`] for an example.
///
/// See [`State`](crate::extract::State) for more details about accessing state.
pub fn from_fn_with_state_arc<F, S, T>(state: Arc<S>, f: F) -> FromFnLayer<F, S, T> {
FromFnLayer {
@ -396,7 +403,7 @@ mod tests {
}
async fn handle(headers: HeaderMap) -> String {
(&headers["x-axum-test"]).to_str().unwrap().to_owned()
headers["x-axum-test"].to_str().unwrap().to_owned()
}
let app = Router::new()

464
axum/src/middleware/map_request.rs Normal file
View file

@ -0,0 +1,464 @@
use crate::response::{IntoResponse, Response};
use axum_core::extract::{FromRequest, FromRequestParts};
use futures_util::future::BoxFuture;
use http::Request;
use std::{
any::type_name,
convert::Infallible,
fmt,
future::Future,
marker::PhantomData,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use tower_layer::Layer;
use tower_service::Service;
/// Create a middleware from an async function that transforms a request.
///
/// This differs from [`tower::util::MapRequest`] in that it allows you to easily run axum-specific
/// extractors.
///
/// # Example
///
/// ```
/// use axum::{
/// Router,
/// routing::get,
/// middleware::map_request,
/// http::Request,
/// };
///
/// async fn set_header<B>(mut request: Request<B>) -> Request<B> {
/// request.headers_mut().insert("x-foo", "foo".parse().unwrap());
/// request
/// }
///
/// async fn handler<B>(request: Request<B>) {
/// // `request` will have an `x-foo` header
/// }
///
/// let app = Router::new()
/// .route("/", get(handler))
/// .layer(map_request(set_header));
/// # let _: Router = app;
/// ```
///
/// # Rejection the request
///
/// The function given to `map_request` is allowed to also return a `Result` which can be used to
/// reject the request and return a response immediately, without calling the remaining
/// middleware.
///
/// Specifically the valid return types are:
///
/// - `Request<B>`
/// - `Request<Request<B>, E> where E: IntoResponse`
///
/// ```
/// use axum::{
/// Router,
/// http::{Request, StatusCode},
/// routing::get,
/// middleware::map_request,
/// };
///
/// async fn auth<B>(request: Request<B>) -> Result<Request<B>, StatusCode> {
/// let auth_header = request.headers()
/// .get(http::header::AUTHORIZATION)
/// .and_then(|header| header.to_str().ok());
///
/// match auth_header {
/// Some(auth_header) if token_is_valid(auth_header) => Ok(request),
/// _ => Err(StatusCode::UNAUTHORIZED),
/// }
/// }
///
/// fn token_is_valid(token: &str) -> bool {
/// // ...
/// # false
/// }
///
/// let app = Router::new()
/// .route("/", get(|| async { /* ... */ }))
/// .route_layer(map_request(auth));
/// # let app: Router = app;
/// ```
///
/// # Running extractors
///
/// ```
/// use axum::{
/// Router,
/// routing::get,
/// middleware::map_request,
/// extract::Path,
/// http::Request,
/// };
/// use std::collections::HashMap;
///
/// async fn log_path_params<B>(
/// Path(path_params): Path<HashMap<String, String>>,
/// request: Request<B>,
/// ) -> Request<B> {
/// tracing::debug!(?path_params);
/// request
/// }
///
/// let app = Router::new()
/// .route("/", get(|| async { /* ... */ }))
/// .layer(map_request(log_path_params));
/// # let _: Router = app;
/// ```
///
/// Note that to access state you must use either [`map_request_with_state`] or [`map_request_with_state_arc`].
pub fn map_request<F, T>(f: F) -> MapRequestLayer<F, (), T> {
map_request_with_state((), f)
}
/// Create a middleware from an async function that transforms a request, with the given state.
///
/// See [`State`](crate::extract::State) for more details about accessing state.
///
/// # Example
///
/// ```rust
/// use axum::{
/// Router,
/// http::{Request, StatusCode},
/// routing::get,
/// response::IntoResponse,
/// middleware::map_request_with_state,
/// extract::State,
/// };
///
/// #[derive(Clone)]
/// struct AppState { /* ... */ }
///
/// async fn my_middleware<B>(
/// State(state): State<AppState>,
/// // you can add more extractors here but the last
/// // extractor must implement `FromRequest` which
/// // `Request` does
/// request: Request<B>,
/// ) -> Request<B> {
/// // do something with `state` and `request`...
/// request
/// }
///
/// let state = AppState { /* ... */ };
///
/// let app = Router::with_state(state.clone())
/// .route("/", get(|| async { /* ... */ }))
/// .route_layer(map_request_with_state(state, my_middleware));
/// # let app: Router<_> = app;
/// ```
pub fn map_request_with_state<F, S, T>(state: S, f: F) -> MapRequestLayer<F, S, T> {
map_request_with_state_arc(Arc::new(state), f)
}
/// Create a middleware from an async function that transforms a request, with the given [`Arc`]'ed
/// state.
///
/// See [`map_request_with_state`] for an example.
///
/// See [`State`](crate::extract::State) for more details about accessing state.
pub fn map_request_with_state_arc<F, S, T>(state: Arc<S>, f: F) -> MapRequestLayer<F, S, T> {
MapRequestLayer {
f,
state,
_extractor: PhantomData,
}
}
/// A [`tower::Layer`] from an async function that transforms a request.
///
/// Created with [`map_request`]. See that function for more details.
pub struct MapRequestLayer<F, S, T> {
f: F,
state: Arc<S>,
_extractor: PhantomData<fn() -> T>,
}
impl<F, S, T> Clone for MapRequestLayer<F, S, T>
where
F: Clone,
{
fn clone(&self) -> Self {
Self {
f: self.f.clone(),
state: Arc::clone(&self.state),
_extractor: self._extractor,
}
}
}
impl<S, I, F, T> Layer<I> for MapRequestLayer<F, S, T>
where
F: Clone,
{
type Service = MapRequest<F, S, I, T>;
fn layer(&self, inner: I) -> Self::Service {
MapRequest {
f: self.f.clone(),
state: Arc::clone(&self.state),
inner,
_extractor: PhantomData,
}
}
}
impl<F, S, T> fmt::Debug for MapRequestLayer<F, S, T>
where
S: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MapRequestLayer")
// Write out the type name, without quoting it as `&type_name::<F>()` would
.field("f", &format_args!("{}", type_name::<F>()))
.field("state", &self.state)
.finish()
}
}
/// A middleware created from an async function that transforms a request.
///
/// Created with [`map_request`]. See that function for more details.
pub struct MapRequest<F, S, I, T> {
f: F,
inner: I,
state: Arc<S>,
_extractor: PhantomData<fn() -> T>,
}
impl<F, S, I, T> Clone for MapRequest<F, S, I, T>
where
F: Clone,
I: Clone,
{
fn clone(&self) -> Self {
Self {
f: self.f.clone(),
inner: self.inner.clone(),
state: Arc::clone(&self.state),
_extractor: self._extractor,
}
}
}
macro_rules! impl_service {
(
[$($ty:ident),*], $last:ident
) => {
#[allow(non_snake_case, unused_mut)]
impl<F, Fut, S, I, B, $($ty,)* $last> Service<Request<B>> for MapRequest<F, S, I, ($($ty,)* $last,)>
where
F: FnMut($($ty,)* $last) -> Fut + Clone + Send + 'static,
$( $ty: FromRequestParts<S> + Send, )*
$last: FromRequest<S, B> + Send,
Fut: Future + Send + 'static,
Fut::Output: IntoMapRequestResult<B> + Send + 'static,
I: Service<Request<B>, Error = Infallible>
+ Clone
+ Send
+ 'static,
I::Response: IntoResponse,
I::Future: Send + 'static,
B: Send + 'static,
S: Send + Sync + 'static,
{
type Response = Response;
type Error = Infallible;
type Future = ResponseFuture;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<B>) -> Self::Future {
let not_ready_inner = self.inner.clone();
let mut ready_inner = std::mem::replace(&mut self.inner, not_ready_inner);
let mut f = self.f.clone();
let state = Arc::clone(&self.state);
let future = Box::pin(async move {
let (mut parts, body) = req.into_parts();
$(
let $ty = match $ty::from_request_parts(&mut parts, &state).await {
Ok(value) => value,
Err(rejection) => return rejection.into_response(),
};
)*
let req = Request::from_parts(parts, body);
let $last = match $last::from_request(req, &state).await {
Ok(value) => value,
Err(rejection) => return rejection.into_response(),
};
match f($($ty,)* $last).await.into_map_request_result() {
Ok(req) => {
ready_inner.call(req).await.into_response()
}
Err(res) => {
res
}
}
});
ResponseFuture {
inner: future
}
}
}
};
}
impl_service!([], T1);
impl_service!([T1], T2);
impl_service!([T1, T2], T3);
impl_service!([T1, T2, T3], T4);
impl_service!([T1, T2, T3, T4], T5);
impl_service!([T1, T2, T3, T4, T5], T6);
impl_service!([T1, T2, T3, T4, T5, T6], T7);
impl_service!([T1, T2, T3, T4, T5, T6, T7], T8);
impl_service!([T1, T2, T3, T4, T5, T6, T7, T8], T9);
impl_service!([T1, T2, T3, T4, T5, T6, T7, T8, T9], T10);
impl_service!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10], T11);
impl_service!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11], T12);
impl_service!([T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12], T13);
impl_service!(
[T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13],
T14
);
impl_service!(
[T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14],
T15
);
impl_service!(
[T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15],
T16
);
impl<F, S, I, T> fmt::Debug for MapRequest<F, S, I, T>
where
S: fmt::Debug,
I: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MapRequest")
.field("f", &format_args!("{}", type_name::<F>()))
.field("inner", &self.inner)
.field("state", &self.state)
.finish()
}
}
/// Response future for [`MapRequest`].
pub struct ResponseFuture {
inner: BoxFuture<'static, Response>,
}
impl Future for ResponseFuture {
type Output = Result<Response, Infallible>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.inner.as_mut().poll(cx).map(Ok)
}
}
impl fmt::Debug for ResponseFuture {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ResponseFuture").finish()
}
}
mod private {
use crate::{http::Request, response::IntoResponse};
pub trait Sealed<B> {}
impl<B, E> Sealed<B> for Result<Request<B>, E> where E: IntoResponse {}
impl<B> Sealed<B> for Request<B> {}
}
/// Trait implemented by types that can be returned from [`map_request`],
/// [`map_request_with_state`], and [`map_request_with_state_arc`].
///
/// This trait is sealed such that it cannot be implemented outside this crate.
pub trait IntoMapRequestResult<B>: private::Sealed<B> {
/// Perform the conversion.
fn into_map_request_result(self) -> Result<Request<B>, Response>;
}
impl<B, E> IntoMapRequestResult<B> for Result<Request<B>, E>
where
E: IntoResponse,
{
fn into_map_request_result(self) -> Result<Request<B>, Response> {
self.map_err(IntoResponse::into_response)
}
}
impl<B> IntoMapRequestResult<B> for Request<B> {
fn into_map_request_result(self) -> Result<Request<B>, Response> {
Ok(self)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{routing::get, test_helpers::TestClient, Router};
use http::{HeaderMap, StatusCode};
#[tokio::test]
async fn works() {
async fn add_header<B>(mut req: Request<B>) -> Request<B> {
req.headers_mut().insert("x-foo", "foo".parse().unwrap());
req
}
async fn handler(headers: HeaderMap) -> Response {
headers["x-foo"]
.to_str()
.unwrap()
.to_owned()
.into_response()
}
let app = Router::new()
.route("/", get(handler))
.layer(map_request(add_header));
let client = TestClient::new(app);
let res = client.get("/").send().await;
assert_eq!(res.text().await, "foo");
}
#[tokio::test]
async fn works_for_short_circutting() {
async fn add_header<B>(_req: Request<B>) -> Result<Request<B>, (StatusCode, &'static str)> {
Err((StatusCode::INTERNAL_SERVER_ERROR, "something went wrong"))
}
async fn handler(_headers: HeaderMap) -> Response {
unreachable!()
}
let app = Router::new()
.route("/", get(handler))
.layer(map_request(add_header));
let client = TestClient::new(app);
let res = client.get("/").send().await;
assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR);
assert_eq!(res.text().await, "something went wrong");
}
}

6
axum/src/middleware/mod.rs
View file

@ -4,6 +4,7 @@
mod from_extractor;
mod from_fn;
mod map_request;
pub use self::from_extractor::{
from_extractor, from_extractor_with_state, from_extractor_with_state_arc, FromExtractor,
@ -12,6 +13,10 @@ pub use self::from_extractor::{
pub use self::from_fn::{
from_fn, from_fn_with_state, from_fn_with_state_arc, FromFn, FromFnLayer, Next,
};
pub use self::map_request::{
map_request, map_request_with_state, map_request_with_state_arc, IntoMapRequestResult,
MapRequest, MapRequestLayer,
};
pub use crate::extension::AddExtension;
pub mod future {
@ -19,4 +24,5 @@ pub mod future {
pub use super::from_extractor::ResponseFuture as FromExtractorResponseFuture;
pub use super::from_fn::ResponseFuture as FromFnResponseFuture;
pub use super::map_request::ResponseFuture as MapRequestResponseFuture;
}