Properly failing a large request while using an extractor

[echelonh]

Summary

Hi!

I have the following problem -

I'm writing a server that expects large files in some POST requests. The serves uses an extractor layer to authenticate the request, or fail early if it fails authentication.

The code looks roughly like this:

// ROUTE INITIALIZATION
Router::new()
        .route("/*path", post(upload_file))
        .layer(from_extractor_with_state::<
            AuthExtractor,
            AuthState,
        >(AuthState::new_from_env()))


...

// UPLOAD FILE FUNCTION
async fn upload_file(
    path: Option<Path<String>>,
    headers: HeaderMap,
    body: String,
) -> Result<Response, Response> {
   ...
}

The problem occurs when a large enough file is uploaded and the request fails authentication. Instead of returning 401 Unauthorized the server returns 503 Service Unavailable.

This is probably due to the issue mentioned here #2168 , which itself is a byproduct of this issue: hyperium/hyper#2384 . TL;DR the request is supposed to fail early, body is not consumed, TCP reset is sent, connection is dropped (...ISTIO layer sends 503).

My question is -
What is the best way to properly send 401 and solve the "body drainage" issue? I'd prefer not to completely re-write the application if possible in some with a "lower-level" infrastructure, BUT - my auth extractor cannot drain the body "early" since the body must be used by the upload_file function.
Do I create some additional middleware that does the auth and provides the body as well? Do I have a way to drain the body early when the extractor fails?

Thanks a lot! 🙏

axum version

0.7.7

[mladedav]

You can use an auth layer instead of just an extractor. Then you can drain the body only if you intend to return 401 and leave it in place for inner services and the handler itself.

[echelonh]

Do you have an example I can look into? What's the difference from what I'm doing today?
I haven't mentioned it, but my auth layer needs to keep a certain state as well so it's not just a stateless middleware - I haven't found any matching examples online. Not sure if it matters.
Thanks!

[mladedav]

You can try using from_fn, but instead of your AuthExtractor, take Result<AuthExtractor, _>. If it is Err, consume the body and return 401, if it is ok, pass through to the next middleware.

I think there should be some middleware examples in this repository, I hope you can make it work with that. Feel free to ask clarifying questions if not.

[echelonh]

tbh I'm not sure why I worked at the time with an Extractor to begin with.
using from_fn_with_state did the trick.

For anyone who might reach this post in the future - I ended up with something like this:

pub async fn auth_middleware(
    State(state): State<AuthState>,
    mut request: Request,
    next: Next,
) -> Response {
    if let Err(e) = auth_middleware_inner(state, &mut request).await {
        let _ = request.into_body().collect().await;
        return e.into_response();
    }
    next.run(request).await
}

// just so that the `?` operator can be used with ease to short-circuit, AuthError impls IntoResponse
async fn auth_middleware_inner(
    state: AuthState,
    request: &mut Request,
) -> Result<(), AuthError> {
 /* ... */
}

And here's how to reproduce the issue (curl seems to work slightly different and the issue didn't reproduce there) - if you need to test your solution before & after a change:

wget -d http://localhost:8080/some_path --post-file=big.data