zesdex/crates/zesdex-middleware/src/auth.rs

//! Authentication middleware — session-lock based auth for Axum.
//!
//! Provides:
//! - [`SessionAuthLayer`]: a tower [`Layer`] that injects session validation
//! - [`SessionIdentity`]: extracted from validated requests
//! - [`validate_session`]: low-level session existence/validity check

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use axum::extract::FromRequestParts;
use axum::http::header;
use axum::http::request::Parts;
use axum::http::{Request, StatusCode};
use axum::response::{IntoResponse, Response};
use serde::{Deserialize, Serialize};
use tower::{Layer, Service};
use zesdex_entities::domain::common::store::Store;

// ---------------------------------------------------------------------------
// SessionIdentity
// ---------------------------------------------------------------------------

/// Identity extracted from a validated session token / lock.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionIdentity {
    /// The validated session id (from `X-Session-Id`).
    pub session_id: String,
    /// User-Agent header value, if present.
    pub user_agent: String,
    /// Unix-epoch timestamp (seconds) when the session was first seen by
    /// this middleware.
    pub connected_at: i64,
}

impl SessionIdentity {
    /// Create a new identity from a validated session id.
    fn new(session_id: String, user_agent: String) -> Self {
        let connected_at = chrono::Utc::now().timestamp();
        Self {
            session_id,
            user_agent,
            connected_at,
        }
    }
}

/// Extractor: pull the identity from request extensions.
///
/// If the identity has not been inserted by the middleware the request is
/// rejected with 401 Unauthorized.
impl<S: Send + Sync> FromRequestParts<S> for SessionIdentity {
    type Rejection = Response;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        parts
            .extensions
            .get::<SessionIdentity>()
            .cloned()
            .ok_or_else(|| (StatusCode::UNAUTHORIZED, "session identity not found").into_response())
    }
}

// ---------------------------------------------------------------------------
// SessionAuthLayer
// ---------------------------------------------------------------------------

/// Tower [`Layer`] that produces [`SessionAuthMiddleware`] services.
///
/// Wraps every request with session validation: if the `X-Session-Id`
/// header points to a valid session, the request passes through and a
/// [`SessionIdentity`] is injected into the request extensions.  Otherwise
/// a 401 response is returned immediately.
#[derive(Debug, Clone)]
pub struct SessionAuthLayer {
    store: Arc<Store>,
}

impl SessionAuthLayer {
    /// Create a new layer backed by the given [`Store`].
    pub fn new(store: Store) -> Self {
        Self {
            store: Arc::new(store),
        }
    }
}

impl Default for SessionAuthLayer {
    fn default() -> Self {
        Self::new(Store::new())
    }
}

impl<S> Layer<S> for SessionAuthLayer {
    type Service = SessionAuthMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        SessionAuthMiddleware {
            inner,
            store: Arc::clone(&self.store),
        }
    }
}

// ---------------------------------------------------------------------------
// SessionAuthMiddleware
// ---------------------------------------------------------------------------

/// Tower [`Service`] that validates `X-Session-Id` before forwarding.
#[derive(Debug, Clone)]
pub struct SessionAuthMiddleware<S> {
    inner: S,
    store: Arc<Store>,
}

// ---------------------------------------------------------------------------
// Session ID helpers
// ---------------------------------------------------------------------------

/// Extract and validate `X-Session-Id` from request headers.
///
/// Flow: read header -> validate non-empty -> return ID or a 401 error response.
fn extract_session_id<ReqBody>(req: &Request<ReqBody>) -> Result<String, Box<Response>> {
    let session_id = req
        .headers()
        .get("X-Session-Id") // custom header carrying the session identifier
        .and_then(|v| v.to_str().ok())
        .map(|s| s.to_string());

    match session_id {
        Some(id) if !id.is_empty() => Ok(id),
        _ => Err(Box::new(
            (StatusCode::UNAUTHORIZED, "missing X-Session-Id header").into_response(),
        )),
    }
}

