|
|
|
|
@ -1,4 +1,4 @@ |
|
|
|
|
use parking_lot::RwLock; |
|
|
|
|
use parking_lot::{Mutex, RwLock, RwLockUpgradableReadGuard}; |
|
|
|
|
use rand::Rng; |
|
|
|
|
|
|
|
|
|
use rocket::{ |
|
|
|
|
@ -8,68 +8,110 @@ use rocket::{ |
|
|
|
|
Outcome, Request, Response, Rocket, State, |
|
|
|
|
}; |
|
|
|
|
|
|
|
|
|
use std::borrow::Cow; |
|
|
|
|
use std::collections::HashMap; |
|
|
|
|
use std::fmt::{self, Display, Formatter}; |
|
|
|
|
use std::marker::PhantomData; |
|
|
|
|
use std::ops::Add; |
|
|
|
|
use std::time::{Duration, Instant}; |
|
|
|
|
|
|
|
|
|
const SESSION_COOKIE: &str = "SESSID"; |
|
|
|
|
const SESSION_ID_LEN: usize = 16; |
|
|
|
|
|
|
|
|
|
/// Session, as stored in the sessions store
|
|
|
|
|
/// Session store (shared state)
|
|
|
|
|
#[derive(Debug)] |
|
|
|
|
struct SessionInstance<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
pub struct SessionStore<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
/// Data object
|
|
|
|
|
data: D, |
|
|
|
|
/// Expiry
|
|
|
|
|
expires: Instant, |
|
|
|
|
/// The internally mutable map of sessions
|
|
|
|
|
inner: RwLock<StoreInner<D>>, |
|
|
|
|
// Session config
|
|
|
|
|
config: SessionConfig, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Session store (shared state)
|
|
|
|
|
#[derive(Default, Debug)] |
|
|
|
|
pub struct SessionStore<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
/// The internaly mutable map of sessions
|
|
|
|
|
inner: RwLock<HashMap<String, SessionInstance<D>>>, |
|
|
|
|
/// Session config object
|
|
|
|
|
#[derive(Debug, Clone)] |
|
|
|
|
struct SessionConfig { |
|
|
|
|
/// Sessions lifespan
|
|
|
|
|
lifespan: Duration, |
|
|
|
|
/// Session cookie name
|
|
|
|
|
cookie_name: Cow<'static, str>, |
|
|
|
|
/// Session cookie path
|
|
|
|
|
cookie_path: Cow<'static, str>, |
|
|
|
|
/// Session ID character length
|
|
|
|
|
cookie_len: usize, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl<D> SessionStore<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
impl Default for SessionConfig { |
|
|
|
|
fn default() -> Self { |
|
|
|
|
Self { |
|
|
|
|
lifespan: Duration::from_secs(3600), |
|
|
|
|
cookie_name: "rocket_session".into(), |
|
|
|
|
cookie_path: "/".into(), |
|
|
|
|
cookie_len: 16, |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Mutable object stored inside SessionStore behind a RwLock
|
|
|
|
|
#[derive(Debug)] |
|
|
|
|
struct StoreInner<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
/// Remove all expired sessions
|
|
|
|
|
pub fn remove_expired(&self) { |
|
|
|
|
let now = Instant::now(); |
|
|
|
|
self.inner.write().retain(|_k, v| v.expires > now); |
|
|
|
|
sessions: HashMap<String, Mutex<SessionInstance<D>>>, |
|
|
|
|
last_expiry_sweep: Instant, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl<D> Default for StoreInner<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
fn default() -> Self { |
|
|
|
|
Self { |
|
|
|
|
sessions: Default::default(), |
|
|
|
|
// the first expiry sweep is scheduled one lifetime from start-up
|
|
|
|
|
last_expiry_sweep: Instant::now(), |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Session, as stored in the sessions store
|
|
|
|
|
#[derive(Debug)] |
|
|
|
|
struct SessionInstance<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
/// Data object
|
|
|
|
|
data: D, |
|
|
|
|
/// Expiry
|
|
|
|
|
expires: Instant, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Session ID newtype for rocket's "local_cache"
|
|
|
|
|
#[derive(PartialEq, Hash, Clone, Debug)] |
|
|
|
|
#[derive(Clone, Debug)] |
|
|
|
|
struct SessionID(String); |
|
|
|
|
|
|
|
|
|
impl SessionID { |
|
|
|
|
fn as_str(&self) -> &str { |
|
|
|
|
self.0.as_str() |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
fn to_string(&self) -> String { |
|
|
|
|
self.0.clone() |
|
|
|
|
impl Display for SessionID { |
|
|
|
|
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { |
|
|
|
|
f.write_str(&self.0) |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Session instance
|
|
|
|
|
///
|
|
|
|
|
/// To access the active session, simply add it as an argument to a route function.
|
|
|
|
|
///
|
|
|
|
|
/// Sessions are started, restored, or expired in the `FromRequest::from_request()` method
|
|
|
|
|
/// when a `Session` is prepared for one of the route functions.
|
|
|
|
|
#[derive(Debug)] |
|
|
|
|
pub struct Session<'a, D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
/// The shared state reference
|
|
|
|
|
store: State<'a, SessionStore<D>>, |
|
|
|
|
@ -78,51 +120,82 @@ pub struct Session<'a, D> |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl<'a, 'r, D> FromRequest<'a, 'r> for Session<'a, D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
type Error = (); |
|
|
|
|
|
|
|
|
|
fn from_request(request: &'a Request<'r>) -> Outcome<Self, (Status, Self::Error), ()> { |
|
|
|
|
let store : State<SessionStore<D>> = request.guard().unwrap(); |
|
|
|
|
let store: State<SessionStore<D>> = request.guard().unwrap(); |
|
|
|
|
Outcome::Success(Session { |
|
|
|
|
id: request.local_cache(|| { |
|
|
|
|
let store_ug = store.inner.upgradable_read(); |
|
|
|
|
|
|
|
|
|
// Resolve session ID
|
|
|
|
|
let id = if let Some(cookie) = request.cookies().get(SESSION_COOKIE) { |
|
|
|
|
SessionID(cookie.value().to_string()) |
|
|
|
|
let id = if let Some(cookie) = request.cookies().get(&store.config.cookie_name) { |
|
|
|
|
Some(SessionID(cookie.value().to_string())) |
|
|
|
|
} else { |
|
|
|
|
SessionID( |
|
|
|
|
rand::thread_rng() |
|
|
|
|
.sample_iter(&rand::distributions::Alphanumeric) |
|
|
|
|
.take(SESSION_ID_LEN) |
|
|
|
|
.collect(), |
|
|
|
|
) |
|
|
|
|
None |
|
|
|
|
}; |
|
|
|
|
|
|
|
|
|
let new_expiration = Instant::now().add(store.lifespan); |
|
|
|
|
let mut wg = store.inner.write(); |
|
|
|
|
match wg.get_mut(id.as_str()) { |
|
|
|
|
Some(ses) => { |
|
|
|
|
// Check expiration
|
|
|
|
|
if ses.expires <= Instant::now() { |
|
|
|
|
ses.data = D::default(); |
|
|
|
|
} |
|
|
|
|
// Update expiry timestamp
|
|
|
|
|
ses.expires = new_expiration; |
|
|
|
|
}, |
|
|
|
|
None => { |
|
|
|
|
// New session
|
|
|
|
|
wg.insert( |
|
|
|
|
id.to_string(), |
|
|
|
|
SessionInstance { |
|
|
|
|
data: D::default(), |
|
|
|
|
expires: new_expiration, |
|
|
|
|
} |
|
|
|
|
); |
|
|
|
|
let expires = Instant::now().add(store.config.lifespan); |
|
|
|
|
|
|
|
|
|
if let Some(m) = id |
|
|
|
|
.as_ref() |
|
|
|
|
.and_then(|token| store_ug.sessions.get(token.as_str())) |
|
|
|
|
{ |
|
|
|
|
// --- ID obtained from a cookie && session found in the store ---
|
|
|
|
|
|
|
|
|
|
let mut inner = m.lock(); |
|
|
|
|
if inner.expires <= Instant::now() { |
|
|
|
|
// Session expired, reuse the ID but drop data.
|
|
|
|
|
inner.data = D::default(); |
|
|
|
|
} |
|
|
|
|
}; |
|
|
|
|
|
|
|
|
|
id |
|
|
|
|
// Session is extended by making a request with valid ID
|
|
|
|
|
inner.expires = expires; |
|
|
|
|
|
|
|
|
|
id.unwrap() |
|
|
|
|
} else { |
|
|
|
|
// --- ID missing or session not found ---
|
|
|
|
|
|
|
|
|
|
// Get exclusive write access to the map
|
|
|
|
|
let mut store_wg = RwLockUpgradableReadGuard::upgrade(store_ug); |
|
|
|
|
|
|
|
|
|
// This branch runs less often, and we already have write access,
|
|
|
|
|
// let's check if any sessions expired. We don't want to hog memory
|
|
|
|
|
// forever by abandoned sessions (e.g. when a client lost their cookie)
|
|
|
|
|
|
|
|
|
|
// Throttle by lifespan - e.g. sweep every hour
|
|
|
|
|
if store_wg.last_expiry_sweep.elapsed() > store.config.lifespan { |
|
|
|
|
let now = Instant::now(); |
|
|
|
|
store_wg.sessions.retain(|_k, v| v.lock().expires > now); |
|
|
|
|
|
|
|
|
|
store_wg.last_expiry_sweep = now; |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// Find a new unique ID - we are still safely inside the write guard
|
|
|
|
|
let new_id = SessionID(loop { |
|
|
|
|
let token: String = rand::thread_rng() |
|
|
|
|
.sample_iter(&rand::distributions::Alphanumeric) |
|
|
|
|
.take(store.config.cookie_len) |
|
|
|
|
.collect(); |
|
|
|
|
|
|
|
|
|
if !store_wg.sessions.contains_key(&token) { |
|
|
|
|
break token; |
|
|
|
|
} |
|
|
|
|
}); |
|
|
|
|
|
|
|
|
|
store_wg.sessions.insert( |
|
|
|
|
new_id.to_string(), |
|
|
|
|
Mutex::new(SessionInstance { |
|
|
|
|
data: Default::default(), |
|
|
|
|
expires, |
|
|
|
|
}), |
|
|
|
|
); |
|
|
|
|
|
|
|
|
|
new_id |
|
|
|
|
} |
|
|
|
|
}), |
|
|
|
|
store, |
|
|
|
|
}) |
|
|
|
|
@ -130,54 +203,100 @@ impl<'a, 'r, D> FromRequest<'a, 'r> for Session<'a, D> |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl<'a, D> Session<'a, D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
/// Get the fairing object
|
|
|
|
|
pub fn fairing(lifespan: Duration) -> impl Fairing { |
|
|
|
|
SessionFairing::<D> { |
|
|
|
|
lifespan, |
|
|
|
|
_phantom: PhantomData, |
|
|
|
|
} |
|
|
|
|
/// Create the session fairing.
|
|
|
|
|
///
|
|
|
|
|
/// You can configure the session store by calling chained methods on the returned value
|
|
|
|
|
/// before passing it to `rocket.attach()`
|
|
|
|
|
pub fn fairing() -> SessionFairing<D> { |
|
|
|
|
SessionFairing::<D>::new() |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Access the session store
|
|
|
|
|
pub fn get_store(&self) -> &SessionStore<D> { |
|
|
|
|
&self.store |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Set the session object to its default state
|
|
|
|
|
pub fn reset(&self) { |
|
|
|
|
self.tap_mut(|m| { |
|
|
|
|
/// Clear session data (replace the value with default)
|
|
|
|
|
pub fn clear(&self) { |
|
|
|
|
self.tap(|m| { |
|
|
|
|
*m = D::default(); |
|
|
|
|
}) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pub fn tap<T>(&self, func: impl FnOnce(&D) -> T) -> T { |
|
|
|
|
let rg = self.store.inner.read(); |
|
|
|
|
let instance = rg.get(self.id.as_str()).unwrap(); |
|
|
|
|
func(&instance.data) |
|
|
|
|
} |
|
|
|
|
/// Access the session's data using a closure.
|
|
|
|
|
///
|
|
|
|
|
/// The closure is called with the data value as a mutable argument,
|
|
|
|
|
/// and can return any value to be is passed up to the caller.
|
|
|
|
|
pub fn tap<T>(&self, func: impl FnOnce(&mut D) -> T) -> T { |
|
|
|
|
// Use a read guard, so other already active sessions are not blocked
|
|
|
|
|
// from accessing the store. New incoming clients may be blocked until
|
|
|
|
|
// the tap() call finishes
|
|
|
|
|
let store_rg = self.store.inner.read(); |
|
|
|
|
|
|
|
|
|
// Unlock the session's mutex.
|
|
|
|
|
// Expiry was checked and prolonged at the beginning of the request
|
|
|
|
|
let mut instance = store_rg |
|
|
|
|
.sessions |
|
|
|
|
.get(self.id.as_str()) |
|
|
|
|
.expect("Session data unexpectedly missing") |
|
|
|
|
.lock(); |
|
|
|
|
|
|
|
|
|
pub fn tap_mut<T>(&self, func: impl FnOnce(&mut D) -> T) -> T { |
|
|
|
|
let mut wg = self.store.inner.write(); |
|
|
|
|
let instance = wg.get_mut(self.id.as_str()).unwrap(); |
|
|
|
|
func(&mut instance.data) |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Fairing struct
|
|
|
|
|
struct SessionFairing<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
#[derive(Default)] |
|
|
|
|
pub struct SessionFairing<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
lifespan: Duration, |
|
|
|
|
_phantom: PhantomData<D>, |
|
|
|
|
config: SessionConfig, |
|
|
|
|
phantom: PhantomData<D>, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl<D> SessionFairing<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
fn new() -> Self { |
|
|
|
|
Self::default() |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Set session lifetime (expiration time).
|
|
|
|
|
///
|
|
|
|
|
/// Call on the fairing before passing it to `rocket.attach()`
|
|
|
|
|
pub fn with_lifetime(mut self, time: Duration) -> Self { |
|
|
|
|
self.config.lifespan = time; |
|
|
|
|
self |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Set session cookie name and length
|
|
|
|
|
///
|
|
|
|
|
/// Call on the fairing before passing it to `rocket.attach()`
|
|
|
|
|
pub fn with_cookie_name(mut self, name: impl Into<Cow<'static, str>>) -> Self { |
|
|
|
|
self.config.cookie_name = name.into(); |
|
|
|
|
self |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Set session cookie name and length
|
|
|
|
|
///
|
|
|
|
|
/// Call on the fairing before passing it to `rocket.attach()`
|
|
|
|
|
pub fn with_cookie_len(mut self, length: usize) -> Self { |
|
|
|
|
self.config.cookie_len = length; |
|
|
|
|
self |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Set session cookie name and length
|
|
|
|
|
///
|
|
|
|
|
/// Call on the fairing before passing it to `rocket.attach()`
|
|
|
|
|
pub fn with_cookie_path(mut self, path: impl Into<Cow<'static, str>>) -> Self { |
|
|
|
|
self.config.cookie_path = path.into(); |
|
|
|
|
self |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl<D> Fairing for SessionFairing<D> |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
where |
|
|
|
|
D: 'static + Sync + Send + Default, |
|
|
|
|
{ |
|
|
|
|
fn info(&self) -> Info { |
|
|
|
|
Info { |
|
|
|
|
@ -190,7 +309,7 @@ impl<D> Fairing for SessionFairing<D> |
|
|
|
|
// install the store singleton
|
|
|
|
|
Ok(rocket.manage(SessionStore::<D> { |
|
|
|
|
inner: Default::default(), |
|
|
|
|
lifespan: self.lifespan, |
|
|
|
|
config: self.config.clone(), |
|
|
|
|
})) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
@ -199,7 +318,11 @@ impl<D> Fairing for SessionFairing<D> |
|
|
|
|
let session = request.local_cache(|| SessionID("".to_string())); |
|
|
|
|
|
|
|
|
|
if !session.0.is_empty() { |
|
|
|
|
response.adjoin_header(Cookie::build(SESSION_COOKIE, session.0.clone()).finish()); |
|
|
|
|
response.adjoin_header( |
|
|
|
|
Cookie::build(self.config.cookie_name.clone(), session.to_string()) |
|
|
|
|
.path("/") |
|
|
|
|
.finish(), |
|
|
|
|
); |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
Ondřej Hruška