factoriauth/src/auth/mod.rs

use std::sync::Arc;

use base64::{prelude::BASE64_STANDARD, Engine};
use hmac::{Hmac, Mac};
use md5::Md5;
use rand::{
    distributions::{Alphanumeric, DistString},
    thread_rng,
};
use secrecy::ExposeSecret;
use serde::Deserialize;
use sha2::Sha256;
use thiserror::Error;
use time::{macros::format_description, OffsetDateTime};
use tokio::sync::Mutex;
use tracing::{event, instrument, Level};
use url::Url;

use crate::{
    config::{AuthBackendConfig, PadlockConfig},
    db::{Database, UserTokenEntry},
    secrets::{
        PadlockGenerationSecret, Password, ServerHash, ServerPadlock, UserServerKey, UserToken,
    },
};

use self::backends::{LdapBackend, ValidateLogin};

pub mod backends;

#[derive(Debug, Error)]
pub enum AuthenticationError {
    #[error("Invalid username or password")]
    InvalidUserOrPassword,
    #[error("Invalid token")]
    InvalidToken,
    #[error("Invalid server hash")]
    InvalidServerHash,
    #[error("Database error")]
    Database(#[from] sqlx::Error),
    #[error("Authentication backend error")]
    Backend(#[from] ldap3::LdapError),
    #[error("Padlock proxy error")]
    PadlockProxy(#[from] PadlockProxyError),
    #[error("No authentication backends available")]
    NoBackends,
}

#[derive(Debug)]
pub enum AuthenticationBackend {
    Ldap(Arc<LdapBackend>),
}

impl AuthenticationBackend {
    pub async fn new(config: AuthBackendConfig) -> Result<Self, AuthenticationError> {
        match config {
            AuthBackendConfig::Ldap(c) => Ok(Self::Ldap(LdapBackend::new(c).await?)),
        }
    }
}

impl ValidateLogin for AuthenticationBackend {
    async fn validate_login(
        &self,
        username: &str,
        password: &Password,
    ) -> Result<String, AuthenticationError> {
        match self {
            AuthenticationBackend::Ldap(b) => b.validate_login(username, password).await,
        }
    }
}

#[derive(Debug)]
pub struct UserAuthenticator {
    db: Arc<Mutex<Box<dyn Database + Send>>>,
    backends: Vec<AuthenticationBackend>,
}

impl UserAuthenticator {
    const TOKEN_LEN: usize = 30;

    pub fn new(
        db: Arc<Mutex<Box<dyn Database + Send>>>,
        backends: Vec<AuthenticationBackend>,
    ) -> Self {
        Self { db, backends }
    }

    /// Attempt to verify credentials against all backends and return the user's most recent token if they match.
    ///
    /// If a valid token already exists, return it; if not, create a new one.
    #[instrument]
    pub async fn get_user_token(
        &self,
        username: &str,
        password: &Password,
    ) -> Result<(String, UserToken), AuthenticationError> {
        let mut auth_result = Err(AuthenticationError::NoBackends);

        for backend in &self.backends {
            auth_result = backend.validate_login(username, password).await;
            match &auth_result {
                Ok(_) => {
                    break;
                }
                Err(err) => {
                    event!(Level::WARN, ?err, "Authentication with backend failed");
                }
            }
        }

        let Ok(username) = auth_result else {
            event!(Level::WARN, "Authentication failed with all backends");
            return Err(AuthenticationError::InvalidUserOrPassword);
        };

        let mut db = self.db.lock().await;
        let token = if let Some(UserTokenEntry::Valid(old_token, _, _)) =
            db.get_user_token(&username).await?
        {
            old_token
        } else {
            let new_token =
                UserToken::from(Alphanumeric.sample_string(&mut thread_rng(), Self::TOKEN_LEN));

            db.save_user_token(&username, &new_token).await?;

            new_token
        };

        Ok((username, token))
    }

    /// Check if the passed token matches the user's currently active one.
    ///
    /// Additionally, update the token's `last_used` value if verification is successful.
    #[instrument]
    pub async fn verify_user_token(
        &self,
        username: &str,
        token: &UserToken,
    ) -> Result<(), AuthenticationError> {
        let mut db = self.db.lock().await;

        if let Some(UserTokenEntry::Valid(user_token, ..)) = &db.get_user_token(username).await? {
            if token == user_token {
                db.update_user_token_last_used(username, token).await?;
                return Ok(());
            }
        }

        Err(AuthenticationError::InvalidToken)
    }
}

#[derive(Debug)]
pub struct ServerPadlockSecret {
    secret: PadlockGenerationSecret,
}

impl ServerPadlockSecret {
    const HASH_LEN: usize = 32;

    pub fn new(secret: PadlockGenerationSecret) -> Self {
        Self { secret }
    }

    #[instrument]
    pub fn generate_hash() -> ServerHash {
        ServerHash(Alphanumeric.sample_string(&mut thread_rng(), Self::HASH_LEN))
    }

    #[instrument]
    pub fn generate_padlock(&self, server_hash: &ServerHash) -> ServerPadlock {
        #[allow(clippy::expect_used)]
        let mut hmac: Hmac<Sha256> = Hmac::new_from_slice(self.secret.0.expose_secret())
            .expect("HMAC should accept key of any length");

        hmac.update(server_hash.0.as_bytes());

        BASE64_STANDARD.encode(hmac.finalize().into_bytes()).into()
    }
}

#[derive(Debug, Error)]
pub enum PadlockProxyError {
    #[error("upstream responded with error")]
    Upstream(#[from] reqwest::Error),
    #[error("database error")]
    Database(#[from] sqlx::Error),
    #[error("invalid upstream URL")]
    InvalidUrl(#[from] url::ParseError),
    #[error("unknown server_hash")]
    UnknownHash(ServerHash),
}

#[derive(Debug)]
pub struct ServerPadlockProxy {
    client: reqwest::Client,
    upstream: Url,
    database: Arc<Mutex<Box<dyn Database + Send>>>,
}

impl ServerPadlockProxy {
    pub fn new(
        upstream: &Url,
        database: Arc<Mutex<Box<dyn Database + Send>>>,
    ) -> Result<Self, PadlockProxyError> {
        let client = reqwest::Client::new();
        let upstream = upstream.join("generate-server-padlock-2")?;

        Ok(Self {
            client,
            upstream,
            database,
        })
    }

    #[instrument]
    pub async fn generate_hash(&self) -> Result<ServerHash, PadlockProxyError> {
        #[derive(Deserialize)]
        struct ServerPadlockResponse {
            server_hash: ServerHash,
            server_padlock: ServerPadlock,
        }

        let response: ServerPadlockResponse = self
            .client
            .post(self.upstream.clone())
            .query(&[("api_version", 6)])
            .send()
            .await?
            .json()
            .await?;

        self.database
            .lock()
            .await
            .save_server_padlock(&response.server_hash, &response.server_padlock)
            .await?;

        Ok(response.server_hash)
    }

    #[instrument]
    pub async fn generate_padlock(
        &self,
        server_hash: &ServerHash,
    ) -> Result<ServerPadlock, PadlockProxyError> {
        let Some(padlock) = self
            .database
            .lock()
            .await
            .get_server_padlock(server_hash)
            .await?
        else {
            return Err(PadlockProxyError::UnknownHash(server_hash.clone()));
        };

        Ok(padlock)
    }
}

#[derive(Debug)]
pub enum ServerPadlockGenerator {
    Secret(ServerPadlockSecret),
    Proxy(ServerPadlockProxy),
}

impl ServerPadlockGenerator {
    pub fn new(
        config: PadlockConfig,
        db: Arc<Mutex<Box<dyn Database + Send>>>,
    ) -> Result<Self, PadlockProxyError> {
        match config {
            PadlockConfig::Secret(s) => Ok(Self::Secret(ServerPadlockSecret::new(s))),
            PadlockConfig::Proxy(u) => ServerPadlockProxy::new(&u, db).map(Self::Proxy),
        }
    }

    pub async fn generate_hash(&self) -> Result<ServerHash, PadlockProxyError> {
        match self {
            ServerPadlockGenerator::Secret(_) => Ok(ServerPadlockSecret::generate_hash()),
            ServerPadlockGenerator::Proxy(p) => p.generate_hash().await,
        }
    }

    pub async fn generate_padlock(
        &self,
        server_hash: &ServerHash,
    ) -> Result<ServerPadlock, PadlockProxyError> {
        match self {
            ServerPadlockGenerator::Secret(s) => Ok(s.generate_padlock(server_hash)),
            ServerPadlockGenerator::Proxy(p) => p.generate_padlock(server_hash).await,
        }
    }
}

#[derive(Debug)]
pub struct UserServerKeyGenerator {
    user_authenticator: Arc<UserAuthenticator>,
    padlock_generator: Arc<ServerPadlockGenerator>,
}

impl UserServerKeyGenerator {
    pub fn new(
        user_authenticator: Arc<UserAuthenticator>,
        padlock_generator: Arc<ServerPadlockGenerator>,
    ) -> Self {
        Self {
            user_authenticator,
            padlock_generator,
        }
    }

    #[instrument]
    pub async fn generate_user_server_key(
        &self,
        username: &str,
        token: &UserToken,
        server_hash: &ServerHash,
    ) -> Result<(UserServerKey, String), AuthenticationError> {
        self.user_authenticator
            .verify_user_token(username, token)
            .await?;

        let padlock = self.padlock_generator.generate_padlock(server_hash).await?;

        #[allow(clippy::expect_used)]
        let timestamp = OffsetDateTime::now_utc()
            .format(format_description!(
                "[year repr:last_two][month][day][hour repr:24][minute][second]"
            ))
            .expect("timestamp format should be validated at compile-time");

        event!(Level::DEBUG, timestamp, "Generating user_server_key");

        let mut mac: Hmac<Md5> = Hmac::new_from_slice(padlock.0.expose_secret().as_bytes())
            .map_err(|_e| AuthenticationError::InvalidServerHash)?;
        mac.update(format!("{}_{}_{}", username, padlock.0.expose_secret(), timestamp).as_bytes());
        let user_server_key =
            UserServerKey(BASE64_STANDARD.encode(mac.finalize().into_bytes()).into());

        Ok((user_server_key, timestamp))
    }
}