Axum Auth: jwt token, accepts either Authz bearer or session cookie

auth_controller.rs

use argon2::{Argon2, PasswordHash};
use axum_extra::extract::cookie::{Cookie, SameSite};
use axum_extra::extract::CookieJar;

async fn login_user(
    ctx: State<ApiContext>,
    Json(data): Json<LoginUser>,
) -> Result<(CookieJar, Json<LoginResponse>)> {
    let user = models::user::get_user_pass_hash(&ctx.db, &data.email)
        .await?
        .ok_or_else(|| Error::unprocessable_entity([("email", "does not exist")]))?;

    verify_password(data.password, user.password_hash).await?;
    let token = AuthUser { user_id: user.id }.to_jwt(&ctx);
    let cookie = Cookie::build(("jwt", token.clone()))
        .http_only(true)
        .secure(true)
        .same_site(SameSite::Strict);

    let jar = CookieJar::new().add(cookie);

    // Return token in both JSON and session cookie
    let response = Json(LoginResponse { token });
    Ok((jar, response))
}

async fn logout_user(jar: CookieJar) -> Result<CookieJar> {
    let jar = jar.remove(Cookie::from("jwt"));
    Ok(jar)
}

async fn verify_password(password: String, password_hash: String) -> Result<()> {
    tokio::task::spawn_blocking(move || -> Result<()> {
        let hash = PasswordHash::new(&password_hash)
            .map_err(|e| anyhow::anyhow!("invalid password hash: {}", e))?;

        hash.verify_password(&[&Argon2::default()], password)
            .map_err(|e| match e {
                argon2::password_hash::Error::Password => Error::Unauthorized,
                _ => anyhow::anyhow!("failed to verify password hash: {}", e).into(),
            })
    })
    .await
    .context("panic in verifying password hash")?
}

auth_extractor.rs

use crate::http::error::Error;
use crate::http::ApiContext;

use async_trait::async_trait;
use axum::extract::{FromRef, FromRequestParts};
use axum::http::request::Parts;
use axum_extra::headers::authorization::Bearer;
use axum_extra::headers::{Authorization, Cookie, HeaderMapExt};
use hmac::{Hmac, Mac};
use jwt::{SignWithKey, VerifyWithKey};
use sha2::Sha384;
use time::OffsetDateTime;
use uuid::Uuid;

const DEFAULT_SESSION_LENGTH: time::Duration = time::Duration::weeks(2);

/// Add this as a parameter to a handler function to require the user to be logged in.
///
/// Parses a JWT from the `Authorization: Bearer <token>` header or `jwt` cookie.
pub struct AuthUser {
    pub user_id: Uuid,
}

/// Add this as a parameter to a handler function to optionally check if the user is logged in.
///
/// If the `Authorization` header and `jwt` cookie are absent then this will be `Self(None)`,
/// otherwise it will validate the token.
///
/// This is in contrast to using `Option<AuthUser>`, which will be `None` if there
/// is *any* auth error, which isn't what we want.
pub struct MaybeAuthUser(pub Option<AuthUser>);

#[derive(serde::Serialize, serde::Deserialize)]
struct AuthUserClaims {
    user_id: Uuid,
    /// Standard JWT `exp` claim.
    exp: i64,
}

impl AuthUser {
    pub fn to_jwt(&self, ctx: &ApiContext) -> String {
        let hmac = Hmac::<Sha384>::new_from_slice(ctx.config.hmac_key.as_bytes())
            .expect("HMAC-SHA-384 can accept any key length");

        AuthUserClaims {
            user_id: self.user_id,
            exp: (OffsetDateTime::now_utc() + DEFAULT_SESSION_LENGTH).unix_timestamp(),
        }
        .sign_with_key(&hmac)
        .expect("HMAC signing should be infallible")
    }

    /// Parse & verify JWT as AuthUserClaims, then construct an AuthUser.
    fn from_authorization(ctx: &ApiContext, token: &str) -> Result<Self, Error> {
        // let token = auth_header.token();

        // SHA-384 (HS-384) as the HMAC is more difficult to brute-force
        // than SHA-256 (recommended by the JWT spec) at the cost of a slightly larger token.
        let hmac = Hmac::<Sha384>::new_from_slice(ctx.config.hmac_key.as_bytes())
            .expect("HMAC-SHA-384 can accept any key length");

        let jwt: jwt::Token<jwt::Header, AuthUserClaims, _> =
            token.verify_with_key(&hmac).map_err(|e| {
                tracing::debug!("JWT failed to verify: {}", e);
                Error::Unauthorized
            })?;

        let (_header, claims) = jwt.into();

        // JWTs are stateless, so we don't have any mechanism here to invalidate them
        // except expiration.
        //
        // Things to consider:
        // - Ensuring user ID isn't deleted/banned/deactivated
        // - Add user's password_hash to the HMAC keyring material so changing password invalidates a session
        // - Store valid tokens in a fast KV (delete to revoke keys)
        if claims.exp < OffsetDateTime::now_utc().unix_timestamp() {
            tracing::debug!("token expired");
            return Err(Error::Unauthorized);
        }

        Ok(Self {
            user_id: claims.user_id,
        })
    }
}

impl MaybeAuthUser {
    pub fn user_id(&self) -> Option<Uuid> {
        self.0.as_ref().map(|auth_user| auth_user.user_id)
    }
}

#[async_trait]
impl<S> FromRequestParts<S> for AuthUser
where
    S: Send + Sync,
    ApiContext: FromRef<S>,
{
    type Rejection = Error;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        match MaybeAuthUser::from_request_parts(parts, state).await? {
            MaybeAuthUser(Some(auth_user)) => Ok(auth_user),
            MaybeAuthUser(None) => Err(Error::Unauthorized),
        }
    }
}

#[async_trait]
impl<S> FromRequestParts<S> for MaybeAuthUser
where
    S: Send + Sync,
    ApiContext: FromRef<S>,
{
    type Rejection = Error;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let ctx: ApiContext = ApiContext::from_ref(state);

        // Get the value of the `Authorization` header, if it was sent at all.
        if let Some(auth_header) = parts.headers.typed_get::<Authorization<Bearer>>() {
            let auth_user = AuthUser::from_authorization(&ctx, auth_header.token())?;
            return Ok(Self(Some(auth_user)));
        }

        // Alternatively, check for a session cookie
        if let Some(cookie) = parts.headers.typed_get::<Cookie>() {
            if let Some(token) = cookie.get("jwt") {
                let auth_user = AuthUser::from_authorization(&ctx, token)?;
                return Ok(Self(Some(auth_user)));
            }
        }

        // None for the case where we had no Auth header or cookie to verify
        Ok(Self(None))
    }
}