alibi_plugins/
custom_session.rs

1//! Application-owned asynchronous session response projection.
2use alibi_core::endpoint::{
3    EndpointCall, EndpointContextPatch, EndpointDefinition, EndpointHook, EndpointInput,
4    EndpointResponse, ServerEndpoint,
5};
6use alibi_core::{
7    AuthContext, AuthError, AuthInitContext, AuthPlugin, AuthRequest, AuthResponse, AuthResult,
8    AuthRoute, AuthSchema, HttpMethod,
9};
10use async_trait::async_trait;
11use serde_json::Value;
12use std::sync::Arc;
13/// The actual transport frame supplied to a session projection callback.
14#[derive(Clone)]
15pub enum SessionTransformContext {
16    Http(Box<AuthRequest>),
17    Endpoint(Box<EndpointCall>),
18}
19impl SessionTransformContext {
20    /// A genuine original Request; header-only SDK calls retain absence.
21    #[must_use]
22    pub fn request(&self) -> Option<&AuthRequest> {
23        match self {
24            Self::Http(request) => Some(request),
25            Self::Endpoint(call) => call.request(),
26        }
27    }
28    #[must_use]
29    pub fn endpoint_call(&self) -> Option<&EndpointCall> {
30        match self {
31            Self::Endpoint(call) => Some(call),
32            Self::Http(_) => None,
33        }
34    }
35    #[must_use]
36    pub fn headers(&self) -> Option<&std::collections::HashMap<String, String>> {
37        match self {
38            Self::Http(request) => Some(&request.headers),
39            Self::Endpoint(call) => call.headers(),
40        }
41    }
42    #[must_use]
43    pub fn header(&self, name: &str) -> Option<&String> {
44        self.headers()?
45            .iter()
46            .find_map(|(key, value)| key.eq_ignore_ascii_case(name).then_some(value))
47    }
48    #[must_use]
49    pub fn path(&self) -> &str {
50        match self {
51            Self::Http(request) => request.path(),
52            Self::Endpoint(call) => call.path().unwrap_or(""),
53        }
54    }
55    #[must_use]
56    pub fn method(&self) -> &HttpMethod {
57        match self {
58            Self::Http(request) => request.method(),
59            Self::Endpoint(call) => call.method().unwrap_or(&HttpMethod::Get),
60        }
61    }
62}
63struct SessionTransformHandle<S: AuthSchema>(Arc<dyn SessionTransform<S>>);
64
65/// Project an authenticated session response into application response data.
66///
67/// The input is the actual public session/user projection, including configured
68/// fields and deferred-refresh information. Transforming it does not change the
69/// stored principal or the authentication context used by subsequent plugins.
70#[async_trait]
71pub trait SessionTransform<S: AuthSchema>: Send + Sync {
72    /// # Errors
73    /// Return an intentional API error or an application callback failure.
74    async fn transform(
75        &self,
76        session: Value,
77        request: &SessionTransformContext,
78        context: &AuthContext<S>,
79    ) -> AuthResult<Value>;
80}
81
82/// Replace GET session responses while retaining core authentication and cookies.
83/// The replacement owns the public getSession API and accepts GET only on HTTP, including
84/// when deferred refresh is enabled.
85pub struct CustomSessionPlugin<S: AuthSchema> {
86    transform: Arc<dyn SessionTransform<S>>,
87    mutate_device_sessions: bool,
88}
89
90impl<S: AuthSchema> CustomSessionPlugin<S> {
91    /// Read the application-projected session through the public server API.
92    #[must_use]
93    pub const fn get_session_endpoint() -> ServerEndpoint<Value> {
94        ServerEndpoint::new("custom-session", "getSession")
95    }
96
97    #[must_use]
98    pub fn new(transform: impl SessionTransform<S> + 'static) -> Self {
99        Self {
100            transform: Arc::new(transform),
101            mutate_device_sessions: false,
102        }
103    }
104
105    /// Apply the same projection to each genuine multiple-device-session entry.
106    #[must_use]
107    pub fn mutate_device_sessions(mut self, enabled: bool) -> Self {
108        self.mutate_device_sessions = enabled;
109        self
110    }
111}
112
113#[async_trait]
114impl<S: AuthSchema> AuthPlugin<S> for CustomSessionPlugin<S> {
115    fn static_openapi_metadata(&self) -> alibi_core::PluginOpenApiMetadata {
116        crate::metadata::plugin_metadata(
117            <Self as alibi_core::AuthPlugin<S>>::name(self),
118            &<Self as alibi_core::AuthPlugin<S>>::routes(self),
119        )
120    }
121
122    async fn on_init(&self, ctx: &mut AuthInitContext<S>) -> AuthResult<()> {
123        ctx.extensions
124            .insert(SessionTransformHandle(self.transform.clone()));
125        Ok(())
126    }
127    fn endpoint_hooks(&self) -> Vec<&dyn EndpointHook<S>> {
128        vec![self]
129    }
130    fn server_endpoints(&self) -> Vec<EndpointDefinition> {
131        vec![crate::endpoint::definition(
132            "getSession",
133            "get_session",
134            Some("/get-session"),
135            HttpMethod::Get,
136        )]
137    }
138    fn validate_endpoint(
139        &self,
140        call: &EndpointCall,
141        ctx: &AuthContext<S>,
142    ) -> AuthResult<EndpointInput> {
143        <super::session_management::SessionManagementPlugin as AuthPlugin<S>>::validate_endpoint(
144            &super::session_management::SessionManagementPlugin::new(),
145            call,
146            ctx,
147        )
148    }
149    async fn on_endpoint(
150        &self,
151        call: &EndpointCall,
152        ctx: &AuthContext<S>,
153    ) -> AuthResult<EndpointResponse> {
154        project_endpoint(call, ctx).await
155    }
156    fn name(&self) -> &'static str {
157        "custom-session"
158    }
159    fn routes(&self) -> Vec<AuthRoute> {
160        vec![
161            AuthRoute::get("/get-session", "get_session"),
162            AuthRoute::post("/get-session", "custom_session_post"),
163        ]
164    }
165
166    fn openapi_metadata(
167        &self,
168        ctx: &alibi_core::AuthInitContext<S>,
169    ) -> alibi_core::openapi::PluginOpenApiMetadata {
170        let mut metadata =
171            crate::metadata::instance_plugin_metadata("session-management", &self.routes(), ctx);
172        // Reuse the base schemas without treating this replacement as a core plugin.
173        metadata.core = false;
174        // Pinned documentation retains the base endpoint and excludes the
175        // replacement that shares its logical API key, despite GET-only HTTP.
176        for (_, _, endpoint) in &mut metadata.endpoints {
177            endpoint.native_extension = true;
178        }
179        metadata
180    }
181
182    async fn on_request(
183        &self,
184        req: &AuthRequest,
185        ctx: &AuthContext<S>,
186    ) -> AuthResult<Option<AuthResponse>> {
187        if req.path() != "/get-session" {
188            return Ok(None);
189        }
190        if req.method() == &HttpMethod::Post {
191            return Ok(Some(AuthResponse::new(404)));
192        }
193        if req.method() != &HttpMethod::Get {
194            return Ok(None);
195        }
196        Ok(Some(project_http(req, ctx).await?))
197    }
198
199    async fn after_request(
200        &self,
201        req: &AuthRequest,
202        ctx: &AuthContext<S>,
203        mut response: AuthResponse,
204    ) -> AuthResult<AuthResponse> {
205        if !self.mutate_device_sessions
206            || req.path() != "/multi-session/list-device-sessions"
207            || response.status >= 400
208        {
209            return Ok(response);
210        }
211        let value: Value = serde_json::from_slice(&response.body)?;
212        if let Value::Array(sessions) = value {
213            let transformed = self
214                .transform_list(
215                    sessions,
216                    &SessionTransformContext::Http(Box::new(req.clone())),
217                    ctx,
218                )
219                .await?;
220            response.body = serde_json::to_vec(&transformed)?;
221        }
222        Ok(response)
223    }
224}
225
226#[async_trait]
227impl<S: AuthSchema> EndpointHook<S> for CustomSessionPlugin<S> {
228    fn matches_after(
229        &self,
230        call: &EndpointCall,
231        _ctx: &AuthContext<S>,
232        _response: &EndpointResponse,
233    ) -> AuthResult<bool> {
234        Ok(self.mutate_device_sessions
235            && call.path() == Some("/multi-session/list-device-sessions"))
236    }
237    async fn after(
238        &self,
239        call: &EndpointCall,
240        ctx: &AuthContext<S>,
241        mut response: EndpointResponse,
242    ) -> AuthResult<EndpointResponse> {
243        if let Ok(alibi_core::utils::json::JsValue::Array(sessions)) = response.result() {
244            let sessions = sessions
245                .iter()
246                .map(alibi_core::utils::json::JsValue::to_json_value)
247                .collect::<Result<Vec<_>, _>>()?;
248            let transformed = self
249                .transform_list(
250                    sessions,
251                    &SessionTransformContext::Endpoint(Box::new(call.clone())),
252                    ctx,
253                )
254                .await?;
255            response.replace(Value::Array(transformed).into());
256        }
257        Ok(response)
258    }
259}
260
261impl<S: AuthSchema> CustomSessionPlugin<S> {
262    async fn transform_list(
263        &self,
264        sessions: Vec<Value>,
265        request: &SessionTransformContext,
266        context: &AuthContext<S>,
267    ) -> AuthResult<Vec<Value>> {
268        if sessions.is_empty() {
269            return Ok(Vec::new());
270        }
271        let mut results = vec![None; sessions.len()];
272        let transform = self.transform.clone();
273        let request = request.clone();
274        let context = AuthContext::<S> {
275            config: context.config.clone(),
276            database: context.database.clone(),
277            email_provider: context.email_provider.clone(),
278            metadata: context.metadata.clone(),
279            extensions: context.extensions.clone(),
280        };
281        let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
282        // Launched callbacks retain ownership after the aggregate rejects,
283        // matching the independent application work of a device-session list.
284        crate::helpers::spawn_in_request_context(async move {
285            _ = futures_util::future::join_all(sessions.into_iter().enumerate().map(
286                |(index, session)| {
287                    let transform = &transform;
288                    let request = &request;
289                    let context = &context;
290                    let sender = &sender;
291                    async move {
292                        let result = transform
293                            .transform(session, request, context)
294                            .await
295                            .map_err(crate::helpers::callback_failure);
296                        _ = sender.send((index, result));
297                    }
298                },
299            ))
300            .await;
301        });
302        for _ in 0..results.len() {
303            let (index, value) = receiver
304                .recv()
305                .await
306                .ok_or_else(|| alibi_core::AuthError::internal("Session transformation stopped"))?;
307            let slot = results.get_mut(index).ok_or_else(|| {
308                alibi_core::AuthError::internal("Invalid session transformation row")
309            })?;
310            *slot = Some(value?);
311        }
312        Ok(results.into_iter().flatten().collect())
313    }
314}
315
316pub(crate) fn configured<S: AuthSchema>(ctx: &AuthContext<S>) -> bool {
317    ctx.extensions.get::<SessionTransformHandle<S>>().is_some()
318}
319pub(crate) async fn project_http<S: AuthSchema>(
320    req: &AuthRequest,
321    ctx: &AuthContext<S>,
322) -> AuthResult<AuthResponse> {
323    if req.method() != &HttpMethod::Get {
324        return Ok(AuthResponse::new(404));
325    }
326    let transform = ctx
327        .extensions
328        .get::<SessionTransformHandle<S>>()
329        .ok_or_else(|| AuthError::internal("Custom session projection is not installed"))?;
330    let transform = &transform.0;
331    let prior_headers = req.take_response_headers();
332    let mut response = match super::session_management::SessionManagementPlugin::new()
333        .handle_base_get_session(req, ctx)
334        .await
335    {
336        Ok(response) => response,
337        Err(_) => AuthResponse::json(200, &Value::Null)?,
338    };
339    let session: Value = if response.status >= 400 {
340        Value::Null
341    } else {
342        serde_json::from_slice(&response.body)?
343    };
344    let core_headers = req.take_response_headers();
345    for (name, value) in prior_headers {
346        req.queue_response_header(name, value);
347    }
348    let (issued_cookies, _) =
349        alibi_core::session::cookie_cache::runtime::take_issuance(req.extensions());
350    if session.is_null() {
351        // The replaced endpoint does not forward headers from a missing or
352        // failed core response. Preserve headers emitted before that read.
353        return Ok(AuthResponse::json(200, &Value::Null)?);
354    }
355    let transformed = transform
356        .transform(
357            session,
358            &SessionTransformContext::Http(Box::new(req.clone())),
359            ctx,
360        )
361        .await
362        .map_err(crate::helpers::callback_failure)?;
363    response.body = serde_json::to_vec(&transformed)?;
364    for (name, value) in core_headers {
365        response.headers.append(name, value);
366    }
367    for cookie in issued_cookies {
368        response.headers.append("Set-Cookie", cookie);
369    }
370    Ok(response)
371}
372pub(crate) async fn project_endpoint<S: AuthSchema>(
373    call: &EndpointCall,
374    ctx: &AuthContext<S>,
375) -> AuthResult<EndpointResponse> {
376    let transform = ctx
377        .extensions
378        .get::<SessionTransformHandle<S>>()
379        .ok_or_else(|| AuthError::internal("Custom session projection is not installed"))?;
380    let prior = call.take_response_headers();
381    let mut read = call.clone();
382    EndpointContextPatch {
383        method: Some(HttpMethod::Get),
384        ..Default::default()
385    }
386    .apply(&mut read);
387    let session = super::session_management::SessionManagementPlugin::new()
388        .get_session_response(&read, ctx)
389        .await
390        .ok()
391        .flatten();
392    let core_headers = read.take_response_headers();
393    for (name, value) in prior {
394        call.queue_response_header(name, value);
395    }
396    let (cookies, _) = alibi_core::session::cookie_cache::runtime::take_issuance(read.extensions());
397    let Some(session) = session else {
398        return EndpointResponse::json(&Value::Null);
399    };
400    let value = transform
401        .0
402        .transform(
403            serde_json::to_value(session)?,
404            &SessionTransformContext::Endpoint(Box::new(call.clone())),
405            ctx,
406        )
407        .await
408        .map_err(crate::helpers::callback_failure)?;
409    let mut response = EndpointResponse::json(&value)?
410        .with_header("cache-control", "no-store")
411        .with_header("pragma", "no-cache");
412    response.merge_headers(core_headers);
413    for cookie in cookies {
414        response = response.with_header("set-cookie", cookie);
415    }
416    Ok(response)
417}
418
419// LCOV_EXCL_START
420#[cfg(test)]
421mod tests {
422    use super::*;
423    use crate::test_helpers::{
424        create_auth_request_no_query, create_test_context, create_user_and_session,
425    };
426    use alibi_core::AuthUser;
427    use alibi_core::hooks::with_request_hook_context;
428    use alibi_core::{CreateUser, UpdateUser};
429    use std::sync::Arc;
430    use tokio::sync::Notify;
431    type Schema = alibi_seaorm::store::__private_test_support::bundled_schema::BundledSchema;
432
433    struct Transform {
434        entered: Arc<Notify>,
435        release: Arc<Notify>,
436        finished: Arc<Notify>,
437    }
438    #[async_trait]
439    impl SessionTransform<Schema> for Transform {
440        async fn transform(
441            &self,
442            value: Value,
443            req: &SessionTransformContext,
444            ctx: &AuthContext<Schema>,
445        ) -> AuthResult<Value> {
446            let id = value
447                .pointer("/user/id")
448                .and_then(Value::as_str)
449                .expect("actual user id");
450            let name = value
451                .pointer("/user/name")
452                .and_then(Value::as_str)
453                .expect("actual user name");
454            if name == "Reject" {
455                self.entered.notified().await;
456                return Err(alibi_core::AuthError::internal("application failure"));
457            }
458            self.entered.notify_one();
459            self.release.notified().await;
460            assert_eq!(req.path(), "/multi-session/list-device-sessions");
461            assert_eq!(
462                alibi_core::hooks::current_request_hook_context()
463                    .expect("retained real request")
464                    .headers
465                    .get("x-application"),
466                Some(&"original".to_owned())
467            );
468            ctx.database
469                .update_user(
470                    id,
471                    UpdateUser {
472                        name: Some("Completed".into()),
473                        ..Default::default()
474                    },
475                )
476                .await?;
477            self.finished.notify_one();
478            Ok(value)
479        }
480    }
481
482    #[tokio::test]
483    async fn rejected_device_list_preserves_launched_callback_request_and_database_ownership() {
484        let ctx = create_test_context().await;
485        let (reject, reject_session) = create_user_and_session(
486            &ctx,
487            CreateUser::new()
488                .with_email("reject@custom.fixture.test")
489                .with_name("Reject"),
490            chrono::Duration::hours(1),
491        )
492        .await;
493        let (slow, slow_session) = create_user_and_session(
494            &ctx,
495            CreateUser::new()
496                .with_email("slow@custom.fixture.test")
497                .with_name("Slow"),
498            chrono::Duration::hours(1),
499        )
500        .await;
501        let entered = Arc::new(Notify::new());
502        let release = Arc::new(Notify::new());
503        let finished = Arc::new(Notify::new());
504        let plugin = CustomSessionPlugin::new(Transform {
505            entered,
506            release: release.clone(),
507            finished: finished.clone(),
508        })
509        .mutate_device_sessions(true);
510        let mut req = create_auth_request_no_query(
511            HttpMethod::Get,
512            "/multi-session/list-device-sessions",
513            Some(&reject_session.token),
514            None,
515        );
516        _ = req
517            .headers
518            .insert("x-application".into(), "original".into());
519        let cookies = [&reject_session.token, &slow_session.token]
520            .into_iter()
521            .map(|token| {
522                format!(
523                    "{}_multi-{}={}",
524                    ctx.config.session.cookie_name,
525                    token.to_lowercase(),
526                    alibi_core::utils::cookie_utils::sign_cookie_value(
527                        token,
528                        ctx.config.current_secret()
529                    )
530                )
531            })
532            .collect::<Vec<_>>();
533        let ordinary = req.header("cookie").expect("signed current cookie").clone();
534        _ = req.headers.insert(
535            "cookie".into(),
536            format!("{ordinary}; {}", cookies.join("; ")),
537        );
538        let response = crate::multi_session::MultiSessionPlugin::new()
539            .on_request(&req, &ctx)
540            .await
541            .expect("real list")
542            .expect("list route");
543        assert_eq!(response.status, 200);
544        let result = tokio::time::timeout(
545            std::time::Duration::from_secs(2),
546            with_request_hook_context(&req, plugin.after_request(&req, &ctx, response)),
547        )
548        .await
549        .expect("aggregate rejects without waiting for held callback");
550        assert!(matches!(
551            result,
552            Err(alibi_core::AuthError::CallbackFailure(_))
553        ));
554        assert_eq!(
555            ctx.database
556                .get_user_by_id(&slow.id)
557                .await
558                .expect("read")
559                .expect("stored slow owner")
560                .name(),
561            Some("Slow")
562        );
563        release.notify_one();
564        tokio::time::timeout(std::time::Duration::from_secs(2), finished.notified())
565            .await
566            .expect("launched callback completes after observer rejection");
567        assert_eq!(
568            ctx.database
569                .get_user_by_id(&slow.id)
570                .await
571                .expect("read")
572                .expect("stored slow owner")
573                .name(),
574            Some("Completed")
575        );
576        assert_eq!(
577            ctx.database
578                .get_user_by_id(&reject.id)
579                .await
580                .expect("read")
581                .expect("stored reject owner")
582                .name(),
583            Some("Reject")
584        );
585    }
586}
587// LCOV_EXCL_STOP
alibi_plugins/custom_session.rs

alibi_plugins/
custom_session.rs
