How to read the state from the router in middleware?

[eisterman]

Summary

Hello.
I'm trying to implement a complex nested router with axum_utoipa.
I need to generate the openapi.json from the router, and the state I'm using is very big and I cannot freely generate one from scratch.
So my plan was to make a function return the Router with the state missing, and:

1. In the main program plug the state properly generated
2. In the openapi.json generator not plug the state, and use the router only for the schema gen.

This is just the reason why I'm asking this strange question.
Basically I need for 3 specific routes a middleware that do some authentication work, and this middleware needs to read information from the State.
The problem is that the middleware requires to be defined a State upfront and cannot read the state from the router itself, breaking completely my present pipeline.

Is there some way/workaround to catch the state inside a middleware WITHOUT needing to give it upfront and instead reading it from the Router object?

axum version

0.8.4

[mladedav]

Do you mean something like this?

[mladedav]

Then I don't understand what you're asking.

You can do this

let app = Router::new()
    .route("/", get(|| async { /* ... */ }))
    .route_layer(middleware::from_fn_with_state(state.clone(), my_middleware));

And then use that in any way you want and the call with_state whenever you need to run the app.

Can you try to describe in different words what you're trying to achieve then?

[eisterman]

It's really a pain to explain, I'm sorry...
When I use from_fn_with_state, I need to have the state struct available. I know seems strange but my function that returns the Router is not guaranteed to have the State always available.
This is because the state is a very big struct with side effects encoded in it plus my need for an OpenAPI spec.
This is my exact code. The Router is wrapped into utoipa-axum OpenApiRouter but inside it's just a Router and an OpenApiSpec struct.

pub fn init_router() -> OpenApiRouter<AppState> {
    let api_routes = OpenApiRouter::new()
        .merge(peers::routes())
        .merge(newcomer::routes())
        .merge(onboarding::routes())
        .merge(auth::routes());
    // We make the router without State because it isn't necessary for OpenApi generation
    // So when the main program use this function it needs to `with_state`,
    // while for Docs we can skip this.

    OpenApiRouter::with_openapi(ApiDoc::openapi())
        // With Utoipa for OpenAPI we use this macro instead of Axum syntax for route declaration
        .routes(routes!(health))
        .routes(routes!(network_info))
        .nest("/api", api_routes)
}

Basically I have this init function that generates the router WITHOUT the state, so that I can in the main binary attach the state and run the program normally, but there is a second binary that when executed just generates an openapi.json. To do that it takes the OpenApiRouter, destructure it and ignore the Axum Router. The problem is that while right now I don't need the state to define the Router, with a middleware with state I need always a State struct available.
So the solution I thought of was to allow the middleware to take the state directly from the request. In fact I can't understand why this is not possible and the middleware needs always a dedicated state instance.

I know this is a strange situation, sorry for my bad explaination skills

[jplatte]

My personal recommendation would be to use a task-local for your state sharing instead of axum's State, similar to this example (the logic for setting the task-local can just be "clone this closure-captured variable, doesn't need to be dynamic per request like there).

[eisterman]

So there is no way to use the Axum state inside a middleware taking it directly from the router instead of explicitly cloning it ahead of time and the solution need to be a workaround of sorts like using a task local state.

Super sad, but ok I guess.
Thank you very much!

[jplatte]

Yeah. This is because axum's state concept is kind of "tacked on"to the tower traits. We could make our own traits that pass state around, then this would be trivial to fix. I don't think there's any interest in that though.

By the way, task locals have some other advantages, and I wanted to remove state entirely in favor of task locals (#3330), but my co-maintainers disagreed :)