athena_rs 2.5.3

Database gateway API
Documentation
use actix_web::{HttpRequest, HttpResponse};
use serde_json::json;
use std::collections::HashMap;
use std::env;
use tracing::warn;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApiKeySource {
    AuthorizationBearer,
    ApiKeyHeader,
    XApiKeyHeader,
    XAthenaKeyHeader,
    QueryParamApiKey,
    QueryParamApikey,
}

fn query_param_value(req: &HttpRequest, key: &str) -> Option<String> {
    req.uri().query().and_then(|query| {
        serde_urlencoded::from_str::<HashMap<String, String>>(query)
            .ok()
            .and_then(|params| params.get(key).cloned())
            .filter(|value| !value.trim().is_empty())
    })
}

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

    if let Some(value) = req
        .headers()
        .get("Authorization")
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.strip_prefix("Bearer ").map(str::to_string))
        .filter(|value| !value.trim().is_empty())
    {
        keys.push((value, ApiKeySource::AuthorizationBearer));
    }

    if let Some(value) = req
        .headers()
        .get("apikey")
        .and_then(|value| value.to_str().ok())
        .map(str::to_string)
        .filter(|value| !value.trim().is_empty())
    {
        keys.push((value, ApiKeySource::ApiKeyHeader));
    }

    if let Some(value) = req
        .headers()
        .get("x-api-key")
        .and_then(|value| value.to_str().ok())
        .map(str::to_string)
        .filter(|value| !value.trim().is_empty())
    {
        keys.push((value, ApiKeySource::XApiKeyHeader));
    }

    if let Some(value) = req
        .headers()
        .get("x-athena-key")
        .and_then(|value| value.to_str().ok())
        .map(str::to_string)
        .filter(|value| !value.trim().is_empty())
    {
        keys.push((value, ApiKeySource::XAthenaKeyHeader));
    }

    if let Some(value) = query_param_value(req, "api_key") {
        warn!("API key query parameter 'api_key' is deprecated; use headers instead");
        keys.push((value, ApiKeySource::QueryParamApiKey));
    }

    if let Some(value) = query_param_value(req, "apikey") {
        warn!("API key query parameter 'apikey' is deprecated; use headers instead");
        keys.push((value, ApiKeySource::QueryParamApikey));
    }

    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 {
    extract_api_keys_with_sources(req)
        .into_iter()
        .any(|(_, source)| {
            matches!(
                source,
                ApiKeySource::QueryParamApiKey | ApiKeySource::QueryParamApikey
            )
        })
}

/// 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)
}

/// 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(HttpResponse::InternalServerError().json(json!({
            "error": "ATHENA_KEY_12 is not configured"
        })));
    };

    let authorized: bool = extract_api_keys_with_sources(req)
        .into_iter()
        .any(|(provided, _)| provided == expected);

    if authorized {
        Ok(())
    } else {
        Err(HttpResponse::Unauthorized().json(json!({
            "error": "Invalid or missing API key. Use Authorization: Bearer <key>, X-Athena-Key, apikey, or X-API-Key. Query-param API keys are deprecated."
        })))
    }
}

#[cfg(test)]
mod tests {
    use super::{ApiKeySource, authorize_static_admin_key, extract_api_key_with_source};
    use actix_web::test::TestRequest;

    #[test]
    fn extract_prefers_authorization_bearer() {
        let req = TestRequest::default()
            .insert_header(("Authorization", "Bearer token-from-auth"))
            .insert_header(("x-api-key", "token-from-x-api"))
            .to_http_request();

        let extracted = extract_api_key_with_source(&req);
        assert_eq!(
            extracted,
            Some((
                "token-from-auth".to_string(),
                ApiKeySource::AuthorizationBearer
            ))
        );
    }

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

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

        let result = 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() {
        unsafe {
            std::env::set_var("ATHENA_KEY_12", "expected-token");
        }

        let req = TestRequest::default()
            .insert_header(("Authorization", "Bearer wrong-1"))
            .insert_header(("x-api-key", "wrong-2"))
            .insert_header(("apikey", "wrong-3"))
            .to_http_request();

        let result = authorize_static_admin_key(&req);

        assert!(result.is_err());

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