Skip to content

Instantly share code, notes, and snippets.

@ceejbot
Last active September 17, 2024 21:45
Show Gist options
  • Save ceejbot/cd61b42aeb865364a36cd811f08e09f7 to your computer and use it in GitHub Desktop.
Save ceejbot/cd61b42aeb865364a36cd811f08e09f7 to your computer and use it in GitHub Desktop.
Middleware, onions, and crying

Notes on axum middleware

A few days ago I complained about the agony of trying to write per-endpoint middleware in Axum using Tower services and layers, and how awful I found the experience every time Axum exposes Tower to me. (I still feel that way.) Josh Triplett was kind enough to point me to a relative newcomer in the Rust http framework scene, Trillium, as a good alternative.

Here are my notes on the spike I did with Trillium.


tl;dr Axum is much more popular and also finished. You are getting a more well-known and understood project when you choose it. You should probably choose it unless you have a reason not to. For now.


I did this spike because what I wanted to do with Axum-- write per-endpoint authorization middleware-- felt difficult. Every time I had to interact with Tower while using Axum was miserable, so this integration felt like a warning signal and not a selling point after a while. However, my normal use of axum was smooth sailing until that need arose. Eventually I figured out that by ignoring the documentation's instructions not use lots of layers, I could make lots of layers and get what was effectively per-endpoint middleware with the signature of ordinary layer middleware.

Tower layers confusion

The problem space: I am transcribing/rewriting a slice of an existing monolithic service, as a proof of concept. One of the things that needs to be replicated is enforcement of the existing overly-complex permission system, which has (among other things) a way to specify per-endpoint authorization requirements for the logged-in user. These are noun-verb pairs like resource + upload or project + admin. There's a fairly long existing list of resources to cope with, and the same for verbs. This permissions system is not particularly unusual.

I could see right away in the Axum documentation that there was a huge difference in support between layer middleware and per-endpoint middleware. The moment Axum dips into Tower and Services etc is the moment I stop enjoying its developer experience, because it goes lower level into the building blocks that someone uses to make an http framework, not something that (IMO) should be exposed in the framework. But! I was not deterred. I like Axum generally and was willing to learn what I had to.

My first guess about implementation: per-endpoint permissions middleware that takes two enums to configure. I assumed I would write something like this, after implementing Authorizer:

let widget_list_authed = widget_list.layer(Authorizer::new(Noun::Widget, Verb::Read));
let widget_by_id_authed = widget_by_id.layer(Authorizer::new(Noun::Widget, Verb::Read));
let widget_create_authed = widget_create.layer(Authorizer::new(Noun::Widget, Verb::Write));
// etc, only in practice a little more varied than this
Router::new()
    .route("/widgets", get(widget_list_authed))
    .route("/widgets/:id", get(widget_by_id_authed))
    .route("/widgets", put(widget_create_authed))
    // etc
    .layer(axum::middleware::from_fn_with_state(
        Arc::new(state.clone()),
        check_authentication,
    ))

Now for Authorizer. It has to retrieve a user object from request extensions and compare the required permissions to the ones the user has. Easy-peasy. My first stop was the ConcurrencyLimit source linked from the docs, to see what the example was doing. I ended up at layer_fn(), which is clearly what I need.

So in the place where I build my router:

// Figure out a type for S from whatever type it is the router has... for each endpoint?
// Or I'm making these on the fly somewhere that I have something bound to `service` already,
// but I don't have one of those when I'm building a router, I think?
// I'm not sure this saves me much over the approach I ended up with.
let authz_layer = layer_fn(|service: S, resource: Widget, permission: Permission| Authorizer {
    service,
    resource,
    permission,
});

And in the middleware implementation:

//. A middleware that checks that the logged-in user has the requested permissions.
pub struct Authorizer<S> {
    pub resource: Resource,
    pub permission: Permission,
    pub service: S,
}

