rsession 0.2.1

A simple session manager for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240

//! Axum session management middleware
//!
//! This module provides Axum integration for session management using Tower middleware,
//! handling session creation, storage operations, and cookie management.
use crate::{Session, SessionBuilder, SessionInner, SessionStatus, SessionStore};
use axum::body::Body;
use axum::http::HeaderMap;
use axum::http::header::COOKIE;
use axum::{extract::Request, response::Response};
use cookie::{Cookie, CookieJar};
use futures::future::BoxFuture;
use http::header::SET_COOKIE;
use std::cell::RefCell;
use std::convert::Infallible;
use std::rc::Rc;
use std::sync::Arc;
use std::task::{Context, Poll};
use tower::{Layer, Service};

/// Axum session management middleware
///
/// This Tower Service implementation integrates session management into Axum applications,
/// handling session loading from storage, request processing, and response finalization.
///
/// # Type Parameters
/// * `S` - The inner service type
/// * `Storage` - The session storage backend implementing SessionStore
#[derive(Clone)]
pub struct AxumSessionMiddleware<S, Storage>
where
    S: Service<Request<Body>, Response = Response, Error = Infallible> + Send + 'static,
    S::Future: Send + 'static,
    Storage: SessionStore + 'static + Send + Sync + Clone,
{
    inner: S,
    builder: Arc<SessionBuilder>,
    store: Arc<Storage>,
}

/// Tower Service implementation for Axum session middleware
///
/// Processes incoming requests to load or create sessions, injects the session into
/// request extensions, and handles responses to persist session changes and update cookies.
impl<S, Storage> Service<Request<Body>> for AxumSessionMiddleware<S, Storage>
where
    Storage: SessionStore + 'static + Send + Sync + Clone,
    S: Service<Request<Body>, Response = Response, Error = Infallible> + Clone + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = Response;
    type Error = Infallible;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut req: Request<Body>) -> Self::Future {
        let store = self.store.clone();
        let not_ready_inner = self.inner.clone();
        let mut ready_inner = std::mem::replace(&mut self.inner, not_ready_inner);
        let builder = self.builder.clone();
        Box::pin(async move {
            let cookies = get_cookies(req.headers());
            let session_key = cookies.get(&builder.key);
            let session_inner = if let Some(session_key) = session_key {
                let session_key = session_key.value().to_string();
                if let Ok(inner) = store.get(&session_key).await {
                    inner
                } else {
                    SessionInner::new(builder.rand_key.generate())
                }
            } else {
                SessionInner::new(builder.rand_key.generate())
            };
            let session = Session::new(Rc::new(RefCell::new(session_inner)));
            req.extensions_mut().insert(session.clone());
            let future = ready_inner.call(req);
            let res = future.await;
            match res {
                Ok(mut res) => {
                    let inner = session.inner();
                    let cookie = builder.build(inner.id.clone());
                    if let Ok(cookie) = cookie.to_string().parse() {
                        res.headers_mut().insert(SET_COOKIE, cookie);
                    }
                    match inner.status {
                        SessionStatus::UnChange => {
                            if builder.auto_expire {
                                store
                                    .expire(&inner.id.to_string(), builder.expire_time)
                                    .await
                                    .ok();
                            }
                            store.set(&inner.id.to_string(), inner.clone()).await.ok();
                        }
                        SessionStatus::Change => {
                            store.remove(&inner.id.to_string()).await.ok();
                            store.set(&inner.id.to_string(), inner.clone()).await.ok();
                            store
                                .expire(&inner.id.to_string(), builder.expire_time)
                                .await
                                .ok();
                        }
                        SessionStatus::Clear => {
                            store.remove(&inner.id.to_string()).await.ok();
                        }
                        SessionStatus::Destroy => {
                            store.remove(&inner.id.to_string()).await.ok();
                        }
                        SessionStatus::Expire => {
                            store
                                .expire(&inner.id.to_string(), builder.expire_time)
                                .await
                                .ok();
                        }
                    }
                    Ok(res)
                }
                Err(err) => Err(err),
            }
        })
    }
}

/// Tower Layer for Axum session middleware
///
/// Provides a configuration layer for adding session management to Axum applications
/// through the Tower middleware system.
///
/// # Type Parameters
/// * `Storage` - The session storage backend implementing SessionStore
#[derive(Clone)]
pub struct AxumSessionMiddlewareLayer<Storage>
where
    Storage: SessionStore + 'static,
{
    builder: Arc<SessionBuilder>,
    store: Arc<Storage>,
}

impl<Storage> AxumSessionMiddlewareLayer<Storage>
where
    Storage: SessionStore + 'static,
{
    /// Creates a new AxumSessionMiddlewareLayer
    ///
    /// # Arguments
    /// * `builder` - Session configuration builder with cookie/session settings
    /// * `store` - Session storage backend implementation
    pub fn new(builder: SessionBuilder, store: Storage) -> Self {
        Self {
            builder: Arc::new(builder),
            store: Arc::new(store),
        }
    }
}

/// Tower Layer implementation for Axum session middleware
///
/// Applies the session middleware to an Axum service, enabling session management
/// for all routes wrapped by this layer.
impl<S, Storage> Layer<S> for AxumSessionMiddlewareLayer<Storage>
where
    Storage: SessionStore + 'static,
    S: Service<Request, Response = Response, Error = Infallible>
        + Send
        + 'static
        + std::marker::Sync,
    S::Future: Send + 'static,
{
    type Service = AxumSessionMiddleware<S, Storage>;

    fn layer(&self, inner: S) -> Self::Service {
        AxumSessionMiddleware {
            inner,
            builder: self.builder.clone(),
            store: self.store.clone(),
        }
    }
}

/// Parses cookies from request headers into a CookieJar
///
/// # Arguments
/// * `headers` - Request headers containing cookie information
///
/// # Returns
/// A CookieJar populated with cookies from the request
pub(crate) fn get_cookies(headers: &HeaderMap) -> CookieJar {
    let mut jar = CookieJar::new();
    let cookie_iter = headers
        .get_all(COOKIE)
        .into_iter()
        .filter_map(|value| value.to_str().ok())
        .flat_map(|value| value.split(';'))
        .filter_map(|cookie| Cookie::parse(cookie.to_owned()).ok());
    for cookie in cookie_iter {
        jar.add_original(cookie);
    }
    jar
}

impl<S> axum::extract::FromRequest<S> for Session {
    type Rejection = (axum::http::status::StatusCode, &'static str);

    fn from_request(
        req: Request,
        _: &S,
    ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
        async move {
            let inner = req.extensions().get::<Session>();
            if let Some(inner) = inner {
                return Ok(inner.clone());
            } else {
                Err((
                    axum::http::status::StatusCode::INTERNAL_SERVER_ERROR,
                    "session not found",
                ))
            }
        }
    }
}

// impl<S> axum::extract::FromRequestParts<S> for Session {
//     type Rejection = (axum::http::status::StatusCode, &'static str);
//     fn from_request_parts(parts: &mut Parts, _: &S) -> impl Future<Output=Result<Self, Self::Rejection>> + Send {
//         async move {
//             let inner = parts.extensions.get::<SessionInner>();
//             if let Some(inner) = inner {
//                 Ok(Session::new(Rc::new(RefCell::new(inner.clone()))))
//             } else {
//                 Err((
//                     axum::http::status::StatusCode::INTERNAL_SERVER_ERROR,
//                     "session not found",
//                 ))
//             }
//         }
//     }
// }