ezidam/crates/ezidam/src/guards/jwt.rs

use crate::database::Database;
use crate::guards::refresh_token::get_refresh_token_from_cookie;
use crate::tokens::{
    JWT_COOKIE_NAME, JWT_DURATION_MINUTES, REFRESH_TOKEN_COOKIE_NAME, REFRESH_TOKEN_DURATION_DAYS,
};
use hash::SecretString;
use id::{KeyID, UserID};
use jwt::database::Key;
use jwt::{JwtClaims, PrivateKey};
use permissions::Permission;
use refresh_tokens::RefreshToken;
use rocket::http::Status;
use rocket::http::{Cookie, CookieJar, SameSite};
use rocket::request::Outcome;
use rocket::time::Duration;
use rocket::tokio::task;
use rocket::Request;
use settings::Settings;
use users::User;

mod admin;
mod admin_not_current;
mod user;

pub use admin::JwtAdmin;
pub use admin_not_current::JwtAdminNotCurrent;
pub use user::JwtUser;

#[derive(Debug)]
pub enum Error {
    GetDatabase,
    GetCookies,
    StartTransaction,
    GetRefreshToken(refresh_tokens::Error),
    UserNotFound,
    GetUser(users::Error),
    RevokeRefreshTokens(refresh_tokens::Error),
    MarkRefreshTokenUsed(refresh_tokens::Error),
    GetSettings(settings::Error),
    FirstAdminNotSet,
    ServerUrlNotSet,
    UnknownIp,
    SaveRefreshToken(refresh_tokens::Error),
    GetKey(jwt::Error),
    MostRecentKeyNotFound,
    MostRecentKeyRevoked,
    SignJwt(jwt::Error),
    CommitTransaction,
    Keys(jwt::Error),
    JwtParsing(jwt::Error),
    NoSigningKey,
    NonExistentKey(String),
    RevokedKey(KeyID),
    ImportKey(jwt::Error),
    JwtValidation(jwt::Error),
    BlockingTask(String),
    GetPermissions(permissions::Error),
}

pub struct SpecificUser<'a> {
    allow_first_admin: bool,
    requested_user: &'a UserID,
}

pub(super) fn get_access_token_from_cookie(request: &Request) -> Option<String> {
    request
        .cookies()
        .get(JWT_COOKIE_NAME)
        .map(|cookie| cookie.value().to_string())
}

pub async fn validate_jwt<T>(
    jwt: String,
    request: &Request<'_>,
    get_admin: Option<bool>,
    specific_user: Option<SpecificUser<'_>>,
) -> Result<Option<JwtClaims>, Outcome<T, Error>> {
    // Get database
    let db = match request.guard::<&Database>().await {
        Outcome::Success(database) => database,
        Outcome::Failure(e) => return Err(Outcome::Failure((e.0, Error::GetDatabase))),
        Outcome::Forward(f) => return Err(Outcome::Forward(f)),
    };

    let mut transaction = match db.begin().await {
        Ok(transaction) => transaction,
        Err(_e) => {
            return Err(Outcome::Failure((
                Status::InternalServerError,
                Error::StartTransaction,
            )));
        }
    };

    // Get keys
    let keys = match Key::get_all(&mut transaction, Some(false)).await {
        Ok(keys) => keys,
        Err(e) => {
            return Err(Outcome::Failure((
                Status::InternalServerError,
                Error::Keys(e),
            )))
        }
    };

    if let Some(specific_user) = specific_user {
        // Get settings
        let settings = match Settings::get(&mut transaction).await {
            Ok(settings) => settings,
            Err(e) => {
                return Err(Outcome::Failure((
                    Status::InternalServerError,
                    Error::GetSettings(e),
                )));
            }
        };

        // Get first admin
        let first_admin = match settings.first_admin() {
            Some(home_page) => UserID(home_page.to_string()),
            None => {
                return Err(Outcome::Failure((
                    Status::InternalServerError,
                    Error::FirstAdminNotSet,
                )));
            }
        };

        // If requested user is the first admin and they're not allowed
        if !specific_user.allow_first_admin && specific_user.requested_user == &first_admin {
            return Ok(None);
        }
    }

    if let Err(_e) = transaction.commit().await {
        return Err(Outcome::Failure((
            Status::InternalServerError,
            Error::CommitTransaction,
        )));
    }

    match task::spawn_blocking(move || -> Result<Option<JwtClaims>, Error> {
        // Parse jwt
        let parsed_jwt = jwt::parse(&jwt).map_err(Error::JwtParsing)?;

        // Get key id
        let jwk_id = parsed_jwt
            .header()
            .key_id
            .as_deref()
            .ok_or(Error::NoSigningKey)?;

        // Get key
        let key = keys
            .iter()
            .find(|&key| key.key_id().as_ref() == jwk_id)
            .ok_or_else(|| Error::NonExistentKey(jwk_id.into()))?;

        // If key has been revoked
        if key.is_revoked() {
            return Err(Error::RevokedKey(key.key_id().to_owned()));
        }

        // Import private key
        let private_key =
            PrivateKey::from_der(key.private_der(), key.key_id()).map_err(Error::ImportKey)?;

        // Validate jwt and get claims
        let jwt_claims = private_key
            .validate_jwt_extract_claims(&parsed_jwt)
            .map_err(Error::JwtValidation)?;

        // Is specific kind of user required?
        match get_admin {
            // Yes, need to get specific kind of user
            Some(get_admin) => {
                if jwt_claims.is_admin == get_admin {
                    Ok(Some(jwt_claims))
                } else {
                    Ok(None)
                }
            }
            // No, any user is good
            None => Ok(Some(jwt_claims)),
        }
    })
    .await
    {
        Ok(result) => match result {
            Ok(claims) => {
                // Return jwt claims
                Ok(claims)
            }
            Err(_e) => Err(Outcome::Forward(())),
        },
        Err(e) => {
            // Failed to run blocking task
            Err(Outcome::Failure((
                Status::InternalServerError,
                Error::BlockingTask(e.to_string()),
            )))
        }
    }
}