impl<S, Request> Service<Request, Future = Result<_, StatusCode>> for Authorizer<S>
where
    S: Service<Request, Future = Result<_, StatusCode>>, // this type is incomplete
{
    type Response = S::Response; // okay so far
    type Error = S::Error;
    type Future = S::Future;

    // stealing boilerplate code here
    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        // The desired logic is:
        // get user from request
        // respond 401 if not found
        // check permissions
        // respond 401 if we fail check
        // run next in the chain if we pass

        // In a regular axum middleware, it would be as simple as:
        let Some(user): Option<&User> = request.extensions().get() else {
            // So I need to figure out how to make sure I'm getting `request` as
            // the type I expect it to be, by implementing this trait on the correct type.
            // I'll come back to this once I've figured it out...

            log::warn!("authorizer: No user found in request data; programmer error?");
            // Now I'm reading docs for a crate I didn't sign up to dive into, but here I am.
            // Here, I should be returning a future for the error instead of the error itself
            return Err(StatusCode::UNAUTHORIZED);
        };

        // `project` is a path parameter: how do I get it in a Tower layer thing? Again, this
        // means I have to figure out how to get Axum's request object here, so I don't have to
        // replicate work that Axum has already done.
        // An example would go a long way toward speeding me up.
        if !user.has_project_permission(project.as_str(), self.resource, self.permission) {
            log::info!(
                "authorizer: user does not have {} permissions for {}={project}; id={:?}; email={:?};",
                self.permission,
                self.resource,
                user.id,
                user.email
            );
            // Here, I should be returning a future for the error instead of the error itself
            return Err(StatusCode::UNAUTHORIZED);
        }

        // the success case in regular axum middleware
        // Ok(next.run(req).await)

        // the success case for Tower layers
        self.service.call(request)
    }
}

As the comment says, an example showing how to do Axum things in a Tower layer would go a long way. I did some searching, but did not find one. Then I complained on the Fediverse in frustration and got an alternative to try.

Trillium's api stays the same

As mentioned above, somebody suggested Trillium, which is quite nice. Trillium made writing the middleware trivial. My router looks like this:

Router::new()
    .get(format!("/widgets"), (read_permissions, widget_list))
    .get(format!("/widgets/:id"), (read_permissions, widget_by_id))
    .post(format!("/widgets"), (write_permissions, widget_create))
    // etc

The middleware and the endpoint handlers all have the exact same signature: take a connection, return a connection.

There were some things I had to give up to get easy middleware, however.

  • A finished, polished product.
  • A deep pre-existing ecosystem.
  • Onion-style middleware.

Onion-style middleware runs before and after the next layer down, with the endpoint handler at the core. Everything else about Trillium was just fine. I could make json API go brr in Rust with only minor adjustments from my existing code. The Trillium "take a connection, return a connection" api is lovely for almost everything. Treated as a learning spike, it was great.

I finally had to learn to write Rust macros to do some meta-programming to keep Result ? sugar in Trillium endpoint handlers while not having to write boilerplate wrappers over and over. For various reasons, my handlers have to do some special error handling, so no existing macros were suitable for my use case. So hooray, my first proc macro!

I became a little tired of writing predictable, repetitive get-and-validate code to extract typed path parameters from a connection. So I wrote a macro to extract path parameters into a struct, similar to the way Axum extracts them into a tuple. Hooray, my second proc macro!

Then I wrote a trait for Trillium's headers to extract integer and string values from them, with error handling. And then one for deserializing query strings into structs with serde_qs. At this point my eyebrow was up. This was a lot of developer experience work that will probably go away as Trillium advances and implements its own spin on all these things. And was this really the right thing to hand to Rust newcomers to try out as their intro? Ceej Hackware™ might make Ceej happy but it's probably not good enough for the team I need to support.

Then I figured it out

And then! And then I figured out how to write the per-endpoint middleware in Axum. Write N variations of the middleware, one per authZ requirement in the spike (N was small for now). Make one layer per group of endpoints that share the authz requirement. Like this, roughly:

Router::new()
    .merge(
        Router::new()
            .route("/widgets", get(widget_list))
            .route("/widgets/:id", get(widget_by_id))
            .layer(axum::middleware::from_fn_with_state(
                Arc::new(state.clone()),
                enforce_read_permissions,
            )),
    )
    .merge(
        Router::new()
            .route("/widgets", put(widget_create))
            .route("/widgets/:id", post(widget_update))
            .route("/widgets/:id", delete(widget_delete))
            .layer(axum::middleware::from_fn_with_state(
                Arc::new(state.clone()),
                enforce_write_permissions,
            )),
    )
    .layer(axum::middleware::from_fn_with_state(
        Arc::new(state.clone()),
        check_authentication,
    ))

The regular middleware has the signature I want, and this follows the general pattern of Axum setup. I eventually ripped out all app state and went with another way to get a db connection pool to endpoint handlers, because I suspect module once-cells will be a little easier to cope with than baton-passing for my audience, but that could go either way.

I felt pretty stupid when I realized I should just use lots of layers and ignore the docs. This is not perfect because I do want to emit a function that does the check but hey, I now know how to write macros to get rid of boilerplate when types defeat me, I guess.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment