Unify middleware & extension into extractor

[dsaghliani]

In most, if not all, cases where you protect a route, you also pass it the user object. Currently, that's split into two parts:

1. wrap the route(s) with the authorization middleware;
2. request the user object in the route handler via an Extension.

The use of the extension system introduces the possibility of runtime panics, and while that's by no means a dealbreaker, I think there's an opportunity for improvement by unifying the two concepts into an extractor.

Possible Implementation

// axum-login:
pub struct User<T, Role = ()>(T, PhantomData<Role>);
/// This could also be an enum. It depends.
pub struct UserRejection;

// Consumer crate:
fn protected_handler(User(user, _): User<MyUserType>) -> impl IntoResponse { /* ... */ }
fn admin_handler(User(user, _): User<MyUserType, Admin>) -> impl IntoResponse { /* ... */ }
fn let_me_handle_it(user: Result<User<MyUserType>, UserRejection>) -> impl IntoResponse { /* .. */ }

This could potentially be made more ergonomic by moving the role into the consumer's user type.

fn route_handler(User(user): User<MyUserType<Admin>>) -> impl IntoResponse { /* ... */ }

Rejection Behavior

By default, the extractor will return 401: Unauthorized just as the middleware currently does. A redirect URL can't be included in the type signature (nor should it), but that can be solved easily enough.

Solution 1: Rely On axum-extra's WithRejection

This is perhaps the purest solution, but I'm not a fan because it would have to appear in every route handler and really bloat the code.

Solution 2: Include A LoginRedirect Extension

Consumers of axum-login can add an extension to the router with the redirection information. (This can be extended to a completely custom rejection, if necessary.) The extractor can then try to grab that and use it if it exists or fall back to the default behavior.

// axum-login:
/// The `&str` is used for the sake of example, and this wouldn't compile as it is,
/// but it should communicate the gist of things.
pub struct LoginRedirect(&str);

