caldwell
/
rocket_session
forked from packages/rocket_session
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add examples, automatic expired removal, better configurability

Browse Source
refactor
Ondřej Hruška 5 years ago
parent af30c552ac
commit cde08fe788
Signed by untrusted user: MightyPork
GPG Key ID: 2C5FD5035250423D
3 changed files with 343 additions and 85 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 70
      examples/dog_list/main.rs
  2. 72
      examples/visit_counter/main.rs
  3. 286
      src/lib.rs

70
examples/dog_list/main.rs
Unescape 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
Unescape 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
))
}

286
src/lib.rs
Unescape View File

@ -1,4 +1,4 @@
use parking_lot::RwLock;
use parking_lot::{RwLock, RwLockUpgradableReadGuard, Mutex};
use rand::Rng;
use rocket::{
@ -12,60 +12,100 @@ use std::collections::HashMap;
use std::marker::PhantomData;
use std::ops::Add;
use std::time::{Duration, Instant};
use std::borrow::Cow;
use std::fmt::{Display, Formatter, self};
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>
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>
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
#[derive(Debug)]
struct SessionInstance<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);
}
/// 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
@ -84,45 +124,76 @@ impl<'a, 'r, D> FromRequest<'a, 'r> for Session<'a, D>
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,
})
@ -133,46 +204,90 @@ impl<'a, D> Session<'a, D>
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>
#[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>
@ -190,7 +305,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 +314,8 @@ 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());
}
}
}

Write Preview
Loading…
Cancel
Save