doctate/server/src/auth.rs

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;

use axum::extract::{FromRef, FromRequestParts};
use axum::http::request::Parts;
use axum_extra::extract::cookie::CookieJar;
use doctate_common::API_KEY_HEADER;
use tracing::warn;

use crate::config::Config;
use crate::error::AppError;
use crate::web_session::{SESSION_COOKIE, SessionStore};

/// Truncate a session token to a non-sensitive prefix for logs.
fn token_prefix(token: &str) -> String {
    token.chars().take(8).collect()
}

/// Extracted from the X-API-Key header on every authenticated request.
pub struct AuthenticatedUser {
    pub slug: String,
    pub role: String,
    pub data_dir: PathBuf,
}

impl AuthenticatedUser {
    pub fn is_admin(&self) -> bool {
        self.role == "admin"
    }
}

impl<S> FromRequestParts<S> for AuthenticatedUser
where
    S: Send + Sync,
    Arc<Config>: FromRef<S>,
{
    type Rejection = AppError;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let config = Arc::<Config>::from_ref(state);

        let api_key = parts
            .headers
            .get(API_KEY_HEADER)
            .and_then(|v| v.to_str().ok())
            .ok_or(AppError::Unauthorized)?;

        let slug = config.api_keys.get(api_key).ok_or(AppError::Unauthorized)?;

        let user = config
            .users
            .iter()
            .find(|u| u.slug == *slug)
            .ok_or(AppError::Unauthorized)?;

        let data_dir = config.data_path.join(slug);
        tokio::fs::create_dir_all(&data_dir).await?;

        Ok(Self {
            slug: slug.clone(),
            role: user.role.clone(),
            data_dir,
        })
    }
}

/// Extracted from the session cookie on web UI requests.
/// On missing/expired session, returns `AppError::Redirect("/web/login")`.
pub struct AuthenticatedWebUser {
    pub slug: String,
    pub role: String,
    pub data_dir: PathBuf,
}

impl AuthenticatedWebUser {
    pub fn is_admin(&self) -> bool {
        self.role == "admin"
    }
}

impl<S> FromRequestParts<S> for AuthenticatedWebUser
where
    S: Send + Sync,
    Arc<Config>: FromRef<S>,
    SessionStore: FromRef<S>,
{
    type Rejection = AppError;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let config = Arc::<Config>::from_ref(state);
        let store = SessionStore::from_ref(state);

        let jar = CookieJar::from_headers(&parts.headers);
        let token = jar
            .get(SESSION_COOKIE)
            .map(|c| c.value().to_owned())
            .ok_or_else(|| AppError::Redirect("/web/login".into()))?;

        // Read lock first to check validity. If expired, upgrade to write and remove.
        let expired = {
            let r = store.read().await;
            match r.get(&token) {
                Some(s) => s.expires_at <= Instant::now(),
                None => {
                    warn!(token_prefix = %token_prefix(&token), "session lookup failed (unknown or already expired)");
                    return Err(AppError::Redirect("/web/login".into()));
                }
            }
        };

        if expired {
            let mut w = store.write().await;
            let removed = w.remove(&token);
            if let Some(s) = removed {
                warn!(slug = %s.slug, "session expired");
            }
            return Err(AppError::Redirect("/web/login".into()));
        }

        let r = store.read().await;
        let session = r
            .get(&token)
            .ok_or_else(|| AppError::Redirect("/web/login".into()))?;

        let data_dir = config.data_path.join(&session.slug);
        Ok(Self {
            slug: session.slug.clone(),
            role: session.role.clone(),
            data_dir,
        })
    }
}