from_fn 与它的兄弟

pub fn from_fn<F, T>(f: F) -> FromFnLayer<F, (), T> {
    from_fn_with_state((), f)
}

        impl<F, Fut, Out, S, I, $($ty,)* $last> Service<Request> for FromFn<F, S, I, ($($ty,)* $last,)>
        where
            F: FnMut($($ty,)* $last, Next) -> Fut + Clone + Send + 'static,
            $( $ty: FromRequestParts<S> + Send, )*
            $last: FromRequest<S> + Send,
            Fut: Future<Output = Out> + Send + 'static,
            Out: IntoResponse + 'static,
            I: Service<Request, Error = Infallible>
            // ...
            fn call(&mut self, req: Request) -> Self::Future {
                let not_ready_inner = self.inner.clone();
                let ready_inner = std::mem::replace(&mut self.inner, not_ready_inner);

                let mut f = self.f.clone();
                let state = self.state.clone();
                let (mut parts, body) = req.into_parts();

                let future = Box::pin(async move {
                    $(
                        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(),
                    };

                    let inner = BoxCloneSyncService::new(MapIntoResponse::new(ready_inner));
                    let next = Next { inner };

                    f($($ty,)* $last, next).await.into_response()
                });

axum = "=0.8.9"
serde = { version = "1", features = ["derive"] }
tower = { version = "=0.5.3", features = ["util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt"] }

use std::sync::atomic::{AtomicUsize, Ordering};

use axum::body::Body;
use axum::extract::{FromRequestParts, Query, Request, State};
use axum::http::{request::Parts, HeaderMap, HeaderValue, StatusCode};
use axum::middleware::{from_extractor, from_fn, from_fn_with_state, map_request, map_response, Next};
use axum::response::Response;
use axum::routing::get;
use axum::Router;
use serde::Deserialize;
use tower::ServiceExt;

#[derive(Deserialize)]
struct Lang {
  lang: String,
}

// Extractors first, then Request, then Next.
async fn with_extractors(Query(q): Query<Lang>, headers: HeaderMap, req: Request, next: Next) -> Response {
  println!("    mw: lang={} ua={:?}", q.lang, headers.get("user-agent"));
  next.run(req).await
}

#[derive(Clone)]
struct Limits {
  max_len: usize,
}

// State given to from_fn_with_state, independent of the Router's state.
async fn check_len(State(limits): State<Limits>, req: Request, next: Next) -> Result<Response, StatusCode> {
  if req.uri().path().len() > limits.max_len {
    return Err(StatusCode::URI_TOO_LONG);
  }
  Ok(next.run(req).await)
}

// map_request: Ok(req) goes on, Err(e) becomes the response.
async fn require_token(req: Request) -> Result<Request, StatusCode> {
  if req.headers().contains_key("token") { Ok(req) } else { Err(StatusCode::UNAUTHORIZED) }
}

async fn add_header(mut res: Response) -> Response {
  res.headers_mut().insert("x-mw", HeaderValue::from_static("map_response"));
  res
}

static AUTH_RUNS: AtomicUsize = AtomicUsize::new(0);

struct Auth;

impl<S: Send + Sync> FromRequestParts<S> for Auth {
  type Rejection = (StatusCode, &'static str);

  async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
    AUTH_RUNS.fetch_add(1, Ordering::SeqCst);
    match parts.headers.get("token") {
      Some(v) if v == "secret" => Ok(Auth),
      _ => Err((StatusCode::UNAUTHORIZED, "bad token")),
    }
  }
}

async fn handler() -> &'static str {
  println!("    handler");
  "ok"
}

async fn handler_with_auth(_auth: Auth) -> &'static str {
  println!("    handler");
  "ok"
}

