athena_rs 3.22.1

Hyper performant polyglot Database driver
Documentation
use actix_web::{HttpRequest, HttpResponse};
use athena_auth_core::{
    ValidateSessionInput, ValidatedSessionContext, extract_session_token_from_headers,
    validate_session,
};
use athena_chat::ChatContext;
use chrono::Utc;

use crate::AppState;
use crate::api::client_context::auth_pool;
use crate::api::client_context::required_client_name;
use crate::api::headers::request_context::athena_request_id;
use crate::api::response::{service_unavailable, unauthorized};

#[derive(Debug, Default)]
pub struct AuthResolver;

impl AuthResolver {
    pub async fn resolve(
        &self,
        req: &HttpRequest,
        state: &AppState,
    ) -> Result<ChatContext, HttpResponse> {
        let client_name = required_client_name(req)?;
        let Some(session_token) = extract_session_token_from_headers(
            req.headers()
                .get("authorization")
                .and_then(|value| value.to_str().ok()),
            req.headers()
                .get("cookie")
                .and_then(|value| value.to_str().ok()),
            None,
        ) else {
            return Err(unauthorized(
                "Chat authentication required",
                "Missing or invalid Athena Auth session",
            ));
        };
        let pool = auth_pool(state)?;
        let validated_session = validate_session(
            &pool,
            ValidateSessionInput {
                token: &session_token.token,
                now: Utc::now(),
                include_user: true,
                include_organization: true,
            },
        )
        .await
        .map_err(|err| {
            service_unavailable(
                "Chat auth store unavailable",
                format!("failed to verify Athena Auth session: {err}"),
            )
        })?;

        let Some(validated_session) = validated_session else {
            return Err(unauthorized(
                "Chat authentication required",
                "Missing or invalid Athena Auth session",
            ));
        };

        build_chat_context(validated_session, client_name, athena_request_id(req))
    }
}

fn build_chat_context(
    validated_session: ValidatedSessionContext,
    client_name: String,
    request_id: Option<String>,
) -> Result<ChatContext, HttpResponse> {
    let organization_id = validated_session
        .active_organization_id
        .clone()
        .or_else(|| validated_session.organization_ids.first().cloned())
        .ok_or_else(|| {
            unauthorized(
                "Chat organization required",
                "No active organization resolved for actor",
            )
        })?;
    let role_ids = validated_session.merged_roles();

    Ok(ChatContext {
        user_id: validated_session.user_id,
        organization_id,
        client_name,
        is_admin: actor_is_admin(&role_ids, validated_session.has_admin_scope),
        role_ids,
        request_id,
    })
}

fn actor_is_admin(role_ids: &[String], has_admin_scope: bool) -> bool {
    has_admin_scope
        || role_ids.iter().any(|role| {
            matches!(
                role.as_str(),
                "admin" | "super_admin" | "owner" | "storage_admin" | "storage-admin"
            )
        })
}

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

    fn sample_session() -> ValidatedSessionContext {
        ValidatedSessionContext {
            session_id: "session-1".to_string(),
            user_id: "user-1".to_string(),
            email: Some("user@example.com".to_string()),
            name: Some("User".to_string()),
            role: Some("member".to_string()),
            active_organization_id: Some("org-2".to_string()),
            organization_ids: vec!["org-1".to_string()],
            member_roles: vec!["owner".to_string()],
            permissions: Vec::new(),
            has_admin_scope: false,
            expires_at: Utc::now(),
        }
    }

    #[test]
    fn build_chat_context_prefers_active_organization() {
        let ctx = build_chat_context(
            sample_session(),
            "acme".to_string(),
            Some("req-1".to_string()),
        )
        .expect("chat context should resolve");

        assert_eq!(ctx.organization_id, "org-2");
        assert_eq!(
            ctx.role_ids,
            vec!["member".to_string(), "owner".to_string()]
        );
        assert!(ctx.is_admin);
    }

    #[test]
    fn build_chat_context_falls_back_to_first_member_organization() {
        let mut session = sample_session();
        session.active_organization_id = None;

        let ctx = build_chat_context(session, "acme".to_string(), None)
            .expect("chat context should resolve");

        assert_eq!(ctx.organization_id, "org-1");
    }

    #[test]
    fn build_chat_context_rejects_missing_organization_membership() {
        let mut session = sample_session();
        session.active_organization_id = None;
        session.organization_ids.clear();

        let response = build_chat_context(session, "acme".to_string(), None)
            .expect_err("missing org should fail");

        assert_eq!(response.status(), actix_web::http::StatusCode::UNAUTHORIZED);
    }
}