1use crate::key::AppKey;
18use crate::session::{Session, is_valid_id};
19use crate::store::{DEFAULT_LIFETIME, SessionStore, SharedStore};
20use crate::{base64, constant_time_eq};
21use hmac::{Hmac, Mac};
22use rustlavel_core::{Config, Json, Result};
23use rustlavel_http::handler::BoxFuture;
24use rustlavel_http::middleware::{Middleware, Next};
25use rustlavel_http::{Cookie, Method, Request, Response, SameSite};
26use sha2::Sha256;
27use std::collections::BTreeMap;
28use std::sync::{Arc, Mutex, MutexGuard};
29use std::time::Duration;
30
31type HmacSha256 = Hmac<Sha256>;
32
33pub const DEFAULT_COOKIE: &str = "rustlavel_session";
35
36#[derive(Clone)]
43pub struct SessionHandle(Arc<Mutex<Session>>);
44
45impl SessionHandle {
46 pub fn new(session: Session) -> Self {
47 SessionHandle(Arc::new(Mutex::new(session)))
48 }
49
50 pub fn lock(&self) -> MutexGuard<'_, Session> {
56 self.0.lock().expect("session lock poisoned")
57 }
58
59 pub fn with<R>(&self, action: impl FnOnce(&mut Session) -> R) -> R {
61 action(&mut self.lock())
62 }
63
64 pub fn id(&self) -> String {
65 self.lock().id().to_string()
66 }
67
68 pub fn get(&self, key: &str) -> Option<Json> {
69 self.lock().get(key).cloned()
70 }
71
72 pub fn get_string(&self, key: &str) -> Option<String> {
73 self.lock().get_string(key)
74 }
75
76 pub fn has(&self, key: &str) -> bool {
77 self.lock().has(key)
78 }
79
80 pub fn put(&self, key: impl Into<String>, value: impl Into<Json>) {
81 self.lock().put(key, value);
82 }
83
84 pub fn forget(&self, key: &str) -> Option<Json> {
85 self.lock().forget(key)
86 }
87
88 pub fn flush(&self) {
89 self.lock().flush();
90 }
91
92 pub fn all(&self) -> BTreeMap<String, Json> {
93 self.lock().all().clone()
94 }
95
96 pub fn flash(&self, key: impl Into<String>, value: impl Into<Json>) {
97 self.lock().flash(key, value);
98 }
99
100 pub fn keep(&self, key: &str) {
101 self.lock().keep(key);
102 }
103
104 pub fn reflash(&self) {
105 self.lock().reflash();
106 }
107
108 pub fn regenerate(&self) -> String {
110 self.lock().regenerate().to_string()
111 }
112
113 pub fn invalidate(&self) -> String {
115 self.lock().invalidate().to_string()
116 }
117
118 pub fn token(&self) -> String {
120 self.lock().token()
121 }
122}
123
124impl std::fmt::Debug for SessionHandle {
125 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126 f.debug_struct("SessionHandle").field("id", &self.id()).finish()
127 }
128}
129
130pub trait SessionExt {
132 fn try_session(&self) -> Option<&SessionHandle>;
134
135 fn session(&self) -> &SessionHandle;
137}
138
139impl SessionExt for Request {
140 fn try_session(&self) -> Option<&SessionHandle> {
141 self.extension::<SessionHandle>()
142 }
143
144 fn session(&self) -> &SessionHandle {
145 self.try_session().expect(
146 "no session on this request. Add the `session` middleware to the router \
147 (`router.middleware(SessionManager::from_config(&config, store)?)`) before any route \
148 that uses `req.session()`.",
149 )
150 }
151}
152
153#[derive(Clone)]
157pub struct SessionManager {
158 inner: Arc<Inner>,
159}
160
161struct Inner {
162 store: SharedStore,
163 cookie_key: [u8; 32],
164 cookie: String,
165 lifetime: Duration,
166 path: String,
167 secure: bool,
168 same_site: SameSite,
169}
170
171impl rustlavel_http::Flash for SessionHandle {
177 fn flash(&self, key: &str, value: Json) {
178 self.lock().flash(key, value);
179 }
180
181 fn take(&self, key: &str) -> Option<Json> {
182 self.lock().forget(key)
183 }
184
185 fn peek(&self, key: &str) -> Option<Json> {
186 self.lock().get(key).cloned()
187 }
188}
189
190impl SessionManager {
191 pub fn new(key: &AppKey, store: impl SessionStore) -> Self {
193 SessionManager { inner: Arc::new(SessionManager::defaults(key, store)) }
194 }
195
196 pub fn from_config(config: &Config, store: impl SessionStore) -> Result<Self> {
202 let mut inner = SessionManager::defaults(&AppKey::from_config(config)?, store);
203 inner.cookie = config.string("session.cookie", DEFAULT_COOKIE);
204 inner.lifetime = SessionManager::lifetime_from_config(config);
205 inner.path = config.string("session.path", "/");
206 inner.secure = config.bool("session.secure", config.is_production());
208
209 Ok(SessionManager { inner: Arc::new(inner) })
210 }
211
212 fn defaults(key: &AppKey, store: impl SessionStore) -> Inner {
213 let lifetime = store.lifetime();
214 Inner {
215 store: Arc::new(store),
216 cookie_key: key.derive("session-cookie"),
217 cookie: DEFAULT_COOKIE.to_string(),
218 lifetime,
219 path: "/".to_string(),
220 secure: false,
225 same_site: SameSite::Lax,
226 }
227 }
228
229 pub fn lifetime_from_config(config: &Config) -> Duration {
231 match config.int("session.lifetime", -1) {
232 minutes if minutes > 0 => Duration::from_secs(minutes as u64 * 60),
233 _ => DEFAULT_LIFETIME,
234 }
235 }
236
237 pub fn cookie_name(mut self, name: impl Into<String>) -> Self {
238 Arc::make_mut(&mut self.inner).cookie = name.into();
239 self
240 }
241
242 pub fn lifetime(mut self, lifetime: Duration) -> Self {
243 Arc::make_mut(&mut self.inner).lifetime = lifetime;
244 self
245 }
246
247 pub fn path(mut self, path: impl Into<String>) -> Self {
248 Arc::make_mut(&mut self.inner).path = path.into();
249 self
250 }
251
252 pub fn secure(mut self, secure: bool) -> Self {
254 Arc::make_mut(&mut self.inner).secure = secure;
255 self
256 }
257
258 pub fn same_site(mut self, same_site: SameSite) -> Self {
259 Arc::make_mut(&mut self.inner).same_site = same_site;
260 self
261 }
262
263 pub fn store(&self) -> &SharedStore {
264 &self.inner.store
265 }
266
267 fn sign(&self, id: &str) -> String {
274 let mut mac =
275 HmacSha256::new_from_slice(&self.inner.cookie_key).expect("HMAC accepts any key length");
276 mac.update(id.as_bytes());
277 format!("{id}.{}", base64::encode_url(&mac.finalize().into_bytes()))
278 }
279
280 fn unsign(&self, value: &str) -> Option<String> {
282 let (id, signature) = value.rsplit_once('.')?;
283 if !is_valid_id(id) {
284 return None;
285 }
286
287 let expected = self.sign(id);
288 let expected = expected.rsplit_once('.')?.1;
289 constant_time_eq(expected.as_bytes(), signature.as_bytes()).then(|| id.to_string())
290 }
291
292 fn cookie_for(&self, id: &str) -> Cookie {
293 Cookie::new(self.inner.cookie.clone(), self.sign(id))
294 .path(self.inner.path.clone())
295 .max_age(self.inner.lifetime)
296 .http_only(true)
297 .secure(self.inner.secure)
298 .same_site(self.inner.same_site)
299 }
300}
301
302impl Clone for Inner {
303 fn clone(&self) -> Self {
304 Inner {
305 store: Arc::clone(&self.store),
306 cookie_key: self.cookie_key,
307 cookie: self.cookie.clone(),
308 lifetime: self.lifetime,
309 path: self.path.clone(),
310 secure: self.secure,
311 same_site: self.same_site,
312 }
313 }
314}
315
316impl Middleware for SessionManager {
317 fn handle(&self, mut request: Request, next: Next) -> BoxFuture<Response> {
318 let manager = self.clone();
319
320 Box::pin(async move {
321 let presented = request
322 .cookie(&manager.inner.cookie)
323 .and_then(|value| manager.unsign(&value));
324
325 let (session, existed) = match &presented {
326 Some(id) => match manager.inner.store.read(id).await {
327 Ok(Some(session)) => (session, true),
328 Ok(None) => (Session::new(), false),
329 Err(error) => {
330 rustlavel_core::error!("could not read the session: {error}");
333 (Session::new(), false)
334 }
335 },
336 None => (Session::new(), false),
337 };
338
339 let loaded_id = session.id().to_string();
340 let handle = SessionHandle::new(session);
341
342 let method = request.method();
344 let wants_json = request.wants_json();
345 let target = request.target().to_string();
346
347 request.extend(handle.clone());
348 let flash: std::sync::Arc<dyn rustlavel_http::Flash> = std::sync::Arc::new(handle.clone());
352 request.extend(flash);
353
354 let response = next.run(request).await;
355
356 if matches!(method, Method::Get)
372 && !wants_json
373 && (200..300).contains(&response.status.code())
374 && handle.lock().is_dirty()
375 {
376 handle.put(rustlavel_http::flash::PREVIOUS_URL_KEY, target.clone());
377 }
378
379 let (session, dirty) = {
382 let mut session = handle.lock();
383 session.age_flash();
384 (session.clone(), session.is_dirty())
385 };
386
387 if !dirty && !existed {
390 return response;
391 }
392
393 if existed && session.id() != loaded_id
396 && let Err(error) = manager.inner.store.destroy(&loaded_id).await
397 {
398 rustlavel_core::error!("could not destroy the previous session: {error}");
399 }
400
401 if let Err(error) = manager.inner.store.write(&session).await {
402 rustlavel_core::error!("could not write the session: {error}");
403 return response;
404 }
405
406 response.with_cookie(manager.cookie_for(session.id()))
407 })
408 }
409}
410
411impl std::fmt::Debug for SessionManager {
412 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
413 f.debug_struct("SessionManager")
414 .field("cookie", &self.inner.cookie)
415 .field("lifetime", &self.inner.lifetime)
416 .field("secure", &self.inner.secure)
417 .finish()
418 }
419}
420
421#[cfg(test)]
422mod tests {
423 use super::*;
424 use crate::store::MemoryStore;
425 use rustlavel_http::{Router, TestClient};
426
427 fn key() -> AppKey {
428 AppKey::from_bytes([3u8; 32])
429 }
430
431 fn manager() -> SessionManager {
432 SessionManager::new(&key(), MemoryStore::new())
433 }
434
435 fn cookie_of(response: &rustlavel_http::TestResponse) -> String {
437 let header = response
438 .header("set-cookie")
439 .unwrap_or_else(|| panic!("no set-cookie header on the response"));
440 header.split(';').next().unwrap().to_string()
441 }
442
443 fn counting_router(manager: SessionManager) -> Router {
444 let mut router = Router::new();
445 router.middleware(manager);
446 router.get("/count", |request: Request| async move {
447 let visits = request.session().get("visits").and_then(|v| v.as_i64()).unwrap_or(0) + 1;
448 request.session().put("visits", visits);
449 format!("{visits}")
450 });
451 router.get("/quiet", |_request: Request| async { "nothing to see" });
452 router.get("/flash", |request: Request| async move {
453 request.session().flash("status", "Saved");
454 "flashed"
455 });
456 router.get("/status", |request: Request| async move {
457 request.session().get_string("status").unwrap_or_else(|| "none".into())
458 });
459 router.get("/rotate", |request: Request| async move {
460 request.session().put("kept", "yes");
461 request.session().regenerate()
462 });
463 router
464 }
465
466 #[tokio::test]
467 async fn a_session_survives_between_requests_through_the_cookie() {
468 let client = TestClient::new(counting_router(manager()));
469
470 let first = client.get("/count").await.assert_ok();
471 assert_eq!(first.body(), "1");
472 let cookie = cookie_of(&first);
473 assert!(cookie.starts_with(DEFAULT_COOKIE));
474
475 let second = client
476 .send(Request::new(Method::Get, "/count").with_header("cookie", cookie.clone()))
477 .await;
478 assert_eq!(second.body(), "2");
479
480 let third = client
481 .send(Request::new(Method::Get, "/count").with_header("cookie", cookie))
482 .await;
483 assert_eq!(third.body(), "3");
484 }
485
486 #[tokio::test]
487 async fn the_session_cookie_is_http_only_and_scoped() {
488 let response = TestClient::new(counting_router(manager())).get("/count").await;
489 let header = response.header("set-cookie").unwrap();
490
491 assert!(header.contains("HttpOnly"), "header was {header}");
492 assert!(header.contains("SameSite=Lax"));
493 assert!(header.contains("Path=/"));
494 assert!(header.contains("Max-Age=7200"));
495 assert!(!header.contains("Secure"), "a plain-http development cookie is not Secure");
496 }
497
498 #[tokio::test]
499 async fn a_request_that_never_touches_the_session_gets_no_cookie() {
500 let response = TestClient::new(counting_router(manager())).get("/quiet").await.assert_ok();
501
502 assert_eq!(
503 response.header("set-cookie"),
504 None,
505 "nothing was stored, so there is nothing to hand the client"
506 );
507 }
508
509 #[tokio::test]
510 async fn a_forged_or_unsigned_cookie_starts_a_fresh_session() {
511 let client = TestClient::new(counting_router(manager()));
512
513 let real = client.get("/count").await;
514 let cookie = cookie_of(&real);
515
516 let (name_and_id, _) = cookie.rsplit_once('.').unwrap();
518 let forged = format!("{name_and_id}.{}", base64::encode_url(&[0u8; 32]));
519 let response = client
520 .send(Request::new(Method::Get, "/count").with_header("cookie", forged))
521 .await;
522 assert_eq!(response.body(), "1", "a forged cookie must not open the real session");
523
524 let response = client
526 .send(Request::new(Method::Get, "/count").with_header(
527 "cookie",
528 format!("{DEFAULT_COOKIE}={}", Session::new_id()),
529 ))
530 .await;
531 assert_eq!(response.body(), "1");
532 }
533
534 #[test]
535 fn a_cookie_signed_with_another_key_is_refused() {
536 let mine = SessionManager::new(&key(), MemoryStore::new());
537 let theirs = SessionManager::new(&AppKey::from_bytes([4u8; 32]), MemoryStore::new());
538 let id = Session::new_id();
539
540 assert_eq!(mine.unsign(&mine.sign(&id)).as_deref(), Some(id.as_str()));
541 assert!(mine.unsign(&theirs.sign(&id)).is_none());
542 assert!(mine.unsign("no-dot-here").is_none());
544 assert!(mine.unsign(&format!("not-an-id.{}", base64::encode_url(&[0u8; 32]))).is_none());
545 }
546
547 #[tokio::test]
548 async fn flashed_data_is_readable_on_the_next_request_and_gone_after_it() {
549 let client = TestClient::new(counting_router(manager()));
550
551 let flashed = client.get("/flash").await;
552 let cookie = cookie_of(&flashed);
553
554 let next = client
555 .send(Request::new(Method::Get, "/status").with_header("cookie", cookie.clone()))
556 .await;
557 assert_eq!(next.body(), "Saved");
558
559 let after = client
560 .send(Request::new(Method::Get, "/status").with_header("cookie", cookie))
561 .await;
562 assert_eq!(after.body(), "none");
563 }
564
565 #[tokio::test]
566 async fn regenerating_issues_a_new_cookie_and_drops_the_old_record() {
567 let manager = manager();
568 let store = Arc::clone(manager.store());
569 let client = TestClient::new(counting_router(manager));
570
571 let first = client.get("/count").await;
572 let old_cookie = cookie_of(&first);
573 let old_id = old_cookie.split_once('=').unwrap().1.rsplit_once('.').unwrap().0.to_string();
574 assert!(store.read(&old_id).await.unwrap().is_some());
575
576 let rotated = client
577 .send(Request::new(Method::Get, "/rotate").with_header("cookie", old_cookie.clone()))
578 .await;
579 let new_id = rotated.body();
580 let new_cookie = cookie_of(&rotated);
581
582 assert_ne!(new_id, old_id);
583 assert!(store.read(&old_id).await.unwrap().is_none(), "the old id must stop working");
584
585 let carried = store.read(&new_id).await.unwrap().expect("the new id holds the session");
586 assert_eq!(carried.get_string("kept").as_deref(), Some("yes"));
587 assert_ne!(new_cookie, old_cookie);
588 }
589
590 #[tokio::test]
591 async fn an_expired_stored_session_starts_over() {
592 let manager = SessionManager::new(&key(), MemoryStore::with_lifetime(Duration::ZERO));
593 let client = TestClient::new(counting_router(manager));
594
595 let first = client.get("/count").await;
596 let cookie = cookie_of(&first);
597
598 let second = client
599 .send(Request::new(Method::Get, "/count").with_header("cookie", cookie))
600 .await;
601 assert_eq!(second.body(), "1", "the stored session had already expired");
602 }
603
604 #[test]
605 fn configuration_drives_the_cookie_name_lifetime_and_secure_flag() {
606 let config = Config::new();
607 config.set("app.key", AppKey::generate());
608 config.set("session.cookie", "my_app_session");
609 config.set("session.lifetime", 30);
610 config.set("app.env", "production");
611
612 let manager = SessionManager::from_config(&config, MemoryStore::new()).unwrap();
613
614 assert_eq!(manager.inner.cookie, "my_app_session");
615 assert_eq!(manager.inner.lifetime, Duration::from_secs(30 * 60));
616 assert!(manager.inner.secure, "production cookies default to Secure");
617 assert_eq!(SessionManager::lifetime_from_config(&Config::new()), DEFAULT_LIFETIME);
618 }
619
620 #[test]
621 fn builders_override_configuration() {
622 let manager = manager()
623 .cookie_name("custom")
624 .lifetime(Duration::from_secs(60))
625 .path("/app")
626 .secure(true)
627 .same_site(SameSite::Strict);
628
629 assert_eq!(manager.inner.cookie, "custom");
630 assert_eq!(manager.inner.path, "/app");
631 assert!(manager.inner.secure);
632
633 let header = manager.cookie_for(&Session::new_id()).to_header();
634 assert!(header.contains("SameSite=Strict"));
635 assert!(header.contains("Secure"));
636 assert!(header.starts_with("custom="));
637 }
638
639 #[tokio::test]
640 async fn a_handler_without_the_middleware_says_what_is_missing() {
641 let mut router = Router::new();
642 router.get("/", |request: Request| async move {
643 assert!(request.try_session().is_none());
644 request.session().id()
645 });
646
647 let response = TestClient::new(router).get("/").await;
649 assert_eq!(response.status(), 500);
650 }
651}