athena_rs 3.4.7

Database driver
Documentation
use crate::api::response::internal_error;
use actix_web::{HttpRequest, HttpResponse};
use serde_json::json;
use std::env;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApiKeySource {
    XAthenaKeyHeader,
}

fn header_trimmed(req: &HttpRequest, name: &str) -> Option<String> {
    req.headers()
        .get(name)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(str::to_string)
}

/// Key sent as `X-Athena-Key`, used for gateway API-key auth (`ath_*` keys).
pub fn extract_gateway_api_key(req: &HttpRequest) -> Option<String> {
    header_trimmed(req, "x-athena-key")
}

fn extract_api_keys_with_sources(req: &HttpRequest) -> Vec<(String, ApiKeySource)> {
    let mut keys: Vec<(String, ApiKeySource)> = Vec::new();

    if let Some(value) = extract_gateway_api_key(req) {
        keys.push((value, ApiKeySource::XAthenaKeyHeader));
    }

    keys
}

pub fn extract_api_key_with_source(req: &HttpRequest) -> Option<(String, ApiKeySource)> {
    extract_api_keys_with_sources(req).into_iter().next()
}

pub fn api_key_query_param_used(req: &HttpRequest) -> bool {
    let _ = req;
    false
}

/// Extract an API key from standard request locations.
///
/// The first non-empty value wins.
pub fn extract_api_key(req: &HttpRequest) -> Option<String> {
    extract_api_key_with_source(req).map(|(value, _)| value)
}

fn static_admin_key_candidates(req: &HttpRequest) -> Vec<String> {
    let mut out: Vec<String> = Vec::new();
    if let Some(value) = extract_gateway_api_key(req) {
        out.push(value);
    }
    if let Some(value) = header_trimmed(req, "x-athena-admin-key") {
        out.push(value);
    }
    out
}

/// Validate the static admin key stored in `ATHENA_KEY_12`.
pub fn authorize_static_admin_key(req: &HttpRequest) -> Result<(), HttpResponse> {
    let expected: Option<String> = env::var("ATHENA_KEY_12")
        .ok()
        .filter(|value| !value.is_empty());

    let Some(expected) = expected else {
        return Err(internal_error(
            "ATHENA_KEY_12 is not configured",
            "ATHENA_KEY_12 environment variable is not set",
        ));
    };

    let authorized: bool = static_admin_key_candidates(req)
        .into_iter()
        .any(|provided| provided == expected);

    if authorized {
        Ok(())
    } else {
        Err(HttpResponse::Unauthorized().json(json!({
            "error": "Invalid or missing API key. Use X-Athena-Key."
        })))
    }
}

#[cfg(test)]
mod tests {
    use super::{ApiKeySource, authorize_static_admin_key, extract_api_key_with_source};
    use actix_web::test::TestRequest;
    use actix_web::{HttpRequest, HttpResponse};
    use std::sync::{Mutex, OnceLock};

    static ADMIN_KEY_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();

    fn admin_key_env_lock() -> std::sync::MutexGuard<'static, ()> {
        ADMIN_KEY_TEST_LOCK
            .get_or_init(|| Mutex::new(()))
            .lock()
            .expect("admin key test lock poisoned")
    }

    #[test]
    fn extract_reads_x_athena_key() {
        let req: HttpRequest = TestRequest::default()
            .insert_header(("x-athena-key", "token-from-x-athena"))
            .to_http_request();

        let extracted: Option<(String, ApiKeySource)> = extract_api_key_with_source(&req);
        assert_eq!(
            extracted,
            Some((
                "token-from-x-athena".to_string(),
                ApiKeySource::XAthenaKeyHeader
            ))
        );
    }

    #[test]
    fn authorize_accepts_when_any_supported_header_matches() {
        let _guard = admin_key_env_lock();
        unsafe {
            std::env::set_var("ATHENA_KEY_12", "expected-token");
        }

        let req: HttpRequest = TestRequest::default()
            .insert_header(("x-athena-key", "expected-token"))
            .to_http_request();

        let result: Result<(), HttpResponse> = authorize_static_admin_key(&req);

        assert!(result.is_ok());

        unsafe {
            std::env::remove_var("ATHENA_KEY_12");
        }
    }

    #[test]
    fn authorize_accepts_x_athena_admin_key_when_gateway_key_differs() {
        let _guard = admin_key_env_lock();
        unsafe {
            std::env::set_var("ATHENA_KEY_12", "admin-token");
        }

        let req: HttpRequest = TestRequest::default()
            .insert_header(("x-athena-key", "ath_pub.secret"))
            .insert_header(("x-athena-admin-key", "admin-token"))
            .to_http_request();

        let result: Result<(), HttpResponse> = authorize_static_admin_key(&req);

        assert!(result.is_ok());

        unsafe {
            std::env::remove_var("ATHENA_KEY_12");
        }
    }

    #[test]
    fn authorize_rejects_when_all_provided_keys_are_wrong() {
        let _guard = admin_key_env_lock();
        unsafe {
            std::env::set_var("ATHENA_KEY_12", "expected-token");
        }

        let req = TestRequest::default()
            .insert_header(("x-athena-key", "wrong"))
            .to_http_request();

        let result: Result<(), HttpResponse> = authorize_static_admin_key(&req);

        assert!(result.is_err());

        unsafe {
            std::env::remove_var("ATHENA_KEY_12");
        }
    }
}