/// Validate session and build identity from request context.
///
/// Flow: validate session in store -> extract User-Agent -> build SessionIdentity.
fn validate_and_build_identity<ReqBody>(
    session_id: &str,
    store: &Store,
    req: &Request<ReqBody>,
) -> Result<SessionIdentity, Box<Response>> {
    match validate_session(session_id, store) {
        Ok(_session) => {
            let user_agent = req
                .headers()
                .get(header::USER_AGENT)
                .and_then(|v| v.to_str().ok())
                .unwrap_or("")
                .to_string();
            Ok(SessionIdentity::new(session_id.to_string(), user_agent))
        }
        Err(e) => Err(Box::new(
            (
                StatusCode::UNAUTHORIZED,
                format!("session validation failed: {e}"),
            )
                .into_response(),
        )),
    }
}

impl<S, ReqBody> Service<Request<ReqBody>> for SessionAuthMiddleware<S>
where
    S: Service<Request<ReqBody>, Response = Response> + Send + 'static,
    S::Future: Send + 'static,
    ReqBody: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future =
        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
        let store = Arc::clone(&self.store);

        let session_id = match extract_session_id(&req) {
            Ok(id) => id,
            Err(resp) => return Box::pin(async move { Ok(*resp) }),
        };

        match validate_and_build_identity(&session_id, &store, &req) {
            Ok(identity) => {
                req.extensions_mut().insert(identity);
            }
            Err(resp) => return Box::pin(async move { Ok(*resp) }),
        };

        let fut = self.inner.call(req);
        Box::pin(fut)
    }
}

// ---------------------------------------------------------------------------
// Helper: `require_session` (convenience middleware function)
// ---------------------------------------------------------------------------

/// Axum middleware function that validates `X-Session-Id` against the
/// [`Store`] extracted from request extensions.
///
/// This is an alternative to [`SessionAuthLayer`] when you want to attach
/// auth to a specific route group via `axum::middleware::from_fn_with_state`.
pub async fn require_session(
    store: axum::extract::State<Store>,
    mut req: Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> Response {
    let session_id = match extract_session_id(&req) {
        Ok(id) => id,
        Err(resp) => return *resp,
    };

    let identity = match validate_and_build_identity(&session_id, &store, &req) {
        Ok(identity) => identity,
        Err(resp) => return *resp,
    };
    req.extensions_mut().insert(identity);
    next.run(req).await
}

// ---------------------------------------------------------------------------
// Session validation
// ---------------------------------------------------------------------------

/// Check whether a session lock exists and is valid, returning the
/// associated [`SessionIdentity`].
///
/// Validation logic:
/// 1. Verify the session id is not a path-traversal attack.
/// 2. Check that `<store.base_dir>/sessions/<id>/session.json` exists.
/// 3. Deserialise the session metadata to confirm it is well-formed.
///
/// This is a synchronous, CPU-light check so it can be called directly
/// inside tower service impls without spawning a blocking task.
pub fn validate_session(session_id: &str, store: &Store) -> anyhow::Result<SessionIdentity> {
    // Directory-traversal prevention.
    if session_id.contains('/') || session_id.contains('\\') || session_id.contains("..") {
        anyhow::bail!("invalid session id: must not contain path separators");
    }

    let session_path = store
        .base_dir
        .join("sessions")
        .join(session_id)
        .join("session.json"); // path to session metadata file

    if !session_path.exists() {
        anyhow::bail!("session not found: {session_id}");
    }

    let _data = std::fs::read_to_string(&session_path)?; // raw session JSON
    // We verify the JSON is well-formed by deserialising it.
    let _session: serde_json::Value = serde_json::from_str(&_data)?;

    let user_agent = String::new();
    Ok(SessionIdentity::new(session_id.to_string(), user_agent))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_validate_session_rejects_path_traversal() {
        let store = Store::new();
        assert!(validate_session("../etc/passwd", &store).is_err());
        assert!(validate_session("foo/bar", &store).is_err());
        assert!(validate_session("foo\\bar", &store).is_err());
        assert!(validate_session("..", &store).is_err());
    }

    #[test]
    fn test_validate_session_nonexistent() {
        let store = Store::new();
        let result = validate_session("nonexistent-session-id", &store);
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("session not found"));
    }

    #[test]
    fn test_session_identity_creation() {
        let identity = SessionIdentity::new("sess-123".into(), "test-agent".into());
        assert_eq!(identity.session_id, "sess-123");
        assert_eq!(identity.user_agent, "test-agent");
        assert!(identity.connected_at > 0);
    }
}