Skip to content

Instantly share code, notes, and snippets.

@anowell
Last active September 29, 2024 01:31
Show Gist options
  • Save anowell/7cb7cb04acfeffeddce181408157cacb to your computer and use it in GitHub Desktop.
Save anowell/7cb7cb04acfeffeddce181408157cacb to your computer and use it in GitHub Desktop.
Axum Auth: jwt token, accepts either Authz bearer or session cookie
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")?
}
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))
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment