web/

jwt_auth.rs

use std::sync::Arc;

use axum::{
    Extension, Json,
    body::Body,
    extract::State,
    http::{HeaderValue, Request, StatusCode, header},
    middleware::Next,
    response::{IntoResponse, Response},
};

use axum_extra::extract::{
    CookieJar,
    cookie::{Cookie, SameSite},
};
use serde::{Deserialize, Serialize};

use crate::{AppState, model::User};
use redis::AsyncCommands;
use server::command::user::{Account, Refusal};

#[derive(Debug, Serialize)]
pub struct ErrorResponse {
    pub status: &'static str,
    pub message: String,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct JWTAuthMiddleware {
    pub user: User,
    pub access_token_uuid: uuid::Uuid,
    /// This request arrived with a dead access token and the middleware
    /// refreshed the session itself, consuming the refresh token in the jar.
    pub refreshed: bool,
}

type AuthError = (StatusCode, Json<ErrorResponse>);

fn unauthorized(message: impl Into<String>) -> AuthError {
    (
        StatusCode::UNAUTHORIZED,
        Json(ErrorResponse {
            status: "fail",
            message: message.into(),
        }),
    )
}

/// A tenant whose database is not at the current schema.
///
/// 503 rather than 500: the account is intact and will come back, and a
/// retryable status is what a caller (and an operator reading logs) needs to
/// see. Checked here because this middleware already re-reads the user row on
/// every authenticated request, so it costs one column rather than a query,
/// and it covers the whole authenticated surface instead of each handler
/// remembering to.
fn unavailable(message: impl Into<String>) -> AuthError {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        Json(ErrorResponse {
            status: "error",
            message: message.into(),
        }),
    )
}

fn internal(message: impl Into<String>) -> AuthError {
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        Json(ErrorResponse {
            status: "error",
            message: message.into(),
        }),
    )
}

fn extract_access_token(cookie_jar: &CookieJar, req: &Request<Body>) -> Option<String> {
    cookie_jar
        .get("access_token")
        .map(|cookie| cookie.value().to_string())
        .or_else(|| {
            req.headers()
                .get(header::AUTHORIZATION)
                .and_then(|auth_header| auth_header.to_str().ok())
                .and_then(|auth_value| {
                    auth_value
                        .strip_prefix("Bearer ")
                        .map(std::borrow::ToOwned::to_owned)
                })
        })
}

/// Why an access token did or did not resolve to a session.
enum Access {
    Valid(uuid::Uuid, uuid::Uuid),
    /// No usable token. The caller may still refresh.
    Absent,
    /// The session store could not be reached. Distinguished from `Absent`
    /// because falling through to the refresh path told the user their session
    /// had expired — and consumed their refresh token — over an outage that
    /// had nothing to do with them.
    Unavailable,
}

async fn resolve_access(data: &Arc<AppState>, access_token: Option<&str>) -> Access {
    let Some(token_str) = access_token else {
        return Access::Absent;
    };
    let details = match crate::auth_keys::verify(token_str, crate::token::TokenType::Access).await {
        Ok(Some(details)) => details,
        Ok(None) => return Access::Absent,
        Err(_) => return Access::Unavailable,
    };
    let Ok(mut redis) = data.redis_client.get_multiplexed_async_connection().await else {
        return Access::Unavailable;
    };
    // `Option<String>`, so a missing session reads as `Ok(None)` rather than
    // sharing an error arm with the store being down.
    let Ok(user_id_str) = redis
        .get::<_, Option<String>>(details.token_uuid.to_string())
        .await
    else {
        return Access::Unavailable;
    };
    let Some(user_id) = user_id_str.and_then(|id| uuid::Uuid::parse_str(&id).ok()) else {
        return Access::Absent;
    };
    // Bind the verified JWT subject to the Redis session owner: the token was
    // signed by `details.user_id`'s key, so a mismatch means the token_uuid was
    // paired with a different user's session — reject it.
    if user_id != details.user_id {
        return Access::Absent;
    }
    Access::Valid(user_id, details.token_uuid)
}

struct RefreshOutcome {
    account: Account,
    access_token_uuid: uuid::Uuid,
    set_cookies: Vec<HeaderValue>,
}

/// What one refresh-token rotation writes.
pub(crate) struct Rotation {
    /// The session being redeemed.
    pub token_uuid: uuid::Uuid,
    /// The user the caller's signed token claims to be.
    pub owner: uuid::Uuid,
    /// The refresh session replacing it.
    pub successor: uuid::Uuid,
    /// The access session issued with it.
    pub access_uuid: uuid::Uuid,
    /// The access session the caller held before rotating, if it still has
    /// one: superseded now, and deleted so a later logout, which only sees the
    /// newest tokens, does not leave it valid.
    pub superseded: Option<uuid::Uuid>,
    pub refresh_ttl: u64,
    pub access_ttl: u64,
}

/// Rotates a refresh session in one step, returning whether it was consumed.
///
/// The replacements are written before the old session is deleted, and a
/// failing `redis.call` aborts the script, so either the new sessions are live
/// or the old one still is. A token whose subject is not the session's owner
/// consumes nothing.
pub(crate) async fn rotate_refresh(
    redis: &mut redis::aio::MultiplexedConnection,
    rotation: &Rotation,
) -> Result<bool, AuthError> {
    let script = redis::Script::new(
        r"if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end
          redis.call('SETEX', KEYS[2], ARGV[2], ARGV[1])
          redis.call('SETEX', KEYS[3], ARGV[3], ARGV[1])
          redis.call('DEL', KEYS[1])
          if KEYS[4] then redis.call('DEL', KEYS[4]) end
          return 1",
    );
    let mut invocation = script.prepare_invoke();
    invocation
        .key(rotation.token_uuid.to_string())
        .key(rotation.successor.to_string())
        .key(rotation.access_uuid.to_string());
    if let Some(superseded) = rotation.superseded {
        invocation.key(superseded.to_string());
    }
    invocation
        .arg(rotation.owner.to_string())
        .arg(rotation.refresh_ttl)
        .arg(rotation.access_ttl)
        .invoke_async::<i64>(redis)
        .await
        .map(|rotated| rotated == 1)
        .map_err(|e| unavailable_err(&e))
}

/// Fetches the account a token names, or refuses the request.
///
/// Read through `server`, which owns every database connection and releases
/// this one before the handler runs: held across `next.run` it would tie up one
/// of the admin pool's five connections for the whole request.
async fn fetch_account(user_id: uuid::Uuid) -> Result<Account, AuthError> {
    server::command::user::fetch_account(user_id)
        .await
        .map_err(|e| internal(format!("Error fetching user from database: {e}")))?
        .ok_or_else(|| unauthorized("The user belonging to this token no longer exists"))
}

/// Refuses an account that may not be served.
fn refuse(refusal: Option<Refusal>) -> Result<(), AuthError> {
    match refusal {
        Some(Refusal::Disabled) => Err(unauthorized("The user is not verified yet")),
        Some(Refusal::Quarantined) => Err(unavailable(
            "This account is undergoing maintenance and is temporarily unavailable",
        )),
        None => Ok(()),
    }
}

/// A mint that failed: quarantined and unreachable tenants are retryable.
fn mint_failed(e: crate::auth_keys::MintError) -> AuthError {
    (
        e.status(),
        Json(ErrorResponse {
            status: "error",
            message: format!("error generating token: {e}"),
        }),
    )
}

/// The refusal for a dependency that is merely unreachable.
fn unavailable_err(e: &dyn std::fmt::Display) -> AuthError {
    unavailable(format!("session store error: {e}"))
}

async fn try_refresh_session(
    data: &Arc<AppState>,
    cookie_jar: &CookieJar,
) -> Result<RefreshOutcome, AuthError> {
    let refresh_token = cookie_jar
        .get("refresh_token")
        .map(|c| c.value().to_string())
        .ok_or_else(|| unauthorized("You are not logged in, please provide token"))?;

    let refresh_details =
        match crate::auth_keys::verify(&refresh_token, crate::token::TokenType::Refresh).await {
            Ok(Some(details)) => details,
            Ok(None) => return Err(unauthorized("Token is invalid or session has expired")),
            Err(e) => return Err(unavailable_err(&e)),
        };

    let mut redis_client = data
        .redis_client
        .get_multiplexed_async_connection()
        .await
        .map_err(|e| unavailable_err(&e))?;

    // Everything that can fail runs BEFORE the rotation, which consumes the
    // session: the account read and its refusal, and both mints, which touch
    // the tenant database. A failure after the rotation would log the user
    // out; a token minted here and never registered is inert.
    let user_id = refresh_details.user_id;
    let account = fetch_account(user_id).await?;
    refuse(account.refusal)?;

    let access_details = crate::auth_keys::mint(
        user_id,
        data.conf.access_token_max_age,
        crate::token::TokenType::Access,
        refresh_details.generation,
    )
    .await
    .map_err(mint_failed)?;
    let new_refresh_details = crate::auth_keys::mint(
        user_id,
        data.conf.refresh_token_max_age,
        crate::token::TokenType::Refresh,
        refresh_details.generation,
    )
    .await
    .map_err(mint_failed)?;

    let rotated = rotate_refresh(
        &mut redis_client,
        &Rotation {
            token_uuid: refresh_details.token_uuid,
            owner: user_id,
            successor: new_refresh_details.token_uuid,
            access_uuid: access_details.token_uuid,
            superseded: None,
            refresh_ttl: (data.conf.refresh_token_max_age * 60) as u64,
            access_ttl: (data.conf.access_token_max_age * 60) as u64,
        },
    )
    .await?;
    if !rotated {
        return Err(unauthorized("Token is invalid or session has expired"));
    }

    let access_cookie = Cookie::build((
        "access_token",
        access_details.token.clone().unwrap_or_default(),
    ))
    .path("/")
    .max_age(time::Duration::minutes(data.conf.access_token_max_age))
    .same_site(SameSite::Lax)
    .secure(crate::auth_keys::secure_cookies())
    .http_only(true)
    .build();

    let refresh_cookie = Cookie::build((
        "refresh_token",
        new_refresh_details.token.unwrap_or_default(),
    ))
    .path("/")
    .max_age(time::Duration::minutes(data.conf.refresh_token_max_age))
    .same_site(SameSite::Lax)
    .secure(crate::auth_keys::secure_cookies())
    .http_only(true)
    .build();

    let logged_in_cookie = Cookie::build(("logged_in", "true"))
        .path("/")
        .max_age(time::Duration::minutes(data.conf.access_token_max_age))
        .same_site(SameSite::Lax)
        .http_only(false)
        .build();

    let mut cookie_headers = axum::http::HeaderMap::new();
    crate::handler::append_set_cookies(
        &mut cookie_headers,
        [access_cookie, refresh_cookie, logged_in_cookie].map(|c| c.to_string()),
    );
    let set_cookies = cookie_headers
        .get_all(header::SET_COOKIE)
        .iter()
        .cloned()
        .collect();

    Ok(RefreshOutcome {
        account,
        access_token_uuid: access_details.token_uuid,
        set_cookies,
    })
}

pub async fn auth(
    cookie_jar: CookieJar,
    State(data): State<Arc<AppState>>,
    mut req: Request<Body>,
    next: Next,
) -> Result<Response, AuthError> {
    let access_token = extract_access_token(&cookie_jar, &req);

    let (row, access_token_uuid, refreshed_cookies) =
        match resolve_access(&data, access_token.as_deref()).await {
            Access::Valid(uid, atid) => {
                let account = fetch_account(uid).await?;
                refuse(account.refusal)?;
                (account, atid, Vec::new())
            }
            Access::Unavailable => {
                return Err(unavailable("The service is temporarily unavailable"));
            }
            // The refresh read the account before rotating, so nothing that can
            // fail is left between consuming the session and answering.
            Access::Absent => {
                let outcome = try_refresh_session(&data, &cookie_jar).await?;
                (
                    outcome.account,
                    outcome.access_token_uuid,
                    outcome.set_cookies,
                )
            }
        };

    let user = User {
        id: row.id,
        name: row.name,
        email: row.email,
        role: row.role,
        photo: row.photo,
        verified: row.verified,
        created_at: row.created_at,
        updated_at: row.updated_at,
    };

    req.extensions_mut().insert(JWTAuthMiddleware {
        user,
        access_token_uuid,
        refreshed: !refreshed_cookies.is_empty(),
    });

    let mut response = next.run(req).await;
    // A handler that ends the session sets its own expiry cookies, and the
    // browser keeps whichever `Set-Cookie` for a name comes last. Appending a
    // refresh minted on the way in would hand the client a live session back:
    // logout does not bump `session_generation`, so unlike the password change
    // that token would still verify, and the middleware's rotated refresh UUID
    // is one the logout handler never saw and so cannot delete from Redis.
    if !ends_session(&response) {
        for cookie in refreshed_cookies {
            response.headers_mut().append(header::SET_COOKIE, cookie);
        }
    }
    Ok(response)
}

/// Whether the handler already expired the access cookie.
pub fn ends_session<B>(response: &Response<B>) -> bool {
    response
        .headers()
        .get_all(header::SET_COOKIE)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .filter_map(|raw| Cookie::parse(raw.to_owned()).ok())
        .any(|cookie| cookie.name() == "access_token" && cookie.value().is_empty())
}

