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 = super::session_management::SessionManagementPlugin::new()
333        .handle_base_get_session(req, ctx)
334        .await?;
335    let core_headers = req.take_response_headers();
336    for (name, value) in prior_headers {
337        req.queue_response_header(name, value);
338    }
339    let (issued_cookies, _) =
340        alibi_core::session::cookie_cache::runtime::take_issuance(req.extensions());
341    if response.status >= 400 {
342        // Since 1.7.12 a failed core read propagates. The thrown error leaves
343        // the core read's headers behind, like the replaced endpoint's.
344        let mut failure = AuthResponse::new(response.status);
345        failure.body = response.body;
346        if let Some(content_type) = response.headers.get("content-type") {
347            failure = failure.with_header("content-type", content_type.to_owned());
348        }
349        return Ok(failure);
350    }
351    let session: Value = serde_json::from_slice(&response.body)?;
352    // A null session still forwards the core read's headers and cookie cleanup.
353    if !session.is_null() {
354        let transformed = transform
355            .transform(
356                session,
357                &SessionTransformContext::Http(Box::new(req.clone())),
358                ctx,
359            )
360            .await
361            .map_err(crate::helpers::callback_failure)?;
362        response.body = serde_json::to_vec(&transformed)?;
363    }
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    let core_headers = read.take_response_headers();
391    for (name, value) in prior {
392        call.queue_response_header(name, value);
393    }
394    let (cookies, _) = alibi_core::session::cookie_cache::runtime::take_issuance(read.extensions());
395    // A null session still forwards the core read's headers and cookie cleanup.
396    let value = match session {
397        Some(session) => transform
398            .0
399            .transform(
400                serde_json::to_value(session)?,
401                &SessionTransformContext::Endpoint(Box::new(call.clone())),
402                ctx,
403            )
404            .await
405            .map_err(crate::helpers::callback_failure)?,
406        None => Value::Null,
407    };
408    let mut response = EndpointResponse::json(&value)?
409        .with_header("cache-control", "no-store")
410        .with_header("pragma", "no-cache");
411    response.merge_headers(core_headers);
412    for cookie in cookies {
413        response = response.with_header("set-cookie", cookie);
414    }
415    Ok(response)
416}
417
418// LCOV_EXCL_START
419#[cfg(test)]
420mod tests {
421    use super::*;
422    use crate::test_helpers::{
423        create_auth_request_no_query, create_test_context, create_user_and_session,
424    };
425    use alibi_core::AuthUser;
426    use alibi_core::hooks::with_request_hook_context;
427    use alibi_core::{CreateUser, UpdateUser};
428    use std::sync::Arc;
429    use tokio::sync::Notify;
430    type Schema = alibi_seaorm::store::__private_test_support::bundled_schema::BundledSchema;
431
432    struct Transform {
433        entered: Arc<Notify>,
434        release: Arc<Notify>,
435        finished: Arc<Notify>,
436    }
437    #[async_trait]
438    impl SessionTransform<Schema> for Transform {
439        async fn transform(
440            &self,
441            value: Value,
442            req: &SessionTransformContext,
443            ctx: &AuthContext<Schema>,
444        ) -> AuthResult<Value> {
445            let id = value
446                .pointer("/user/id")
447                .and_then(Value::as_str)
448                .expect("actual user id");
449            let name = value
450                .pointer("/user/name")
451                .and_then(Value::as_str)
452                .expect("actual user name");
453            if name == "Reject" {
454                self.entered.notified().await;
455                return Err(alibi_core::AuthError::internal("application failure"));
456            }
457            self.entered.notify_one();
458            self.release.notified().await;
459            assert_eq!(req.path(), "/multi-session/list-device-sessions");
460            assert_eq!(
461                alibi_core::hooks::current_request_hook_context()
462                    .expect("retained real request")
463                    .headers
464                    .get("x-application"),
465                Some(&"original".to_owned())
466            );
467            ctx.database
468                .update_user(
469                    id,
470                    UpdateUser {
471                        name: Some("Completed".into()),
472                        ..Default::default()
473                    },
474                )
475                .await?;
476            self.finished.notify_one();
477            Ok(value)
478        }
479    }
480
481    #[tokio::test]
482    async fn rejected_device_list_preserves_launched_callback_request_and_database_ownership() {
483        let ctx = create_test_context().await;
484        let (reject, reject_session) = create_user_and_session(
485            &ctx,
486            CreateUser::new()
487                .with_email("reject@custom.fixture.test")
488                .with_name("Reject"),
489            chrono::Duration::hours(1),
490        )
491        .await;
492        let (slow, slow_session) = create_user_and_session(
493            &ctx,
494            CreateUser::new()
495                .with_email("slow@custom.fixture.test")
496                .with_name("Slow"),
497            chrono::Duration::hours(1),
498        )
499        .await;
500        let entered = Arc::new(Notify::new());
501        let release = Arc::new(Notify::new());
502        let finished = Arc::new(Notify::new());
503        let plugin = CustomSessionPlugin::new(Transform {
504            entered,
505            release: release.clone(),
506            finished: finished.clone(),
507        })
508        .mutate_device_sessions(true);
509        let mut req = create_auth_request_no_query(
510            HttpMethod::Get,
511            "/multi-session/list-device-sessions",
512            Some(&reject_session.token),
513            None,
514        );
515        _ = req
516            .headers
517            .insert("x-application".into(), "original".into());
518        let cookies = [&reject_session.token, &slow_session.token]
519            .into_iter()
520            .map(|token| {
521                format!(
522                    "{}_multi-{}={}",
523                    ctx.config.session.cookie_name,
524                    token.to_lowercase(),
525                    alibi_core::utils::cookie_utils::sign_cookie_value(
526                        token,
527                        ctx.config.current_secret()
528                    )
529                )
530            })
531            .collect::<Vec<_>>();
532        let ordinary = req.header("cookie").expect("signed current cookie").clone();
533        _ = req.headers.insert(
534            "cookie".into(),
535            format!("{ordinary}; {}", cookies.join("; ")),
536        );
537        let response = crate::multi_session::MultiSessionPlugin::new()
538            .on_request(&req, &ctx)
539            .await
540            .expect("real list")
541            .expect("list route");
542        assert_eq!(response.status, 200);
543        let result = tokio::time::timeout(
544            std::time::Duration::from_secs(2),
545            with_request_hook_context(&req, plugin.after_request(&req, &ctx, response)),
546        )
547        .await
548        .expect("aggregate rejects without waiting for held callback");
549        assert!(matches!(
550            result,
551            Err(alibi_core::AuthError::CallbackFailure(_))
552        ));
553        assert_eq!(
554            ctx.database
555                .get_user_by_id(&slow.id)
556                .await
557                .expect("read")
558                .expect("stored slow owner")
559                .name(),
560            Some("Slow")
561        );
562        release.notify_one();
563        tokio::time::timeout(std::time::Duration::from_secs(2), finished.notified())
564            .await
565            .expect("launched callback completes after observer rejection");
566        assert_eq!(
567            ctx.database
568                .get_user_by_id(&slow.id)
569                .await
570                .expect("read")
571                .expect("stored slow owner")
572                .name(),
573            Some("Completed")
574        );
575        assert_eq!(
576            ctx.database
577                .get_user_by_id(&reject.id)
578                .await
579                .expect("read")
580                .expect("stored reject owner")
581                .name(),
582            Some("Reject")
583        );
584    }
585}
586// LCOV_EXCL_STOP
alibi_plugins/custom_session.rs

alibi_plugins/
custom_session.rs
