alibi-plugins 0.10.0

//! Application-owned asynchronous session response projection.
use alibi_core::endpoint::{
    EndpointCall, EndpointContextPatch, EndpointDefinition, EndpointHook, EndpointInput,
    EndpointResponse, ServerEndpoint,
};
use alibi_core::{
    AuthContext, AuthError, AuthInitContext, AuthPlugin, AuthRequest, AuthResponse, AuthResult,
    AuthRoute, AuthSchema, HttpMethod,
};
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
/// The actual transport frame supplied to a session projection callback.
#[derive(Clone)]
pub enum SessionTransformContext {
    Http(Box<AuthRequest>),
    Endpoint(Box<EndpointCall>),
}
impl SessionTransformContext {
    /// A genuine original Request; header-only SDK calls retain absence.
    #[must_use]
    pub fn request(&self) -> Option<&AuthRequest> {
        match self {
            Self::Http(request) => Some(request),
            Self::Endpoint(call) => call.request(),
        }
    }
    #[must_use]
    pub fn endpoint_call(&self) -> Option<&EndpointCall> {
        match self {
            Self::Endpoint(call) => Some(call),
            Self::Http(_) => None,
        }
    }
    #[must_use]
    pub fn headers(&self) -> Option<&std::collections::HashMap<String, String>> {
        match self {
            Self::Http(request) => Some(&request.headers),
            Self::Endpoint(call) => call.headers(),
        }
    }
    #[must_use]
    pub fn header(&self, name: &str) -> Option<&String> {
        self.headers()?
            .iter()
            .find_map(|(key, value)| key.eq_ignore_ascii_case(name).then_some(value))
    }
    #[must_use]
    pub fn path(&self) -> &str {
        match self {
            Self::Http(request) => request.path(),
            Self::Endpoint(call) => call.path().unwrap_or(""),
        }
    }
    #[must_use]
    pub fn method(&self) -> &HttpMethod {
        match self {
            Self::Http(request) => request.method(),
            Self::Endpoint(call) => call.method().unwrap_or(&HttpMethod::Get),
        }
    }
}
struct SessionTransformHandle<S: AuthSchema>(Arc<dyn SessionTransform<S>>);

/// Project an authenticated session response into application response data.
///
/// The input is the actual public session/user projection, including configured
/// fields and deferred-refresh information. Transforming it does not change the
/// stored principal or the authentication context used by subsequent plugins.
#[async_trait]
pub trait SessionTransform<S: AuthSchema>: Send + Sync {
    /// # Errors
    /// Return an intentional API error or an application callback failure.
    async fn transform(
        &self,
        session: Value,
        request: &SessionTransformContext,
        context: &AuthContext<S>,
    ) -> AuthResult<Value>;
}

/// Replace GET session responses while retaining core authentication and cookies.
/// The replacement owns the public getSession API and accepts GET only on HTTP, including
/// when deferred refresh is enabled.
pub struct CustomSessionPlugin<S: AuthSchema> {
    transform: Arc<dyn SessionTransform<S>>,
    mutate_device_sessions: bool,
}

impl<S: AuthSchema> CustomSessionPlugin<S> {
    /// Read the application-projected session through the public server API.
    #[must_use]
    pub const fn get_session_endpoint() -> ServerEndpoint<Value> {
        ServerEndpoint::new("custom-session", "getSession")
    }

    #[must_use]
    pub fn new(transform: impl SessionTransform<S> + 'static) -> Self {
        Self {
            transform: Arc::new(transform),
            mutate_device_sessions: false,
        }
    }

    /// Apply the same projection to each genuine multiple-device-session entry.
    #[must_use]
    pub fn mutate_device_sessions(mut self, enabled: bool) -> Self {
        self.mutate_device_sessions = enabled;
        self
    }
}

#[async_trait]
impl<S: AuthSchema> AuthPlugin<S> for CustomSessionPlugin<S> {
    fn static_openapi_metadata(&self) -> alibi_core::PluginOpenApiMetadata {
        crate::metadata::plugin_metadata(
            <Self as alibi_core::AuthPlugin<S>>::name(self),
            &<Self as alibi_core::AuthPlugin<S>>::routes(self),
        )
    }