async fn send(app: &Router, uri: &str, token: Option<&str>) {
  let mut b = Request::get(uri).header("user-agent", "demo");
  if let Some(t) = token {
    b = b.header("token", t);
  }
  let res = app.clone().oneshot(b.body(Body::empty()).unwrap()).await.unwrap();
  let status = res.status();
  let x = res.headers().get("x-mw").cloned();
  let body = axum::body::to_bytes(res.into_body(), usize::MAX).await.unwrap();
  println!("  {uri} token={token:?} -> {} x-mw={x:?} body={body:?}", status.as_u16());
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  println!("[from_fn with Query + HeaderMap]");
  let app = Router::new().route("/", get(handler)).layer(from_fn(with_extractors));
  send(&app, "/?lang=zh", None).await;
  send(&app, "/", None).await;

  println!("[from_fn_with_state]");
  let app = Router::new()
    .route("/{*rest}", get(handler))
    .layer(from_fn_with_state(Limits { max_len: 8 }, check_len));
  send(&app, "/short", None).await;
  send(&app, "/much-too-long", None).await;

  println!("[map_request + map_response]");
  let app = Router::new()
    .route("/", get(handler))
    .layer(map_request(require_token))
    .layer(map_response(add_header));
  send(&app, "/", Some("x")).await;
  send(&app, "/", None).await;

  println!("[from_extractor::<Auth>, handler also takes Auth]");
  let app = Router::new()
    .route("/", get(handler_with_auth))
    .route_layer(from_extractor::<Auth>());
  for token in [Some("secret"), Some("wrong")] {
    AUTH_RUNS.store(0, Ordering::SeqCst);
    send(&app, "/", token).await;
    println!("  Auth::from_request_parts ran {} time(s)", AUTH_RUNS.load(Ordering::SeqCst));
  }
}

[from_fn with Query + HeaderMap]
    mw: lang=zh ua=Some("demo")
    handler
  /?lang=zh token=None -> 200 x-mw=None body=b"ok"
  / token=None -> 400 x-mw=None body=b"Failed to deserialize query string: missing field `lang`"

async fn mw(req: Request, _headers: HeaderMap, next: Next) -> Response {
  next.run(req).await
}

error[E0277]: the trait bound `FromFn<fn(Request<Body>, _, Next) -> _ {mw}, (), Route, _>: Service<_>` is not satisfied
   --> src/bin/req_not_last.rs:8:77
    |
  8 |   let _app: Router = Router::new().route("/", get(|| async { "ok" })).layer(from_fn(mw));
    |                                                                       ----- ^^^^^^^^^^^ unsatisfied trait bound
    |                                                                       |
    |                                                                       required by a bound introduced by this call
    |
    = help: the trait `tower_service::Service<axum::http::Request<Body>>` is not implemented for `FromFn<fn(Request<Body>, HeaderMap, Next) -> _ {mw}, (), Route, _>`
    = help: the following other types implement trait `tower_service::Service<Request>`:
              axum::middleware::FromFn<F, S, I, (T1, T2)>
              axum::middleware::FromFn<F, S, I, (T1, T2, T3)>
              axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4)>
              axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5)>
              axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6)>
              axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6, T7)>
              axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6, T7, T8)>
              axum::middleware::FromFn<F, S, I, (T1, T2, T3, T4, T5, T6, T7, T8, T9)>
            and 8 others
note: required by a bound in `Router::<S>::layer`
   --> axum-0.8.9/src/routing/mod.rs:306:21
    |
303 |     pub fn layer<L>(self, layer: L) -> Router<S>
    |            ----- required by a bound in this associated function
...
306 |         L::Service: Service<Request> + Clone + Send + Sync + 'static,
    |                     ^^^^^^^^^^^^^^^^ required by this bound in `Router::<S>::layer`

error: `Request<_>` consumes the request body and thus must be the last argument to the handler function
 --> src/bin/req_not_last_dbg.rs:4:18
  |
4 | async fn mw(req: Request, _headers: HeaderMap, next: Next) -> Response {
  |                  ^^^^^^^

#[derive(Debug, Clone)]
pub struct Next {
    inner: BoxCloneSyncService<Request, Response, Infallible>,
}

impl Next {
    /// Execute the remaining middleware stack.
    pub async fn run(mut self, req: Request) -> Response {
        match self.inner.call(req).await {
            Ok(res) => res,
            Err(err) => match err {},
        }
    }
}

axum = "=0.8.9"
tower = { version = "=0.5.3", features = ["util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt"] }

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering};

use axum::body::Body;
use axum::extract::Request;
use axum::middleware::{from_fn, map_request, Next};
use axum::response::Response;
use axum::routing::get;
use axum::Router;
use tower::{ServiceBuilder, ServiceExt};
use tower::util::MapRequestLayer;

struct Counting;

static ALLOCS: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
  unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    ALLOCS.fetch_add(1, Ordering::Relaxed);
    unsafe { System.alloc(layout) }
  }
  unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    unsafe { System.dealloc(ptr, layout) }
  }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

