actix_web_flash_messages/
middleware.rs

1use std::cell::RefCell;
2use std::future::Future;
3use std::pin::Pin;
4
5use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
6
7use crate::builder::FlashMessagesFramework;
8use crate::{storage::FlashMessageStore, FlashMessage, Level};
9use actix_web::body::MessageBody;
10use actix_web::HttpMessage;
11use std::sync::Arc;
12
13tokio::task_local! {
14    pub(crate) static OUTGOING_MAILBOX: OutgoingMailbox;
15}
16
17#[derive(Clone)]
18pub(crate) struct OutgoingMailbox {
19    pub(crate) messages: RefCell<Vec<FlashMessage>>,
20    pub(crate) minimum_level: Level,
21}
22
23impl OutgoingMailbox {
24    pub(crate) fn new(minimum_level: Level) -> Self {
25        Self {
26            messages: RefCell::new(vec![]),
27            minimum_level,
28        }
29    }
30}
31
32impl<S, B> Transform<S, ServiceRequest> for FlashMessagesFramework
33where
34    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
35    S::Future: 'static,
36    B: MessageBody + 'static,
37{
38    type Response = ServiceResponse<B>;
39    type Error = actix_web::Error;
40    type Transform = FlashMessagesMiddleware<S>;
41    type InitError = ();
42    type Future = std::future::Ready<Result<Self::Transform, Self::InitError>>;
43
44    fn new_transform(&self, service: S) -> Self::Future {
45        std::future::ready(Ok(FlashMessagesMiddleware {
46            service,
47            storage_backend: self.storage_backend.clone(),
48            minimum_level: self.minimum_level,
49        }))
50    }
51}
52
53#[non_exhaustive]
54#[doc(hidden)]
55pub struct FlashMessagesMiddleware<S> {
56    service: S,
57    storage_backend: Arc<dyn FlashMessageStore>,
58    minimum_level: Level,
59}
60
61#[allow(clippy::type_complexity)]
62impl<S, B> Service<ServiceRequest> for FlashMessagesMiddleware<S>
63where
64    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
65    S::Future: 'static,
66    B: MessageBody + 'static,
67{
68    type Response = ServiceResponse<B>;
69    type Error = actix_web::Error;
70    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
71
72    actix_web::dev::forward_ready!(service);
73
74    fn call(&self, req: ServiceRequest) -> Self::Future {
75        req.extensions_mut().insert(self.storage_backend.clone());
76        let outgoing_mailbox = OutgoingMailbox::new(self.minimum_level);
77        // Working with task-locals in actix-web middlewares is a bit annoying.
78        // We need to make the task local value available to the rest of the middleware chain, which
79        // generates the `future` which will in turn return us a response.
80        // This generation process is synchronous, so we must use `sync_scope`.
81        let future =
82            OUTGOING_MAILBOX.sync_scope(outgoing_mailbox.clone(), move || self.service.call(req));
83        // We can then make the task local value available to the asynchronous execution context
84        // using `scope` without losing the messages that might have been recorded by the middleware
85        // chain.
86        let storage_backend = self.storage_backend.clone();
87        Box::pin(OUTGOING_MAILBOX.scope(outgoing_mailbox, async move {
88            let response: Result<Self::Response, Self::Error> = future.await;
89            response.map(|mut response| {
90                OUTGOING_MAILBOX
91                    .with(|m| {
92                        storage_backend.store(
93                            &m.messages.borrow(),
94                            // This `.clone()` is cheap because `HttpRequest` is just an `Rc` pointer
95                            // around the actual request data.
96                            response.request().clone(),
97                            response.response_mut().head_mut(),
98                        )
99                    })
100                    .unwrap();
101                response
102            })
103        }))
104    }
105}