    async fn on_init(&self, ctx: &mut AuthInitContext<S>) -> AuthResult<()> {
        ctx.extensions
            .insert(SessionTransformHandle(self.transform.clone()));
        Ok(())
    }
    fn endpoint_hooks(&self) -> Vec<&dyn EndpointHook<S>> {
        vec![self]
    }
    fn server_endpoints(&self) -> Vec<EndpointDefinition> {
        vec![crate::endpoint::definition(
            "getSession",
            "get_session",
            Some("/get-session"),
            HttpMethod::Get,
        )]
    }
    fn validate_endpoint(
        &self,
        call: &EndpointCall,
        ctx: &AuthContext<S>,
    ) -> AuthResult<EndpointInput> {
        <super::session_management::SessionManagementPlugin as AuthPlugin<S>>::validate_endpoint(
            &super::session_management::SessionManagementPlugin::new(),
            call,
            ctx,
        )
    }
    async fn on_endpoint(
        &self,
        call: &EndpointCall,
        ctx: &AuthContext<S>,
    ) -> AuthResult<EndpointResponse> {
        project_endpoint(call, ctx).await
    }
    fn name(&self) -> &'static str {
        "custom-session"
    }
    fn routes(&self) -> Vec<AuthRoute> {
        vec![
            AuthRoute::get("/get-session", "get_session"),
            AuthRoute::post("/get-session", "custom_session_post"),
        ]
    }

    fn openapi_metadata(
        &self,
        ctx: &alibi_core::AuthInitContext<S>,
    ) -> alibi_core::openapi::PluginOpenApiMetadata {
        let mut metadata =
            crate::metadata::instance_plugin_metadata("session-management", &self.routes(), ctx);
        // Reuse the base schemas without treating this replacement as a core plugin.
        metadata.core = false;
        // Pinned documentation retains the base endpoint and excludes the
        // replacement that shares its logical API key, despite GET-only HTTP.
        for (_, _, endpoint) in &mut metadata.endpoints {
            endpoint.native_extension = true;
        }
        metadata
    }

    async fn on_request(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<S>,
    ) -> AuthResult<Option<AuthResponse>> {
        if req.path() != "/get-session" {
            return Ok(None);
        }
        if req.method() == &HttpMethod::Post {
            return Ok(Some(AuthResponse::new(404)));
        }
        if req.method() != &HttpMethod::Get {
            return Ok(None);
        }
        Ok(Some(project_http(req, ctx).await?))
    }

    async fn after_request(
        &self,
        req: &AuthRequest,
        ctx: &AuthContext<S>,
        mut response: AuthResponse,
    ) -> AuthResult<AuthResponse> {
        if !self.mutate_device_sessions
            || req.path() != "/multi-session/list-device-sessions"
            || response.status >= 400
        {
            return Ok(response);
        }
        let value: Value = serde_json::from_slice(&response.body)?;
        if let Value::Array(sessions) = value {
            let transformed = self
                .transform_list(
                    sessions,
                    &SessionTransformContext::Http(Box::new(req.clone())),
                    ctx,
                )
                .await?;
            response.body = serde_json::to_vec(&transformed)?;
        }
        Ok(response)
    }
}

#[async_trait]
impl<S: AuthSchema> EndpointHook<S> for CustomSessionPlugin<S> {
    fn matches_after(
        &self,
        call: &EndpointCall,
        _ctx: &AuthContext<S>,
        _response: &EndpointResponse,
    ) -> AuthResult<bool> {
        Ok(self.mutate_device_sessions
            && call.path() == Some("/multi-session/list-device-sessions"))
    }
    async fn after(
        &self,
        call: &EndpointCall,
        ctx: &AuthContext<S>,
        mut response: EndpointResponse,
    ) -> AuthResult<EndpointResponse> {
        if let Ok(alibi_core::utils::json::JsValue::Array(sessions)) = response.result() {
            let sessions = sessions
                .iter()
                .map(alibi_core::utils::json::JsValue::to_json_value)
                .collect::<Result<Vec<_>, _>>()?;
            let transformed = self
                .transform_list(
                    sessions,
                    &SessionTransformContext::Endpoint(Box::new(call.clone())),
                    ctx,
                )
                .await?;
            response.replace(Value::Array(transformed).into());
        }
        Ok(response)
    }
}

