YinMo19/rocket_session
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add examples, automatic expired removal, better configurability

Browse Source
This commit is contained in:
Ondřej Hruška 2019-12-31 20:31:38 +01:00
parent af30c552ac
commit cde08fe788
No known key found for this signature in database
GPG Key ID: 2C5FD5035250423D
3 changed files with 345 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
examples/dog_list/main.rs Normal file
View File

@ -0,0 +1,70 @@
#![feature(proc_macro_hygiene, decl_macro)]
#[macro_use]
extern crate rocket;
use rocket::response::content::Html;
use rocket::response::Redirect;
use rocket::request::Form;
type Session<'a> = rocket_session::Session<'a, Vec<String>>;
fn main() {
rocket::ignite()
.attach(Session::fairing())
.mount("/", routes![index, add, remove])
.launch();
}
#[get("/")]
fn index(session: Session) -> Html<String> {
let mut page = String::new();
page.push_str(r#"
<!DOCTYPE html>
<h1>My Dogs</h1>
<form method="POST" action="/add">
Add Dog: <input type="text" name="name"> <input type="submit" value="Add">
</form>
<ul>
"#);
session.tap(|sess| {
for (n, dog) in sess.iter().enumerate() {
page.push_str(&format!(r#"
<li>&#x1F436; {} <a href="/remove/{}">Remove</a></li>
"#, dog, n));
}
});
page.push_str(r#"
</ul>
"#);
Html(page)
}
#[derive(FromForm)]
struct AddForm {
name: String,
}
#[post("/add", data="<dog>")]
fn add(session: Session, dog : Form<AddForm>) -> Redirect {
session.tap(move |sess| {
sess.push(dog.into_inner().name);
});
Redirect::found("/")
}
#[get("/remove/<dog>")]
fn remove(session: Session, dog : usize) -> Redirect {
session.tap(|sess| {
if dog < sess.len() {
sess.remove(dog);
}
});
Redirect::found("/")
}

72
examples/visit_counter/main.rs Normal file
View File

@ -0,0 +1,72 @@
//! This demo is a page visit counter, with a custom cookie name, length, and expiry time.
//!
//! The expiry time is set to 10 seconds to illustrate how a session is cleared if inactive.
#![feature(proc_macro_hygiene, decl_macro)]
#[macro_use]
extern crate rocket;
use std::time::Duration;
use rocket::response::content::Html;
#[derive(Default, Clone)]
struct SessionData {
visits1: usize,
visits2: usize,
}
// It's convenient to define a type alias:
type Session<'a> = rocket_session::Session<'a, SessionData>;
fn main() {
rocket::ignite()
.attach(Session::fairing()
// 10 seconds of inactivity until session expires
// (wait 10s and refresh, the numbers will reset)
.with_lifetime(Duration::from_secs(10))
// custom cookie name and length
.with_cookie_name("my_cookie")
.with_cookie_len(20)
)
.mount("/", routes![index, about])
.launch();
}
#[get("/")]
fn index(session: Session) -> Html<String> {
// Here we build the entire response inside the 'tap' closure.
// While inside, the session is locked to parallel changes, e.g.
// from a different browser tab.
session.tap(|sess| {
sess.visits1 += 1;
Html(format!(r##"
<!DOCTYPE html>
<h1>Home</h1>
<a href="/">Refresh</a> &bull; <a href="/about/">go to About</a>
<p>Visits: home {}, about {}</p>
"##,
sess.visits1,
sess.visits2
))
})
}
#[get("/about")]
fn about(session: Session) -> Html<String> {
// Here we return a value from the tap function and use it below
let count = session.tap(|sess| {
sess.visits2 += 1;
sess.visits2
});
Html(format!(r##"
<!DOCTYPE html>
<h1>About</h1>
<a href="/about">Refresh</a> &bull; <a href="/">go home</a>
<p>Page visits: {}</p>
"##,
count
))
}

288
src/lib.rs
View File

@ -1,4 +1,4 @@
use parking_lot::RwLock; use parking_lot::{RwLock, RwLockUpgradableReadGuard, Mutex};
use rand::Rng; use rand::Rng;
use rocket::{ use rocket::{
@ -12,9 +12,65 @@ use std::collections::HashMap;
use std::marker::PhantomData; use std::marker::PhantomData;
use std::ops::Add; use std::ops::Add;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use std::borrow::Cow;
use std::fmt::{Display, Formatter, self};
const SESSION_COOKIE: &str = "SESSID"; /// Session store (shared state)
const SESSION_ID_LEN: usize = 16; #[derive(Debug)]
pub struct SessionStore<D>
where
D: 'static + Sync + Send + Default,
{
/// The internally mutable map of sessions
inner: RwLock<StoreInner<D>>,
// Session config
config: SessionConfig,
}
/// 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 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 {
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 /// Session, as stored in the sessions store
#[derive(Debug)] #[derive(Debug)]
@ -28,44 +84,28 @@ struct SessionInstance<D>
expires: Instant, expires: Instant,
} }
/// 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>>>,
/// Sessions lifespan
lifespan: Duration,
}
impl<D> SessionStore<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);
}
}
/// Session ID newtype for rocket's "local_cache" /// Session ID newtype for rocket's "local_cache"
#[derive(PartialEq, Hash, Clone, Debug)] #[derive(Clone, Debug)]
struct SessionID(String); struct SessionID(String);
impl SessionID { impl SessionID {
fn as_str(&self) -> &str { fn as_str(&self) -> &str {
self.0.as_str() self.0.as_str()
} }
}
fn to_string(&self) -> String { impl Display for SessionID {
self.0.clone() fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
} }
} }
/// Session instance /// 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)] #[derive(Debug)]
pub struct Session<'a, D> pub struct Session<'a, D>
where where
@ -84,45 +124,76 @@ impl<'a, 'r, D> FromRequest<'a, 'r> for Session<'a, D>
type Error = (); type Error = ();
fn from_request(request: &'a Request<'r>) -> Outcome<Self, (Status, Self::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 { Outcome::Success(Session {
id: request.local_cache(|| { id: request.local_cache(|| {
let store_ug = store.inner.upgradable_read();
// Resolve session ID // Resolve session ID
let id = if let Some(cookie) = request.cookies().get(SESSION_COOKIE) { let id = if let Some(cookie) = request.cookies().get(&store.config.cookie_name) {
SessionID(cookie.value().to_string()) Some(SessionID(cookie.value().to_string()))
} else { } else {
SessionID( None
rand::thread_rng() };
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();
}
// 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) .sample_iter(&rand::distributions::Alphanumeric)
.take(SESSION_ID_LEN) .take(store.config.cookie_len)
.collect(), .collect();
)
};
let new_expiration = Instant::now().add(store.lifespan); if !store_wg.sessions.contains_key(&token) {
let mut wg = store.inner.write(); break token;
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,
} }
});
store_wg.sessions.insert(
new_id.to_string(),
Mutex::new(SessionInstance {
data: Default::default(),
expires,
}),
); );
}
};
id new_id
}
}), }),
store, store,
}) })
@ -133,46 +204,90 @@ impl<'a, D> Session<'a, D>
where where
D: 'static + Sync + Send + Default, D: 'static + Sync + Send + Default,
{ {
/// Get the fairing object /// Create the session fairing.
pub fn fairing(lifespan: Duration) -> impl Fairing { ///
SessionFairing::<D> { /// You can configure the session store by calling chained methods on the returned value
lifespan, /// before passing it to `rocket.attach()`
_phantom: PhantomData, pub fn fairing() -> SessionFairing<D> {
} SessionFairing::<D>::new()
} }
/// Access the session store /// Clear session data (replace the value with default)
pub fn get_store(&self) -> &SessionStore<D> { pub fn clear(&self) {
&self.store self.tap(|m| {
}
/// Set the session object to its default state
pub fn reset(&self) {
self.tap_mut(|m| {
*m = D::default(); *m = D::default();
}) })
} }
pub fn tap<T>(&self, func: impl FnOnce(&D) -> T) -> T { /// Access the session's data using a closure.
let rg = self.store.inner.read(); ///
let instance = rg.get(self.id.as_str()).unwrap(); /// The closure is called with the data value as a mutable argument,
func(&instance.data) /// 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) func(&mut instance.data)
} }
} }
/// Fairing struct /// Fairing struct
struct SessionFairing<D> #[derive(Default)]
pub struct SessionFairing<D>
where where
D: 'static + Sync + Send + Default, D: 'static + Sync + Send + Default,
{ {
lifespan: Duration, config: SessionConfig,
_phantom: PhantomData<D>, 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> impl<D> Fairing for SessionFairing<D>
@ -190,7 +305,7 @@ impl<D> Fairing for SessionFairing<D>
// install the store singleton // install the store singleton
Ok(rocket.manage(SessionStore::<D> { Ok(rocket.manage(SessionStore::<D> {
inner: Default::default(), inner: Default::default(),
lifespan: self.lifespan, config: self.config.clone(),
})) }))
} }
@ -199,7 +314,8 @@ impl<D> Fairing for SessionFairing<D>
let session = request.local_cache(|| SessionID("".to_string())); let session = request.local_cache(|| SessionID("".to_string()));
if !session.0.is_empty() { 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());
} }
} }
} }