warp_sessions/
session.rs

1use crate::cookie::CookieOptions;
2use crate::error::SessionError;
3use async_session::{Session, SessionStore};
4use async_trait::async_trait;
5use std::{ops::Deref, sync::Arc};
6use warp::{Rejection, Reply};
7
8#[derive(Debug, Clone)]
9pub struct ArcSessionStore<T: SessionStore>(pub Arc<T>);
10
11#[async_trait]
12impl<T> SessionStore for ArcSessionStore<T>
13where
14    T: SessionStore,
15{
16    async fn load_session(&self, cookie_value: String) -> async_session::Result<Option<Session>> {
17        self.0.deref().load_session(cookie_value).await
18    }
19    async fn store_session(&self, session: Session) -> async_session::Result<Option<String>> {
20        self.0.deref().store_session(session).await
21    }
22    async fn destroy_session(&self, session: Session) -> async_session::Result {
23        self.0.deref().destroy_session(session).await
24    }
25    async fn clear_store(&self) -> async_session::Result {
26        self.0.deref().clear_store().await
27    }
28}
29
30/// SessionWithStore binds a session object with its backing store and some cookie options.
31/// This is passed around by routes wanting to do things with a session.
32#[derive(Clone)]
33pub struct SessionWithStore<S: SessionStore> {
34    pub session: Session,
35    pub session_store: S,
36    pub cookie_options: CookieOptions,
37}
38
39/// WithSession is a warp::Reply that attaches a session ID in the form of
40/// a Set-Cookie header to an existing warp::Reply
41pub struct WithSession<T: Reply> {
42    reply: T,
43    cookie_options: CookieOptions,
44}
45
46impl<T> WithSession<T>
47where
48    T: Reply,
49{
50    /// This function binds a session to a warp::Reply. It takes the given
51    /// session and binds it to the given warp::Reply by attaching a Set-Cookie
52    /// header to it. This cookie contains the session ID. If the session was
53    /// destroyed, it handles destroying the session in the store and removing
54    /// the cookie.
55    pub async fn new<S: SessionStore>(
56        reply: T,
57        session_with_store: SessionWithStore<S>,
58    ) -> Result<WithSession<T>, Rejection> {
59        let mut cookie_options = session_with_store.cookie_options;
60
61        if session_with_store.session.is_destroyed() {
62            cookie_options.cookie_value = Some("".to_string());
63            cookie_options.max_age = Some(0);
64
65            session_with_store
66                .session_store
67                .destroy_session(session_with_store.session)
68                .await
69                .map_err(|source| SessionError::DestroyError { source })?;
70        } else {
71            if session_with_store.session.data_changed() {
72                match session_with_store
73                    .session_store
74                    .store_session(session_with_store.session)
75                    .await
76                    .map_err(|source| SessionError::StoreError { source })?
77                {
78                    Some(sid) => cookie_options.cookie_value = Some(sid),
79                    None => (),
80                }
81            }
82        }
83
84        Ok(WithSession {
85            reply,
86            cookie_options,
87        })
88    }
89}
90
91impl<T> Reply for WithSession<T>
92where
93    T: Reply,
94{
95    fn into_response(self) -> warp::reply::Response {
96        let mut res = self.reply.into_response();
97        if let Some(_) = self.cookie_options.cookie_value {
98            res.headers_mut().append(
99                "Set-Cookie",
100                http::header::HeaderValue::from_str(&self.cookie_options.to_string()).unwrap(),
101            );
102        }
103
104        res
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::{SessionWithStore, WithSession};
111    use crate::cookie::CookieOptions;
112    use async_session::{MemoryStore, Session};
113
114    #[tokio::test]
115    async fn test_session_reply_with_no_data_changed() {
116        let html_reply = warp::reply::html("".to_string());
117        let session = Session::new();
118        let session_store = MemoryStore::new();
119        let cookie_options = CookieOptions::default();
120        let session_with_store = SessionWithStore {
121            session,
122            session_store,
123            cookie_options,
124        };
125
126        assert_eq!(session_with_store.session.data_changed(), false);
127        WithSession::new(html_reply, session_with_store)
128            .await
129            .unwrap();
130    }
131
132    #[tokio::test]
133    async fn test_session_reply_with_data_changed() {
134        let html_reply = warp::reply::html("".to_string());
135        let mut session = Session::new();
136        session.insert("key", "value").unwrap();
137        let session_store = MemoryStore::new();
138        let cookie_options = CookieOptions::default();
139        let session_with_store = SessionWithStore {
140            session,
141            session_store,
142            cookie_options,
143        };
144
145        assert_eq!(session_with_store.session.data_changed(), true);
146        WithSession::new(html_reply, session_with_store)
147            .await
148            .unwrap();
149    }
150
151    #[tokio::test]
152    async fn test_session_reply_with_session_destroyed() {
153        let html_reply = warp::reply::html("".to_string());
154        let mut session = Session::new();
155        session.destroy();
156        let session_store = MemoryStore::new();
157        let cookie_options = CookieOptions::default();
158        let session_with_store = SessionWithStore {
159            session,
160            session_store,
161            cookie_options,
162        };
163
164        assert_eq!(session_with_store.session.is_destroyed(), true);
165        WithSession::new(html_reply, session_with_store)
166            .await
167            .unwrap();
168    }
169}