impl<S: AuthSchema> CustomSessionPlugin<S> {
    async fn transform_list(
        &self,
        sessions: Vec<Value>,
        request: &SessionTransformContext,
        context: &AuthContext<S>,
    ) -> AuthResult<Vec<Value>> {
        if sessions.is_empty() {
            return Ok(Vec::new());
        }
        let mut results = vec![None; sessions.len()];
        let transform = self.transform.clone();
        let request = request.clone();
        let context = AuthContext::<S> {
            config: context.config.clone(),
            database: context.database.clone(),
            email_provider: context.email_provider.clone(),
            metadata: context.metadata.clone(),
            extensions: context.extensions.clone(),
        };
        let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
        // Launched callbacks retain ownership after the aggregate rejects,
        // matching the independent application work of a device-session list.
        crate::helpers::spawn_in_request_context(async move {
            _ = futures_util::future::join_all(sessions.into_iter().enumerate().map(
                |(index, session)| {
                    let transform = &transform;
                    let request = &request;
                    let context = &context;
                    let sender = &sender;
                    async move {
                        let result = transform
                            .transform(session, request, context)
                            .await
                            .map_err(crate::helpers::callback_failure);
                        _ = sender.send((index, result));
                    }
                },
            ))
            .await;
        });
        for _ in 0..results.len() {
            let (index, value) = receiver
                .recv()
                .await
                .ok_or_else(|| alibi_core::AuthError::internal("Session transformation stopped"))?;
            let slot = results.get_mut(index).ok_or_else(|| {
                alibi_core::AuthError::internal("Invalid session transformation row")
            })?;
            *slot = Some(value?);
        }
        Ok(results.into_iter().flatten().collect())
    }
}

pub(crate) fn configured<S: AuthSchema>(ctx: &AuthContext<S>) -> bool {
    ctx.extensions.get::<SessionTransformHandle<S>>().is_some()
}
pub(crate) async fn project_http<S: AuthSchema>(
    req: &AuthRequest,
    ctx: &AuthContext<S>,
) -> AuthResult<AuthResponse> {
    if req.method() != &HttpMethod::Get {
        return Ok(AuthResponse::new(404));
    }
    let transform = ctx
        .extensions
        .get::<SessionTransformHandle<S>>()
        .ok_or_else(|| AuthError::internal("Custom session projection is not installed"))?;
    let transform = &transform.0;
    let prior_headers = req.take_response_headers();
    let mut response = match super::session_management::SessionManagementPlugin::new()
        .handle_base_get_session(req, ctx)
        .await
    {
        Ok(response) => response,
        Err(_) => AuthResponse::json(200, &Value::Null)?,
    };
    let session: Value = if response.status >= 400 {
        Value::Null
    } else {
        serde_json::from_slice(&response.body)?
    };
    let core_headers = req.take_response_headers();
    for (name, value) in prior_headers {
        req.queue_response_header(name, value);
    }
    let (issued_cookies, _) =
        alibi_core::session::cookie_cache::runtime::take_issuance(req.extensions());
    if session.is_null() {
        // The replaced endpoint does not forward headers from a missing or
        // failed core response. Preserve headers emitted before that read.
        return Ok(AuthResponse::json(200, &Value::Null)?);
    }
    let transformed = transform
        .transform(
            session,
            &SessionTransformContext::Http(Box::new(req.clone())),
            ctx,
        )
        .await
        .map_err(crate::helpers::callback_failure)?;
    response.body = serde_json::to_vec(&transformed)?;
    for (name, value) in core_headers {
        response.headers.append(name, value);
    }
    for cookie in issued_cookies {
        response.headers.append("Set-Cookie", cookie);
    }
    Ok(response)
}
pub(crate) async fn project_endpoint<S: AuthSchema>(
    call: &EndpointCall,
    ctx: &AuthContext<S>,
) -> AuthResult<EndpointResponse> {
    let transform = ctx
        .extensions
        .get::<SessionTransformHandle<S>>()
        .ok_or_else(|| AuthError::internal("Custom session projection is not installed"))?;
    let prior = call.take_response_headers();
    let mut read = call.clone();
    EndpointContextPatch {
        method: Some(HttpMethod::Get),
        ..Default::default()
    }
    .apply(&mut read);
    let session = super::session_management::SessionManagementPlugin::new()
        .get_session_response(&read, ctx)
        .await
        .ok()
        .flatten();
    let core_headers = read.take_response_headers();
    for (name, value) in prior {
        call.queue_response_header(name, value);
    }
    let (cookies, _) = alibi_core::session::cookie_cache::runtime::take_issuance(read.extensions());
    let Some(session) = session else {
        return EndpointResponse::json(&Value::Null);
    };
    let value = transform
        .0
        .transform(
            serde_json::to_value(session)?,
            &SessionTransformContext::Endpoint(Box::new(call.clone())),
            ctx,
        )
        .await
        .map_err(crate::helpers::callback_failure)?;
    let mut response = EndpointResponse::json(&value)?
        .with_header("cache-control", "no-store")
        .with_header("pragma", "no-cache");
    response.merge_headers(core_headers);
    for cookie in cookies {
        response = response.with_header("set-cookie", cookie);
    }
    Ok(response)
}