// Consumer crate:
Router::new()
    // ...
    .layer(Extension(Arc::new(LoginRedirect("/login")))

Concerns

Protected Routes That Don't Need The User

To be honest, I can't think of any cases where you'd want to protect a route but not need any user information. Still, there are bound to be some. This entire proposal is predicated on the assumption that these cases are rare.

The occasional exception can be solved simply by requesting but not using the user object (_: User<MyUser>), but if this is more common than expected, it might make sense to leave the authorization middleware so it can cover a swathe of routes.

What About Role Ordering?

The thorniest part of this proposal is that by breaking up roles from an enum into multiple structs, we can no longer use PartialOrd. This can be mitigated by introducing a trait for priority ranking that would simply return a number. Of course, implementing this trait by hand is error-prone and absolutely a no-go, but that can be solved by a macro.

define_roles! {
   Regular,
   Admin
}

// Expands to:
struct Regular;

impl RolePriority for Regular {
   fn priority(&self) -> u8 {
      0
   }
}

struct Admin;

impl RolePriority for Admin {
   fn priority(&self) -> u8 {
      1
   }
}

In appearance and effect, this is equivalent to the PartialOrd enum implementation, but there is certainly an appeal in keeping things macro-free. It's more elegant. However, I think a macro as small and simple as this is worth the elegance and expressiveness of an extractor implementation.

What do you folks think?

[maxcountryman]

Thank you for this wonderful and detailed write up!

I am very much in favor of this and agree with your assessment.

[maxcountryman]

Another detail I'm curious about here: how should user stores be handled?

[dsaghliani]

I haven't thought too hard about that, but I think it could be converted into an extension?

It's not perfect because that could still panic, but at least it only has to be done once (unlike using extensions in route handlers, which could panic for every route not covered by the middleware) and would be caught in development the first time the User AuthContext extractor is hit.

We could also define and require users to implement a trait with functions for accessing the store on the state struct. It's probably more performant to clone a field than to get it from the extensions hashmap.

I'm still blurry on the details of axum-login's implementation, though, so I'll take a good look tomorrow.

[dsaghliani]

OK, so a deep dive later, I've come up with something resembling a preliminary game plan. First, though, let's make sure we're on the same page. If I understood it correctly, this is how axum-login currently works:

1. AuthService gets the current user from the session and adds the user object to request extensions.
   • Getting the user is delegated to AuthContext, which also invalidates the session if it's stale (i.e. the password was changed).
   • The authorization middleware (Login) also inserts the user into the request extensions. What's the reason behind the duplication?
2. The authorization middleware grabs the user from the extensions and makes sure it satisfies the given role bounds. Otherwise, it rejects the request.
3. Route handlers can access the user object by requesting the extension.

Game Plan

Steps 2 and 3 can be merged into the User extractor. That can be the MVP, and the auth. layer & service can be left alone.

If and when that's done, we can move on to AuthContext. In all honesty, I don't like the name. "Context" is vague.

• IMO, AuthContext should be renamed to the more expressive Authenticator. Also, it doesn't make sense for the authenticator to do the job of getting the current user from the session—its responsibilities should be only to log users in and out—so get_user() should be moved into AuthService (or a pure function in the same module).
• Currently, the FromRequestParts implementation for AuthContext is effectively syntactic sugar for Extension<AuthContext>. To me, that feels awkward. We can probably leave it as is at first and make changes one step at a time, but afterward, maybe we'll figure out a better way.

Keep in mind that this is just early speculation that could turn out fantastically wrong when it's time to actually implement it. Please, let me know your input on this. In the meantime, if you don't mind, I'll take a crack at this and see if the theory holds up in practice.

[dsaghliani]

Update: Implementing role checking has proven troublesome.

I said the extractor could be made more ergonomic by moving the Role generic parameter into the user type, but actually, that's not possible because Rust requires that:

• it appears in the implementing type of the impl, e.g. impl<T> Foo<T>;
• for a trait impl, it appears in the implemented trait, e.g.
  impl<T> SomeTrait<T> for Foo;
• it is bound as an associated type, e.g. impl<T, U> SomeTrait for T where T: AnotherTrait<AssocType=U>.

The role can't be an associated type, so even if it were turned into a generic parameter of the user type, it would still have to appear in the extractor. Well, fine.

The next problem is that I couldn't figure out how to specify that the extractor's Role is different from whatever the user has. impl<R: Role, User: AuthUser<R>> FromRequestParts for UserExtractor<User, R> means that the extractor and the user have the same exact type to begin with, which isn't what we want. This approach was doomed from the start.

Then let's use dynamic dispatch. But there's still an issue: the priority method was defined as fn priority() -> u8 (notice the lack of &self) and used as Role::priority(), but now Role isn't object-safe and can't be made into a trait object, so we have to create two traits, one with priority() and another with priority(&self), both returning the same value.

So long as we define the roles with a macro, this should work, but I'm still unhappy with it. I'll sleep on it and see if there's a better implementation (or maybe a better API). Input is welcome.

Here's the current implementation:

use std::marker::PhantomData;

pub trait Role {
    fn priority(&self) -> u8;
}

pub trait ObjectUnsafeRole {
    fn priority() -> u8;
}

pub trait Authorize {
    fn role(&self) -> &dyn Role;
}

/// This is a convenience trait to simulate getting the user from the request's extensions.
#[cfg_attr(test, mockall::automock)]
pub trait Extensions {
    /// mockall requires the `'static` lifetime.
    fn get<T: 'static>(&self) -> T;
}

pub trait FromRequestParts {
    fn from_request_parts(parts: impl Extensions) -> Self;
}

pub struct Extractor<T, Role>(T, PhantomData<Role>);

impl<User, R> FromRequestParts for Extractor<User, R>
where
    User: Authorize + 'static,
    R: ObjectUnsafeRole,
{
    fn from_request_parts(parts: impl Extensions) -> Self {
        let user = parts.get::<User>();

        if user.role().priority() >= R::priority() {
            Self(user, PhantomData)
        } else {
            panic!("the real thing would return an error instead");
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{Authorize, Extractor, FromRequestParts, MockExtensions, ObjectUnsafeRole, Role};

    struct Regular;

    impl Role for Regular {
        fn priority(&self) -> u8 {
            1
        }
    }

    impl ObjectUnsafeRole for Regular {
        fn priority() -> u8 {
            1
        }
    }

    struct Admin;

    impl Role for Admin {
        fn priority(&self) -> u8 {
            2
        }
    }

    impl ObjectUnsafeRole for Admin {
        fn priority() -> u8 {
            2
        }
    }

    struct User {
        role: Box<dyn Role>,
    }

    impl Authorize for User {
        fn role(&self) -> &dyn Role {
            self.role.as_ref()
        }
    }

    #[test]
    fn regular_allowed_in_regular() {
        let mut parts = MockExtensions::new();

        parts.expect_get().returning(|| User {
            role: Box::new(Regular),
        });

        let _: Extractor<User, Regular> = Extractor::from_request_parts(parts);
    }

    #[test]
    #[should_panic]
    fn regular_not_allowed_in_admin() {
        let mut parts = MockExtensions::new();

        parts.expect_get().returning(|| User {
            role: Box::new(Regular),
        });

        let _: Extractor<User, Admin> = Extractor::from_request_parts(parts);
    }

    #[test]
    fn admin_allowed_in_regular() {
        let mut parts = MockExtensions::new();

        parts.expect_get().returning(|| User {
            role: Box::new(Admin),
        });

        let _: Extractor<User, Regular> = Extractor::from_request_parts(parts);
    }

    #[test]
    fn admin_allowed_in_admin() {
        let mut parts = MockExtensions::new();

        parts.expect_get().returning(|| User {
            role: Box::new(Admin),
        });

        let _: Extractor<User, Admin> = Extractor::from_request_parts(parts);
    }
}

[dsaghliani]

After some thought, I believe I've figured out what shape the API can take.

The role type cannot be moved into the user type. The simplest possible form for an extractor is Extractor<T>(T), so it can only ever wrap T. This means that routes with the parameter, User<MyUser<Regular>>, can't receive MyUser<Admin>. Fundamentally, the user's role must be dynamic; i.e., a trait object. Our user extractor can only be defined as User<T, Role>(T, PhantomData<Role>) and used as, for example, User(user, _): User<MyUser, Admin>. Though, perhaps, Authorized is a better name for the extractor, leaving "User" available.

Because we know the type of the lowest possible role that may hit a route but don't have an instance of it, the priority function must be defined as fn priority() -> u8 and not fn priority(&self) -> u8. This makes it impossible to call on a trait object, so we need two traits, which I've so far been calling Role and ObjectUnsafeRole, that effectively do the same thing.

The two priority function must return the same value for the same type, but while hacky, a macro will take care of this. The traits will have to be public since they need to be implemented for the consumer's types, but the macro will hide them. We can even use the #[doc(hidden)] attribute to minimize the odds of people stumbling upon them.

Regarding AuthContext::get_user

I'm still a bit blurry on this part of the code, but as I see it, the get_user method performs user authentication. However, requests to unprotected routes don't need that, right? If I'm right, then we can move that out of the AuthContext (and therefore AuthService, which uses that method) and into the extractor so that it's only executed when actually necessary. Theoretically, this should improve performance.

---

What do you think of this, @maxcountryman? Do you think this is a direction axum-login should take? Let me know, and I can get started and submit a PR.

[maxcountryman]

Thanks for taking the time to dig into this. I don't think I'll be able to provide good direction here until I've resolved maxcountryman/axum-sessions#56.

I'm unsure how much this will change axum-login or if it will change it at all. One direction we could go is to create a separate crate, tower-login to complement tower-sessions. Alternatively we could migrate axum-login to use this new crate (assuming it ends up replacing axum-sessions).

While I think the exact implementations of sessions is potentially separate from this, there's also the possibility that tower-sessions might influence this design.

[maxcountryman]

@dsaghliani maxcountryman/tower-sessions#2 might be relevant here, specifically this idea of "buckets".

[maxcountryman]

Here's an extractor-only approach: #100 (comment)

[maxcountryman]

@dsaghliani we've redesigned the crate around tower-sessions--let me know what you think about this given that new design.