MatrixGW/matrixgw_backend/src/extractors/auth_extractor.rs

use crate::app_config::AppConfig;
use crate::constants;
use crate::extractors::session_extractor::MatrixGWSession;
use crate::users::{APIToken, APITokenID, User, UserEmail};
use crate::utils::time_utils::time_secs;
use actix_remote_ip::RemoteIP;
use actix_web::dev::Payload;
use actix_web::error::ErrorPreconditionFailed;
use actix_web::{FromRequest, HttpRequest};
use anyhow::Context;
use bytes::Bytes;
use jwt_simple::common::VerificationOptions;
use jwt_simple::prelude::{Duration, HS256Key, MACLike};
use jwt_simple::reexports::serde_json;
use serde::de::DeserializeOwned;
use sha2::{Digest, Sha256};
use std::fmt::Display;
use std::net::IpAddr;
use std::str::FromStr;

#[derive(Debug, Clone)]
pub enum AuthenticatedMethod {
    /// User is authenticated using a cookie
    Cookie,
    /// User is authenticated through command line, for debugging purposes only
    Dev,
    /// User is authenticated using an API token
    Token(APIToken),
}

impl AuthenticatedMethod {
    pub fn light_str(&self) -> String {
        match self {
            AuthenticatedMethod::Cookie => "Cookie".to_string(),
            AuthenticatedMethod::Dev => "DevAuthentication".to_string(),
            AuthenticatedMethod::Token(t) => format!("Token({:?} - {})", t.id, t.base.name),
        }
    }
}

pub struct AuthExtractor {
    pub user: User,
    pub method: AuthenticatedMethod,
    pub payload: Option<Vec<u8>>,
}

impl AsRef<User> for AuthExtractor {
    fn as_ref(&self) -> &User {
        &self.user
    }
}

impl AuthExtractor {
    pub fn decode_json_body<E: DeserializeOwned + Send>(&self) -> anyhow::Result<E> {
        let payload = self
            .payload
            .as_ref()
            .context("Failed to decode request as json: missing payload!")?;
        Ok(serde_json::from_slice(payload)?)
    }
}

#[derive(Debug, Eq, PartialEq)]
pub struct MatrixJWTKID {
    pub user_email: UserEmail,
    pub id: APITokenID,
}

impl Display for MatrixJWTKID {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}#{}", self.user_email.0, self.id.0)
    }
}

impl FromStr for MatrixJWTKID {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (mail, token_id) = s
            .split_once("#")
            .context("Failed to decode KID in two parts!")?;

        let mail = UserEmail(mail.to_string());

        if !mail.is_valid() {
            anyhow::bail!("Given email is invalid!")
        }

        Ok(Self {
            user_email: mail,
            id: token_id.parse().context("Failed to parse API token ID")?,
        })
    }
}

#[derive(Debug, serde::Serialize, serde::Deserialize)]
pub struct TokenClaims {
    #[serde(rename = "met")]
    pub method: String,
    pub uri: String,
    #[serde(rename = "pay", skip_serializing_if = "Option::is_none")]
    pub payload_sha256: Option<String>,
}

impl AuthExtractor {
    async fn extract_auth(
        req: &HttpRequest,
        remote_ip: IpAddr,
        payload_bytes: Option<Bytes>,
    ) -> Result<Self, actix_web::Error> {
        // Check for authentication using API token
        if let Some(token) = req.headers().get(constants::API_AUTH_HEADER) {
            let Ok(jwt_token) = token.to_str() else {
                return Err(actix_web::error::ErrorBadRequest(
                    "Failed to decode token as string!",
                ));
            };

            let metadata = match jwt_simple::token::Token::decode_metadata(jwt_token) {
                Ok(m) => m,
                Err(e) => {
                    log::error!("Failed to decode JWT header metadata! {e}");
                    return Err(actix_web::error::ErrorBadRequest(
                        "Failed to decode JWT header metadata!",
                    ));
                }
            };

            // Extract token ID
            let Some(kid) = metadata.key_id() else {
                return Err(actix_web::error::ErrorBadRequest(
                    "Missing key id in request!",
                ));
            };

            let jwt_kid = match MatrixJWTKID::from_str(kid) {
                Ok(i) => i,
                Err(e) => {
                    log::error!("Failed to parse token id! {e}");
                    return Err(actix_web::error::ErrorBadRequest(
                        "Failed to parse token id!",
                    ));
                }
            };

            // Get token information
            let Ok(mut token) = APIToken::load(&jwt_kid.user_email, &jwt_kid.id).await else {
                log::error!("Token not found!");
                return Err(actix_web::error::ErrorForbidden("Token not found!"));
            };

            // Decode JWT
            let key = HS256Key::from_bytes(token.secret.as_ref());
            let verif = VerificationOptions {
                max_validity: Some(Duration::from_secs(constants::API_TOKEN_JWT_MAX_DURATION)),
                ..Default::default()
            };

            let claims = match key.verify_token::<TokenClaims>(jwt_token, Some(verif)) {
                Ok(t) => t,
                Err(e) => {
                    log::error!("JWT validation failed! {e}");
                    return Err(actix_web::error::ErrorForbidden("JWT validation failed!"));
                }
            };

            // Check for nonce
            if claims.nonce.is_none() {
                return Err(actix_web::error::ErrorBadRequest(
                    "A nonce is required in auth JWT!",
                ));
            }

            // Check IP restriction
            if let Some(nets) = &token.base.networks
                && !nets.is_empty()
                && !nets.iter().any(|n| n.contains(&remote_ip))
            {
                log::error!(
                    "Trying to use token {:?} from unauthorized IP address: {remote_ip:?}",
                    token.id
                );
                return Err(actix_web::error::ErrorForbidden(
                    "This token cannot be used from this IP address!",
                ));
            }

            // Check for write access
            if token.base.read_only && !req.method().is_safe() {
                return Err(actix_web::error::ErrorBadRequest(
                    "Read only token cannot perform write operations!",
                ));
            }

            // Get user information
            let Ok(user) = User::get_by_mail(&jwt_kid.user_email).await else {
                return Err(actix_web::error::ErrorBadRequest(
                    "Failed to get user information from token!",
                ));
            };

            // Update last use (if needed)
            if token.shall_update_time_used() {
                token.last_used = time_secs();
                if let Err(e) = token.write(&jwt_kid.user_email).await {
                    log::error!("Failed to refresh last usage of token! {e}");
                }
            }

            // Handle tokens expiration
            if token.is_expired() {
                log::error!("Attempted to use expired token! {token:?}");
                return Err(actix_web::error::ErrorBadRequest("Token has expired!"));
            }

            // Check payload
            let payload = match (payload_bytes, claims.custom.payload_sha256) {
                (None, _) => None,
                (Some(_), None) => {
                    return Err(actix_web::error::ErrorBadRequest(
                        "A payload digest must be included in the JWT when the request has a payload!",
                    ));
                }
                (Some(payload), Some(provided_digest)) => {
                    let computed_digest = base16ct::lower::encode_string(&Sha256::digest(&payload));
                    if computed_digest != provided_digest {
                        log::error!(
                            "Expected digest {provided_digest} for payload but computed {computed_digest}!"
                        );
                        return Err(actix_web::error::ErrorBadRequest(
                            "Computed digest is different from the one provided in the JWT!",
                        ));
                    }

                    Some(payload.to_vec())
                }
            };

            return Ok(Self {
                method: AuthenticatedMethod::Token(token),
                user,
                payload,
            });
        }

        // Check if login is hard-coded as program argument
        if let Some(email) = &AppConfig::get().unsecure_auto_login_email() {
            let user = User::get_by_mail(email).await.map_err(|e| {
                log::error!("Failed to retrieve dev user: {e}");
                ErrorPreconditionFailed("Unable to retrieve dev user!")
            })?;
            return Ok(Self {
                method: AuthenticatedMethod::Dev,
                user,
                payload: payload_bytes.map(|bytes| bytes.to_vec()),
            });
        }

        // Check for cookie authentication
        let session = MatrixGWSession::extract(req).await?;
        if let Some(mail) = session.current_user().map_err(|e| {
            log::error!("Failed to retrieve user id: {e}");
            ErrorPreconditionFailed("Failed to read session information!")
        })? {
            let user = User::get_by_mail(&mail).await.map_err(|e| {
                log::error!("Failed to retrieve user from cookie session: {e}");
                ErrorPreconditionFailed("Failed to retrieve user information!")
            })?;
            return Ok(Self {
                method: AuthenticatedMethod::Cookie,
                user,
                payload: payload_bytes.map(|bytes| bytes.to_vec()),
            });
        };

        Err(ErrorPreconditionFailed("Authentication required!"))
    }
}

impl FromRequest for AuthExtractor {
    type Error = actix_web::Error;
    type Future = futures_util::future::LocalBoxFuture<'static, Result<Self, Self::Error>>;

    fn from_request(req: &HttpRequest, payload: &mut Payload) -> Self::Future {
        let req = req.clone();

        let remote_ip = match RemoteIP::from_request(&req, &mut Payload::None).into_inner() {
            Ok(ip) => ip,
            Err(e) => return Box::pin(async { Err(e) }),
        };

        let mut payload = payload.take();

        Box::pin(async move {
            let payload_bytes = match Bytes::from_request(&req, &mut payload).await {
                Ok(b) => {
                    if b.is_empty() {
                        None
                    } else {
                        Some(b)
                    }
                }
                Err(e) => {
                    log::error!("Failed to extract request payload! {e}");
                    None
                }
            };

            Self::extract_auth(&req, remote_ip.0, payload_bytes).await
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::extractors::auth_extractor::MatrixJWTKID;
    use crate::users::{APITokenID, UserEmail};
    use std::str::FromStr;

    #[test]
    fn encode_decode_jwt_kid() {
        let src = MatrixJWTKID {
            user_email: UserEmail("test@mail.com".to_string()),
            id: APITokenID::default(),
        };
        let encoded = src.to_string();
        let decoded = encoded.parse::<MatrixJWTKID>().unwrap();
        assert_eq!(src, decoded);

        MatrixJWTKID::from_str("bad").unwrap_err();
        MatrixJWTKID::from_str("ba#d").unwrap_err();
        MatrixJWTKID::from_str("test@valid.com#d").unwrap_err();
    }
}