// LCOV_EXCL_START
#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::{
        create_auth_request_no_query, create_test_context, create_user_and_session,
    };
    use alibi_core::AuthUser;
    use alibi_core::hooks::with_request_hook_context;
    use alibi_core::{CreateUser, UpdateUser};
    use std::sync::Arc;
    use tokio::sync::Notify;
    type Schema = alibi_seaorm::store::__private_test_support::bundled_schema::BundledSchema;

    struct Transform {
        entered: Arc<Notify>,
        release: Arc<Notify>,
        finished: Arc<Notify>,
    }
    #[async_trait]
    impl SessionTransform<Schema> for Transform {
        async fn transform(
            &self,
            value: Value,
            req: &SessionTransformContext,
            ctx: &AuthContext<Schema>,
        ) -> AuthResult<Value> {
            let id = value
                .pointer("/user/id")
                .and_then(Value::as_str)
                .expect("actual user id");
            let name = value
                .pointer("/user/name")
                .and_then(Value::as_str)
                .expect("actual user name");
            if name == "Reject" {
                self.entered.notified().await;
                return Err(alibi_core::AuthError::internal("application failure"));
            }
            self.entered.notify_one();
            self.release.notified().await;
            assert_eq!(req.path(), "/multi-session/list-device-sessions");
            assert_eq!(
                alibi_core::hooks::current_request_hook_context()
                    .expect("retained real request")
                    .headers
                    .get("x-application"),
                Some(&"original".to_owned())
            );
            ctx.database
                .update_user(
                    id,
                    UpdateUser {
                        name: Some("Completed".into()),
                        ..Default::default()
                    },
                )
                .await?;
            self.finished.notify_one();
            Ok(value)
        }
    }

    #[tokio::test]
    async fn rejected_device_list_preserves_launched_callback_request_and_database_ownership() {
        let ctx = create_test_context().await;
        let (reject, reject_session) = create_user_and_session(
            &ctx,
            CreateUser::new()
                .with_email("reject@custom.fixture.test")
                .with_name("Reject"),
            chrono::Duration::hours(1),
        )
        .await;
        let (slow, slow_session) = create_user_and_session(
            &ctx,
            CreateUser::new()
                .with_email("slow@custom.fixture.test")
                .with_name("Slow"),
            chrono::Duration::hours(1),
        )
        .await;
        let entered = Arc::new(Notify::new());
        let release = Arc::new(Notify::new());
        let finished = Arc::new(Notify::new());
        let plugin = CustomSessionPlugin::new(Transform {
            entered,
            release: release.clone(),
            finished: finished.clone(),
        })
        .mutate_device_sessions(true);
        let mut req = create_auth_request_no_query(
            HttpMethod::Get,
            "/multi-session/list-device-sessions",
            Some(&reject_session.token),
            None,
        );
        _ = req
            .headers
            .insert("x-application".into(), "original".into());
        let cookies = [&reject_session.token, &slow_session.token]
            .into_iter()
            .map(|token| {
                format!(
                    "{}_multi-{}={}",
                    ctx.config.session.cookie_name,
                    token.to_lowercase(),
                    alibi_core::utils::cookie_utils::sign_cookie_value(
                        token,
                        ctx.config.current_secret()
                    )
                )
            })
            .collect::<Vec<_>>();
        let ordinary = req.header("cookie").expect("signed current cookie").clone();
        _ = req.headers.insert(
            "cookie".into(),
            format!("{ordinary}; {}", cookies.join("; ")),
        );
        let response = crate::multi_session::MultiSessionPlugin::new()
            .on_request(&req, &ctx)
            .await
            .expect("real list")
            .expect("list route");
        assert_eq!(response.status, 200);
        let result = tokio::time::timeout(
            std::time::Duration::from_secs(2),
            with_request_hook_context(&req, plugin.after_request(&req, &ctx, response)),
        )
        .await
        .expect("aggregate rejects without waiting for held callback");
        assert!(matches!(
            result,
            Err(alibi_core::AuthError::CallbackFailure(_))
        ));
        assert_eq!(
            ctx.database
                .get_user_by_id(&slow.id)
                .await
                .expect("read")
                .expect("stored slow owner")
                .name(),
            Some("Slow")
        );
        release.notify_one();
        tokio::time::timeout(std::time::Duration::from_secs(2), finished.notified())
            .await
            .expect("launched callback completes after observer rejection");
        assert_eq!(
            ctx.database
                .get_user_by_id(&slow.id)
                .await
                .expect("read")
                .expect("stored slow owner")
                .name(),
            Some("Completed")
        );
        assert_eq!(
            ctx.database
                .get_user_by_id(&reject.id)
                .await
                .expect("read")
                .expect("stored reject owner")
                .name(),
            Some("Reject")
        );
    }
}
// LCOV_EXCL_STOP