modo-rs 0.8.0

use serde::Deserialize;

use http::request::Parts;

/// Trait for extracting JWT token strings from HTTP requests.
///
/// Middleware tries sources in order and uses the first `Some(token)`.
/// Implement this trait to support custom token locations.
pub trait TokenSource: Send + Sync {
    /// Attempts to extract a token string from request parts.
    /// Returns `None` if this source does not find a token.
    fn extract(&self, parts: &Parts) -> Option<String>;
}

/// Extracts a token from the `Authorization: Bearer <token>` header.
///
/// The scheme comparison is case-insensitive per RFC 7235 — `Bearer`,
/// `bearer`, `BEARER`, etc. are all accepted.
pub struct BearerSource;

impl TokenSource for BearerSource {
    fn extract(&self, parts: &Parts) -> Option<String> {
        let value = parts
            .headers
            .get(http::header::AUTHORIZATION)?
            .to_str()
            .ok()?;
        let (scheme, rest) = value.split_once(' ')?;
        if !scheme.eq_ignore_ascii_case("Bearer") {
            return None;
        }
        let token = rest.trim_start();
        if token.is_empty() {
            return None;
        }
        Some(token.to_string())
    }
}

/// Extracts a token from a named query parameter (e.g., `?token=xxx`).
///
/// The inner `&'static str` is the parameter name to look up.
pub struct QuerySource(pub &'static str);

impl TokenSource for QuerySource {
    fn extract(&self, parts: &Parts) -> Option<String> {
        let query = parts.uri.query()?;
        for pair in query.split('&') {
            if let Some((key, value)) = pair.split_once('=')
                && key == self.0
                && !value.is_empty()
            {
                return Some(value.to_string());
            }
        }
        None
    }
}

/// Extracts a token from a named cookie (e.g., `Cookie: jwt=xxx`).
///
/// The inner `&'static str` is the cookie name. Parses the raw `Cookie`
/// header directly — no dependency on session middleware or `axum_extra`.
pub struct CookieSource(pub &'static str);

impl TokenSource for CookieSource {
    fn extract(&self, parts: &Parts) -> Option<String> {
        let cookie_header = parts.headers.get(http::header::COOKIE)?.to_str().ok()?;
        for cookie in cookie_header.split(';') {
            let cookie = cookie.trim();
            if let Some((name, value)) = cookie.split_once('=')
                && name.trim() == self.0
                && !value.is_empty()
            {
                return Some(value.trim().to_string());
            }
        }
        None
    }
}

/// Extracts a token from a custom request header (e.g., `X-API-Token: xxx`).
///
/// The inner `&'static str` is the header name.
pub struct HeaderSource(pub &'static str);

impl TokenSource for HeaderSource {
    fn extract(&self, parts: &Parts) -> Option<String> {
        let value = parts.headers.get(self.0)?.to_str().ok()?;
        if value.is_empty() {
            return None;
        }
        Some(value.to_string())
    }
}

// ── Owned-string variants used internally by TokenSourceConfig::build() ────────

struct OwnedQuerySource(String);

impl TokenSource for OwnedQuerySource {
    fn extract(&self, parts: &Parts) -> Option<String> {
        let query = parts.uri.query()?;
        for pair in query.split('&') {
            if let Some((key, value)) = pair.split_once('=')
                && key == self.0
                && !value.is_empty()
            {
                return Some(value.to_string());
            }
        }
        None
    }
}

struct OwnedCookieSource(String);

impl TokenSource for OwnedCookieSource {
    fn extract(&self, parts: &Parts) -> Option<String> {
        let cookie_header = parts.headers.get(http::header::COOKIE)?.to_str().ok()?;
        for cookie in cookie_header.split(';') {
            let cookie = cookie.trim();
            if let Some((name, value)) = cookie.split_once('=')
                && name.trim() == self.0
                && !value.is_empty()
            {
                return Some(value.trim().to_string());
            }
        }
        None
    }
}

struct OwnedHeaderSource(String);

impl TokenSource for OwnedHeaderSource {
    fn extract(&self, parts: &Parts) -> Option<String> {
        let value = parts.headers.get(self.0.as_str())?.to_str().ok()?;
        if value.is_empty() {
            return None;
        }
        Some(value.to_string())
    }
}

// ── TokenSourceConfig ────────────────────────────────────────────────────────

/// YAML-deserialized configuration that selects and constructs a [`TokenSource`].
///
/// Used in [`JwtSessionsConfig`](super::config::JwtSessionsConfig) for
/// `access_source` and `refresh_source`.
///
/// # Examples
///
/// ```yaml
/// access_source:
///   kind: bearer
///
/// refresh_source:
///   kind: cookie
///   name: refresh_jwt
/// ```
#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum TokenSourceConfig {
    /// `Authorization: Bearer <token>` header.
    Bearer,
    /// Named cookie (e.g., `Cookie: jwt=xxx`).
    Cookie { name: String },
    /// Custom request header (e.g., `X-API-Token: xxx`).
    Header { name: String },
    /// Named query parameter (e.g., `?token=xxx`).
    Query { name: String },
    /// JSON body field — the token is read from the request body during
    /// session operations; not extracted by [`JwtLayer`](super::middleware::JwtLayer).
    Body { field: String },
}

impl TokenSourceConfig {
    /// Build a boxed [`TokenSource`] for use in middleware.
    ///
    /// `Body` variants fall back to `BearerSource` because body extraction is
    /// handled at the session-handler level, not in request-parts middleware.
    pub fn build(&self) -> Box<dyn TokenSource> {
        match self {
            Self::Bearer => Box::new(BearerSource),
            Self::Cookie { name } => Box::new(OwnedCookieSource(name.clone())),
            Self::Header { name } => Box::new(OwnedHeaderSource(name.clone())),
            Self::Query { name } => Box::new(OwnedQuerySource(name.clone())),
            // Body tokens are read in the session handler, not by JwtLayer.
            Self::Body { .. } => Box::new(BearerSource),
        }
    }
}

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

    fn parts_with_header(name: &str, value: &str) -> Parts {
        let (parts, _) = http::Request::builder()
            .header(name, value)
            .body(())
            .unwrap()
            .into_parts();
        parts
    }

    fn parts_with_uri(uri: &str) -> Parts {
        let (parts, _) = http::Request::builder()
            .uri(uri)
            .body(())
            .unwrap()
            .into_parts();
        parts
    }

    fn empty_parts() -> Parts {
        let (parts, _) = http::Request::builder().body(()).unwrap().into_parts();
        parts
    }

    // BearerSource tests
    #[test]
    fn bearer_extracts_token() {
        let parts = parts_with_header("Authorization", "Bearer my-token");
        assert_eq!(BearerSource.extract(&parts), Some("my-token".into()));
    }

    #[test]
    fn bearer_case_insensitive_prefix() {
        let parts = parts_with_header("Authorization", "bearer my-token");
        assert_eq!(BearerSource.extract(&parts), Some("my-token".into()));
    }

    #[test]
    fn bearer_uppercase_scheme_works() {
        let parts = parts_with_header("Authorization", "BEARER my-token");
        assert_eq!(BearerSource.extract(&parts), Some("my-token".into()));
    }

    #[test]
    fn bearer_mixed_case_scheme_works() {
        let parts = parts_with_header("Authorization", "BeArEr my-token");
        assert_eq!(BearerSource.extract(&parts), Some("my-token".into()));
    }

    #[test]
    fn bearer_returns_none_when_missing() {
        assert!(BearerSource.extract(&empty_parts()).is_none());
    }

    #[test]
    fn bearer_returns_none_for_non_bearer_scheme() {
        let parts = parts_with_header("Authorization", "Basic abc123");
        assert!(BearerSource.extract(&parts).is_none());
    }

    #[test]
    fn bearer_returns_none_for_empty_token() {
        let parts = parts_with_header("Authorization", "Bearer ");
        assert!(BearerSource.extract(&parts).is_none());
    }

    // QuerySource tests
    #[test]
    fn query_extracts_token() {
        let parts = parts_with_uri("/path?token=my-token&other=val");
        assert_eq!(
            QuerySource("token").extract(&parts),
            Some("my-token".into())
        );
    }

    #[test]
    fn query_returns_none_when_missing() {
        let parts = parts_with_uri("/path?other=val");
        assert!(QuerySource("token").extract(&parts).is_none());
    }

    #[test]
    fn query_returns_none_for_empty_value() {
        let parts = parts_with_uri("/path?token=");
        assert!(QuerySource("token").extract(&parts).is_none());
    }

    #[test]
    fn query_returns_none_without_query_string() {
        let parts = parts_with_uri("/path");
        assert!(QuerySource("token").extract(&parts).is_none());
    }

    // CookieSource tests
    #[test]
    fn cookie_extracts_token() {
        let parts = parts_with_header("Cookie", "jwt=my-token; other=val");
        assert_eq!(CookieSource("jwt").extract(&parts), Some("my-token".into()));
    }

    #[test]
    fn cookie_returns_none_when_missing() {
        let parts = parts_with_header("Cookie", "other=val");
        assert!(CookieSource("jwt").extract(&parts).is_none());
    }

    #[test]
    fn cookie_returns_none_without_cookie_header() {
        assert!(CookieSource("jwt").extract(&empty_parts()).is_none());
    }

    // HeaderSource tests
    #[test]
    fn header_extracts_token() {
        let parts = parts_with_header("X-API-Token", "my-token");
        assert_eq!(
            HeaderSource("X-API-Token").extract(&parts),
            Some("my-token".into())
        );
    }

    #[test]
    fn header_returns_none_when_missing() {
        assert!(
            HeaderSource("X-API-Token")
                .extract(&empty_parts())
                .is_none()
        );
    }

    #[test]
    fn header_returns_none_for_empty_value() {
        let parts = parts_with_header("X-API-Token", "");
        assert!(HeaderSource("X-API-Token").extract(&parts).is_none());
    }
}