async fn pass(req: Request, next: Next) -> Response {
  next.run(req).await
}

async fn pass_map(req: Request) -> Request {
  req
}

fn ident(req: Request) -> Request {
  req
}

async fn measure(name: &str, app: Router) {
  let app = app.with_state(()); // turn handlers into Route up front
  let req = || Request::get("/").body(Body::empty()).unwrap();
  app.clone().oneshot(req()).await.unwrap(); // warm up

  const N: usize = 10_000;
  let mut reqs: Vec<_> = (0..N).map(|_| req()).collect();
  let before = ALLOCS.load(Ordering::Relaxed);
  for r in reqs.drain(..) {
    let res = app.clone().oneshot(r).await.unwrap();
    drop(res);
  }
  let allocs = (ALLOCS.load(Ordering::Relaxed) - before) as f64 / N as f64;
  println!("{name:<26} allocs/req = {allocs:.1}");
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let base = || Router::new().route("/", get(|| async { "ok" }));
  measure("no middleware", base()).await;
  measure("tower MapRequestLayer", base().layer(MapRequestLayer::new(ident))).await;
  measure("map_request", base().layer(map_request(pass_map))).await;
  measure("from_fn x1", base().layer(from_fn(pass))).await;
  measure("from_fn x2", base().layer(from_fn(pass)).layer(from_fn(pass))).await;
  measure("from_fn x3", base().layer(from_fn(pass)).layer(from_fn(pass)).layer(from_fn(pass))).await;
  measure("MapRequestLayer x2", base().layer(MapRequestLayer::new(ident)).layer(MapRequestLayer::new(ident))).await;
  measure("MapRequestLayer x3", base().layer(MapRequestLayer::new(ident)).layer(MapRequestLayer::new(ident)).layer(MapRequestLayer::new(ident))).await;
  let sb = ServiceBuilder::new()
    .layer(MapRequestLayer::new(ident))
    .layer(MapRequestLayer::new(ident))
    .layer(MapRequestLayer::new(ident));
  measure("ServiceBuilder MapReq x3", base().layer(sb)).await;
  let sb = ServiceBuilder::new().layer(from_fn(pass)).layer(from_fn(pass)).layer(from_fn(pass));
  measure("ServiceBuilder from_fn x3", base().layer(sb)).await;
}

no middleware              allocs/req = 11.0
tower MapRequestLayer      allocs/req = 13.0
map_request                allocs/req = 14.0
from_fn x1                 allocs/req = 16.0
from_fn x2                 allocs/req = 23.0
from_fn x3                 allocs/req = 32.0
MapRequestLayer x2         allocs/req = 16.0
MapRequestLayer x3         allocs/req = 20.0
ServiceBuilder MapReq x3   allocs/req = 13.0
ServiceBuilder from_fn x3  allocs/req = 22.0

[from_fn_with_state]
    handler
  /short token=None -> 200 x-mw=None body=b"ok"
  /much-too-long token=None -> 414 x-mw=None body=b""

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)
    }
}

                    match f($($ty,)* $last).await.into_map_request_result() {
                        Ok(req) => {
                            ready_inner.call(req).await.into_response()
                        }
                        Err(res) => {
                            res
                        }
                    }

                    match ready_inner.call(req).await {
                        Ok(res) => {
                            f($($ty,)* res).await.into_response()

[map_request + map_response]
    handler
  / token=Some("x") -> 200 x-mw=Some("map_response") body=b"ok"
  / token=None -> 401 x-mw=Some("map_response") body=b""

    fn call(&mut self, req: Request<B>) -> Self::Future {
        let state = self.state.clone();
        let (mut parts, body) = req.into_parts();

        let extract_future = Box::pin(async move {
            let extracted = E::from_request_parts(&mut parts, &state).await;
            let req = Request::from_parts(parts, body);
            (req, extracted)
        });

        ResponseFuture {
            state: State::Extracting {
                future: extract_future,
            },
            svc: Some(self.inner.clone()),
        }
    }

[from_extractor::<Auth>, handler also takes Auth]
    handler
  / token=Some("secret") -> 200 x-mw=None body=b"ok"
  Auth::from_request_parts ran 2 time(s)
  / token=Some("wrong") -> 401 x-mw=None body=b"bad token"
  Auth::from_request_parts ran 1 time(s)