pub async fn admin(
    Extension(jwtauth): Extension<JWTAuthMiddleware>,
    req: Request<Body>,
    next: Next,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
    if jwtauth.user.role != "admin" {
        let error_response = ErrorResponse {
            status: "fail",
            message: "The user does not have admin permissions".to_string(),
        };
        return Err((StatusCode::UNAUTHORIZED, Json(error_response)));
    }

    Ok(next.run(req).await)
}

/// The rotation script against a real session store: it is all that stands
/// between a refresh and a logged-out user, so it is proved rather than argued.
#[cfg(test)]
mod rotation_tests {
    use super::{Rotation, rotate_refresh};
    use redis::AsyncCommands;

    async fn store() -> redis::aio::MultiplexedConnection {
        let url =
            std::env::var("REDIS_URL").unwrap_or_else(|_| "redis://127.0.0.1:6379".to_string());
        redis::Client::open(url)
            .expect("a valid session-store URL")
            .get_multiplexed_async_connection()
            .await
            .expect("the session store must be running for these tests")
    }

    async fn get(redis: &mut redis::aio::MultiplexedConnection, key: uuid::Uuid) -> Option<String> {
        redis.get(key.to_string()).await.expect("get")
    }

    /// A live session owned by `owner`, and the rotation that would replace it.
    async fn seeded(redis: &mut redis::aio::MultiplexedConnection, owner: uuid::Uuid) -> Rotation {
        let rotation = Rotation {
            token_uuid: uuid::Uuid::new_v4(),
            owner,
            successor: uuid::Uuid::new_v4(),
            access_uuid: uuid::Uuid::new_v4(),
            superseded: None,
            refresh_ttl: 60,
            access_ttl: 60,
        };
        let _: () = redis
            .set_ex(rotation.token_uuid.to_string(), owner.to_string(), 60)
            .await
            .expect("seed");
        rotation
    }

    #[tokio::test]
    async fn a_rotation_replaces_the_session_once() {
        let mut redis = store().await;
        let owner = uuid::Uuid::new_v4();
        let rotation = seeded(&mut redis, owner).await;

        assert!(rotate_refresh(&mut redis, &rotation).await.expect("rotate"));
        assert_eq!(get(&mut redis, rotation.token_uuid).await, None);
        assert_eq!(
            get(&mut redis, rotation.successor).await,
            Some(owner.to_string())
        );
        assert_eq!(
            get(&mut redis, rotation.access_uuid).await,
            Some(owner.to_string())
        );

        assert!(
            !rotate_refresh(&mut redis, &rotation).await.expect("rotate"),
            "a refresh token is single-use"
        );
    }

    /// Refreshing while still holding an access token supersedes it. Logout
    /// only ever sees the newest pair, so a superseded token left alive would
    /// outlive the session.
    #[tokio::test]
    async fn a_rotation_ends_the_access_token_it_supersedes() {
        let mut redis = store().await;
        let owner = uuid::Uuid::new_v4();
        let previous = uuid::Uuid::new_v4();
        let _: () = redis
            .set_ex(previous.to_string(), owner.to_string(), 60)
            .await
            .expect("seed previous access");
        let rotation = Rotation {
            superseded: Some(previous),
            ..seeded(&mut redis, owner).await
        };

        assert!(rotate_refresh(&mut redis, &rotation).await.expect("rotate"));
        assert_eq!(get(&mut redis, previous).await, None);
        assert_eq!(
            get(&mut redis, rotation.access_uuid).await,
            Some(owner.to_string())
        );
    }

    #[tokio::test]
    async fn a_foreign_owner_consumes_nothing() {
        let mut redis = store().await;
        let owner = uuid::Uuid::new_v4();
        let rotation = seeded(&mut redis, owner).await;
        let (session, successor) = (rotation.token_uuid, rotation.successor);

        let impostor = Rotation {
            owner: uuid::Uuid::new_v4(),
            ..rotation
        };
        assert!(!rotate_refresh(&mut redis, &impostor).await.expect("rotate"));
        assert_eq!(get(&mut redis, session).await, Some(owner.to_string()));
        assert_eq!(get(&mut redis, successor).await, None);
    }

    /// A failing `redis.call` aborts the script with earlier writes applied, so
    /// the delete must come last or a bad TTL consumes the session.
    #[tokio::test]
    async fn a_rotation_that_cannot_write_leaves_the_session_intact() {
        let mut redis = store().await;
        let owner = uuid::Uuid::new_v4();
        let rotation = seeded(&mut redis, owner).await;
        let session = rotation.token_uuid;

        let broken = Rotation {
            refresh_ttl: 0,
            ..rotation
        };
        assert!(rotate_refresh(&mut redis, &broken).await.is_err());
        assert_eq!(get(&mut redis, session).await, Some(owner.to_string()));
    }
}