bazzar/actors/token_manager/src/lib.rs

use std::collections::BTreeMap;
use std::str::FromStr;

use actix::{Addr, Message};
use chrono::prelude::*;
use config::SharedAppConfig;
use database_manager::{query_db, Database};
use hmac::digest::KeyInit;
use hmac::Hmac;
use model::{AccountId, Audience, Role, Token, TokenString};
use sha2::Sha256;

#[macro_export]
macro_rules! token_async_handler {
    ($msg: ty, $async: ident, $res: ty) => {
        impl actix::Handler<$msg> for TokenManager {
            type Result = actix::ResponseActFuture<Self, Result<$res>>;

            fn handle(&mut self, msg: $msg, _ctx: &mut Self::Context) -> Self::Result {
                use actix::WrapFuture;
                let db = self.db.clone();
                let config = self.config.clone();
                Box::pin(async { $async(msg, db, config).await }.into_actor(self))
            }
        }
    };
}

/*struct Jwt {
    /// cti (customer id): Customer uuid identifier used by payment service
    pub cti: uuid::Uuid,
    /// arl (account role): account role
    pub arl: Role,
    /// iss (issuer): Issuer of the JWT
    pub iss: String,
    /// sub (subject): Subject of the JWT (the user)
    pub sub: i32,
    /// aud (audience): Recipient for which the JWT is intended
    pub aud: Audience,
    /// exp (expiration time): Time after which the JWT expires
    pub exp: chrono::NaiveDateTime,
    /// nbt (not before time): Time before which the JWT must not be accepted
    /// for processing
    pub nbt: chrono::NaiveDateTime,
    /// iat (issued at time): Time at which the JWT was issued; can be used to
    /// determine age of the JWT,
    pub iat: chrono::NaiveDateTime,
    /// jti (JWT ID): Unique identifier; can be used to prevent the JWT from
    /// being replayed (allows a token to be used only once)
    pub jti: uuid::Uuid,
}*/

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Unable to save new token")]
    Save,
    #[error("Unable to save new token. Can't connect to database")]
    SaveInternal,
    #[error("Unable to validate token")]
    Validate,
    #[error("Unable to validate token. Can't connect to database")]
    ValidateInternal,
}

pub type Result<T> = std::result::Result<T, Error>;

pub struct TokenManager {
    db: Addr<Database>,
    config: SharedAppConfig,
}

impl actix::Actor for TokenManager {
    type Context = actix::Context<Self>;
}

impl TokenManager {
    pub fn new(config: SharedAppConfig, db: Addr<Database>) -> Self {
        Self { db, config }
    }
}

#[derive(Message)]
#[rtype(result = "Result<(Token, TokenString)>")]
pub struct CreateToken {
    pub customer_id: uuid::Uuid,
    pub role: Role,
    pub subject: AccountId,
    pub audience: Option<Audience>,
}

token_async_handler!(CreateToken, create_token, (Token, TokenString));

pub(crate) async fn create_token(
    msg: CreateToken,
    db: Addr<Database>,
    config: SharedAppConfig,
) -> Result<(Token, TokenString)> {
    let CreateToken {
        customer_id,
        role,
        subject,
        audience,
    } = msg;
    let audience = audience.unwrap_or_default();

    let token: Token = query_db!(
        db,
        database_manager::CreateToken {
            customer_id,
            role,
            subject,
            audience,
        },
        Error::Save,
        Error::SaveInternal
    );

    let token_string = {
        use jwt::SignWithKey;

        let secret = config.lock().web().jwt_secret();
        let key: Hmac<Sha256> = build_key(secret)?;
        let mut claims = BTreeMap::new();

        // cti (customer id): Customer uuid identifier used by payment service
        claims.insert("cti", format!("{}", token.customer_id));
        // arl (account role): account role
        claims.insert("arl", String::from(token.role.as_str()));
        // iss (issuer): Issuer of the JWT
        claims.insert("iss", token.issuer.to_string());
        // sub (subject): Subject of the JWT (the user)
        claims.insert("sub", format!("{}", token.subject));
        // aud (audience): Recipient for which the JWT is intended
        claims.insert("aud", String::from(token.audience.as_str()));
        // exp (expiration time): Time after which the JWT expires
        claims.insert(
            "exp",
            format!(
                "{}",
                Utc.from_utc_datetime(&token.expiration_time).format("%+")
            ),
        );
        // nbt (not before time): Time before which the JWT must not be accepted
        // for processing
        claims.insert(
            "nbt",
            format!(
                "{}",
                Utc.from_utc_datetime(&token.not_before_time).format("%+")
            ),
        );
        // iat (issued at time): Time at which the JWT was issued; can be used
        // to determine age of the JWT,
        claims.insert(
            "iat",
            format!(
                "{}",
                Utc.from_utc_datetime(&token.issued_at_time).format("%+")
            ),
        );
        // jti (JWT ID): Unique identifier; can be used to prevent the JWT from
        // being replayed (allows a token to be used only once)
        claims.insert("jti", format!("{}", token.jwt_id));

        let s = match claims.sign_with_key(&key) {
            Ok(s) => s,
            Err(e) => {
                log::error!("{e:?}");
                return Err(Error::SaveInternal);
            }
        };
        TokenString::from(s)
    };
    Ok((token, token_string))
}