pub async fn use_refresh_token(
    refresh: String,
    request: &Request<'_>,
    get_admin: Option<bool>,
) -> Outcome<JwtClaims, Error> {
    // Get database
    let db = match request.guard::<&Database>().await {
        Outcome::Success(database) => database,
        Outcome::Failure(e) => return Outcome::Failure((e.0, Error::GetDatabase)),
        Outcome::Forward(f) => return Outcome::Forward(f),
    };

    // Get cookies
    let cookie_jar = match request.guard::<&CookieJar>().await {
        Outcome::Success(cookie_jar) => cookie_jar,
        Outcome::Failure(e) => return Outcome::Failure((e.0, Error::GetCookies)),
        Outcome::Forward(f) => return Outcome::Forward(f),
    };

    let mut transaction = match db.begin().await {
        Ok(transaction) => transaction,
        Err(_e) => {
            return Outcome::Failure((Status::InternalServerError, Error::StartTransaction));
        }
    };

    let refresh_token = match RefreshToken::get_one(&mut transaction, &refresh).await {
        Ok(refresh_token) => match refresh_token {
            Some(refresh_token) => refresh_token,
            None => return Outcome::Forward(()),
        },
        Err(e) => {
            return Outcome::Failure((Status::InternalServerError, Error::GetRefreshToken(e)));
        }
    };

    let user = match User::get_one_from_refresh_token(&mut transaction, &refresh).await {
        Ok(user) => match user {
            Some(user) => user,
            None => {
                return Outcome::Failure((Status::InternalServerError, Error::UserNotFound));
            }
        },
        Err(e) => {
            return Outcome::Failure((Status::InternalServerError, Error::GetUser(e)));
        }
    };

    // make sure that `get_admin` is respected, dont generate token for unwanted users!
    if let Some(get_admin) = get_admin {
        if user.is_admin() != get_admin {
            return Outcome::Forward(());
        }
    }

    if refresh_token.has_been_used() {
        // Revoke all tokens for user
        if let Err(e) =
            RefreshToken::revoke_all_for_user(&mut transaction, refresh_token.user()).await
        {
            return Outcome::Failure((Status::InternalServerError, Error::RevokeRefreshTokens(e)));
        }

        if let Err(_e) = transaction.commit().await {
            return Outcome::Failure((Status::InternalServerError, Error::CommitTransaction));
        }

        return Outcome::Forward(());
    }

    if refresh_token.is_revoked() {
        return Outcome::Forward(());
    }

    if refresh_token.has_expired() {
        return Outcome::Forward(());
    }

    if let Err(e) = refresh_token.use_token(&mut transaction).await {
        return Outcome::Failure((Status::InternalServerError, Error::MarkRefreshTokenUsed(e)));
    }

    // Get base url
    let settings = match Settings::get(&mut transaction).await {
        Ok(settings) => settings,
        Err(e) => {
            return Outcome::Failure((Status::InternalServerError, Error::GetSettings(e)));
        }
    };

    let home_page = match settings.url().map(String::from) {
        Some(home_page) => home_page,
        None => {
            return Outcome::Failure((Status::InternalServerError, Error::ServerUrlNotSet));
        }
    };

    // Generate refresh token
    let new_refresh_token = match task::spawn_blocking(SecretString::default).await {
        Ok(new_refresh_token) => new_refresh_token,
        Err(e) => {
            return Outcome::Failure((
                Status::InternalServerError,
                Error::BlockingTask(e.to_string()),
            ));
        }
    };

    // Attempt to get ip address
    let ip_address = match request.client_ip() {
        Some(ip) => ip.to_string(),
        None => {
            return Outcome::Failure((Status::BadRequest, Error::UnknownIp));
        }
    };

    // Insert refresh token in database
    if let Err(e) = RefreshToken::insert(
        &mut transaction,
        new_refresh_token.as_ref(),
        ip_address,
        user.id(),
        refresh_token.app(),
        REFRESH_TOKEN_DURATION_DAYS,
    )
    .await
    {
        return Outcome::Failure((Status::InternalServerError, Error::SaveRefreshToken(e)));
    }

    // Add refresh token as a cookie
    let mut cookie = Cookie::new(
        REFRESH_TOKEN_COOKIE_NAME,
        new_refresh_token.as_ref().to_string(),
    );
    cookie.set_secure(true);
    cookie.set_http_only(true);
    cookie.set_same_site(SameSite::Strict);
    cookie.set_max_age(Duration::days(REFRESH_TOKEN_DURATION_DAYS));
    cookie_jar.add(cookie);

    // Get latest key from database
    let key = match Key::get_most_recent(&mut transaction).await {
        Ok(key) => match key {
            Some(key) => key,
            None => {
                return Outcome::Failure((
                    Status::InternalServerError,
                    Error::MostRecentKeyNotFound,
                ));
            }
        },
        Err(e) => {
            return Outcome::Failure((Status::InternalServerError, Error::GetKey(e)));
        }
    };

    // Make sure key has not been revoked
    if key.is_revoked() {
        return Outcome::Failure((Status::InternalServerError, Error::MostRecentKeyRevoked));
    }

    // Import private key
    let private_key =
        match task::spawn_blocking(move || PrivateKey::from_der(key.private_der(), key.key_id()))
            .await
        {
            Ok(private_key) => match private_key {
                Ok(private_key) => private_key,
                Err(e) => {
                    return Outcome::Failure((Status::InternalServerError, Error::ImportKey(e)));
                }
            },
            Err(e) => {
                return Outcome::Failure((
                    Status::InternalServerError,
                    Error::BlockingTask(e.to_string()),
                ));
            }
        };

    // Get user roles
    let roles = match Permission::get_all(&mut transaction, Some(user.id()), None).await {
        Ok(roles) => roles
            .into_iter()
            .map(|role| role.role().to_string())
            .collect(),
        Err(e) => {
            return Outcome::Failure((Status::InternalServerError, Error::GetPermissions(e)));
        }
    };

    // Create jwt, sign and serialize
    let jwt_claims = JwtClaims::new(home_page.clone(), "ezidam", &user, roles);
    let jwt = match jwt_claims
        .clone()
        .sign_serialize(&private_key, JWT_DURATION_MINUTES)
    {
        Ok(jwt) => jwt,
        Err(e) => {
            return Outcome::Failure((Status::InternalServerError, Error::SignJwt(e)));
        }
    };

    // Add jwt as a cookie
    let mut cookie = Cookie::new(JWT_COOKIE_NAME, jwt);
    cookie.set_secure(true);
    cookie.set_http_only(true);
    cookie.set_same_site(SameSite::Strict);
    cookie.set_max_age(Duration::minutes(JWT_DURATION_MINUTES));
    cookie_jar.add(cookie);

    if let Err(_e) = transaction.commit().await {
        return Outcome::Failure((Status::InternalServerError, Error::CommitTransaction));
    }

    Outcome::Success(jwt_claims)
}

pub async fn use_access_token_or_refresh_token(
    request: &Request<'_>,
    get_admin: Option<bool>,
    specific_user: Option<SpecificUser<'_>>,
) -> Outcome<JwtClaims, Error> {
    let access_token = get_access_token_from_cookie(request);
    let refresh_token = get_refresh_token_from_cookie(request);

    match (access_token, refresh_token) {
        (Some(access), _) => {
            // If there is an access token, attempt to validate it
            match validate_jwt(access, request, get_admin, specific_user).await {
                Ok(jwt_claims) => match jwt_claims {
                    Some(jwt_claims) => Outcome::Success(jwt_claims),
                    None => Outcome::Forward(()),
                },
                Err(e) => e,
            }
        }
        (None, Some(refresh)) => {
            // Use refresh token to get new access token
            use_refresh_token(refresh, request, get_admin).await
        }
        (None, None) => {
            // Nothing to do
            Outcome::Forward(())
        }
    }
}