axum_security/cookie/store/
memory.rs

1use std::{collections::HashMap, convert::Infallible, sync::Arc};
2
3use tokio::sync::RwLock;
4
5use crate::cookie::{CookieSession, CookieStore, SessionId};
6
7pub struct MemStore<S> {
8    inner: Arc<RwLock<HashMap<SessionId, CookieSession<S>>>>,
9}
10
11impl<S> Default for MemStore<S> {
12    fn default() -> Self {
13        Self::new()
14    }
15}
16
17impl<S> MemStore<S> {
18    pub fn new() -> Self {
19        Self {
20            inner: RwLock::new(HashMap::new()).into(),
21        }
22    }
23}
24
25impl<S> Clone for MemStore<S> {
26    fn clone(&self) -> Self {
27        MemStore {
28            inner: self.inner.clone(),
29        }
30    }
31}
32
33impl<S: Send + Sync + Clone + 'static> CookieStore for MemStore<S> {
34    type State = S;
35    type Error = Infallible;
36
37    async fn store_session(&self, session: CookieSession<Self::State>) -> Result<(), Self::Error> {
38        let mut lock = self.inner.write().await;
39        let id = session.session_id.clone();
40        lock.insert(id, session);
41        Ok(())
42    }
43
44    async fn remove_session(
45        &self,
46        id: &SessionId,
47    ) -> Result<Option<CookieSession<Self::State>>, Self::Error> {
48        let mut lock = self.inner.write().await;
49        Ok(lock.remove(id))
50    }
51
52    async fn load_session(
53        &self,
54        id: &SessionId,
55    ) -> Result<Option<CookieSession<Self::State>>, Self::Error> {
56        let lock = self.inner.read().await;
57        Ok(lock.get(id).cloned())
58    }
59
60    async fn remove_before(&self, deadline: u64) -> Result<(), Self::Error> {
61        let mut lock = self.inner.write().await;
62        lock.retain(|_, v| v.created_at > deadline);
63        Ok(())
64    }
65}
66
67#[cfg(test)]
68mod mem_store {
69    use crate::cookie::{CookieSession, CookieStore, MemStore, SessionId};
70
71    #[tokio::test]
72    async fn create() {
73        let store = MemStore::<i32>::new();
74
75        let session_id = SessionId::new();
76        let created_at = 100;
77        let value = 1;
78
79        let session = CookieSession::new(session_id.clone(), created_at, value);
80
81        store.store_session(session).await.unwrap();
82
83        let session = store
84            .load_session(&session_id)
85            .await
86            .unwrap()
87            .expect("session to be created");
88
89        assert!(session.created_at == created_at);
90        assert!(session.state == value);
91
92        let session = store
93            .remove_session(&session_id)
94            .await
95            .unwrap()
96            .expect("session should exists");
97
98        assert!(session.created_at == created_at);
99        assert!(session.state == value);
100
101        let session = store.load_session(&session_id).await.unwrap();
102        assert!(session.is_none());
103    }
104
105    #[tokio::test]
106    async fn remove() {
107        let store = MemStore::<i32>::new();
108
109        let session_id = SessionId::new();
110        let created_at = 100;
111        let value = 1;
112
113        let session = CookieSession::new(session_id.clone(), created_at, value);
114
115        store.store_session(session).await.unwrap();
116
117        store.remove_before(101).await;
118
119        let session = store.load_session(&session_id).await.unwrap();
120        assert!(session.is_none());
121    }
122}
axum_security/cookie/store/memory.rs

axum_security/cookie/store/
memory.rs