#[derive(Message)]
#[rtype(result = "Result<(Token, bool)>")]
pub struct Validate {
    pub token: TokenString,
}

token_async_handler!(Validate, validate, (Token, bool));

pub(crate) async fn validate(
    msg: Validate,
    db: Addr<Database>,
    config: SharedAppConfig,
) -> Result<(Token, bool)> {
    use jwt::VerifyWithKey;

    log::info!("Validating token {:?}", msg.token);

    let secret = config.lock().web().jwt_secret();
    let key: Hmac<Sha256> = build_key(secret)?;
    let claims: BTreeMap<String, String> = match msg.token.verify_with_key(&key) {
        Ok(claims) => claims,
        _ => return Err(Error::Validate),
    };
    let jti = match claims.get("jti") {
        Some(jti) => jti,
        _ => return Err(Error::Validate),
    };

    let token: Token = query_db!(
        db,
        database_manager::TokenByJti {
            jti: match uuid::Uuid::from_str(jti) {
                Ok(uid) => uid,
                _ => return Err(Error::Validate),
            },
        },
        Error::Validate,
        Error::ValidateInternal
    );

    if !validate_pair(&claims, "cti", token.customer_id, validate_uuid) {
        return Ok((token, false));
    }
    if !validate_pair(&claims, "arl", token.role, |left, right| right == left) {
        return Ok((token, false));
    }
    if !validate_pair(&claims, "iss", &token.issuer, |left, right| right == left) {
        return Ok((token, false));
    }
    if !validate_pair(&claims, "sub", token.subject, validate_num) {
        return Ok((token, false));
    }
    if !validate_pair(&claims, "aud", token.audience, |left, right| right == left) {
        return Ok((token, false));
    }
    if !validate_pair(&claims, "exp", &token.expiration_time, validate_time) {
        return Ok((token, false));
    }
    if !validate_pair(&claims, "nbt", &token.not_before_time, validate_time) {
        return Ok((token, false));
    }
    if !validate_pair(&claims, "iat", &token.issued_at_time, validate_time) {
        return Ok((token, false));
    }

    log::info!("JWT token valid");
    Ok((token, true))
}

fn build_key(secret: String) -> Result<Hmac<Sha256>> {
    match Hmac::new_from_slice(secret.as_bytes()) {
        Ok(key) => Ok(key),
        Err(e) => {
            log::error!("{e:?}");
            Err(Error::ValidateInternal)
        }
    }
}

fn validate_pair<F, V>(claims: &BTreeMap<String, String>, key: &str, v: V, cmp: F) -> bool
where
    F: FnOnce(&str, V) -> bool,
    V: PartialEq,
{
    claims.get(key).map(|s| cmp(s, v)).unwrap_or_default()
}

fn validate_time(left: &str, right: &NaiveDateTime) -> bool {
    chrono::DateTime::parse_from_str(left, "%+")
        .map(|t| t.naive_utc() == *right)
        .unwrap_or_default()
}

fn validate_num(left: &str, right: i32) -> bool {
    left.parse::<i32>().map(|n| n == right).unwrap_or_default()
}

fn validate_uuid(left: &str, right: uuid::Uuid) -> bool {
    uuid::Uuid::from_str(left)
        .map(|u| u == right)
        .unwrap_or_default()
}