parent
52676cf227
commit
d24af144fa
GPG Key ID:
2C5FD5035250423D
@ -0,0 +1,168 @@ |
||||
use rocket::request::FromRequest; |
||||
use rocket::{Outcome, Request, State, http::{Status, Cookies, Cookie}, Response, Data, Rocket}; |
||||
use std::ops::{Deref, DerefMut}; |
||||
use std::collections::HashMap; |
||||
use parking_lot::{Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard}; |
||||
use serde_json::{Value, Map}; |
||||
use rocket::fairing::{self, Fairing, Info}; |
||||
use rand::Rng; |
||||
use std::borrow::Cow; |
||||
use serde::{Deserialize, Serialize}; |
||||
use serde::de::DeserializeOwned; |
||||
use json_dotpath::DotPaths; |
||||
use rocket::response::ResponseBuilder; |
||||
|
||||
const SESSION_ID : &'static str = "SESSID"; |
||||
|
||||
type SessionsMap = HashMap<String, SessionInstance>; |
||||
|
||||
#[derive(Debug)] |
||||
struct SessionInstance { |
||||
data: serde_json::Map<String, Value>, |
||||
// TODO expiration
|
||||
} |
||||
|
||||
#[derive(Default, Debug)] |
||||
pub struct SessionStore { |
||||
inner: RwLock<SessionsMap>, |
||||
} |
||||
|
||||
#[derive(PartialEq, Hash, Clone, Debug)] |
||||
pub struct SessionID(String); |
||||
|
||||
impl<'a, 'r> FromRequest<'a, 'r> for &'a SessionID { |
||||
type Error = (); |
||||
|
||||
fn from_request(request: &'a Request<'r>) -> Outcome<Self, (Status, Self::Error), ()> { |
||||
Outcome::Success(request.local_cache(|| { |
||||
println!("get id"); |
||||
if let Some(cookie) = request.cookies().get(SESSION_ID) { |
||||
println!("from cookie"); |
||||
SessionID(cookie.value().to_string()) // FIXME avoid cloning
|
||||
} else { |
||||
println!("new id"); |
||||
SessionID(rand::thread_rng() |
||||
.sample_iter(&rand::distributions::Alphanumeric) |
||||
.take(16) |
||||
.collect()) |
||||
} |
||||
})) |
||||
} |
||||
} |
||||
|
||||
#[derive(Debug)] |
||||
pub struct Session<'a> { |
||||
store: State<'a, SessionStore>, |
||||
id : &'a SessionID, |
||||
} |
||||
|
||||
impl<'a, 'r> FromRequest<'a, 'r> for Session<'a> { |
||||
type Error = (); |
||||
|
||||
fn from_request(request: &'a Request<'r>) -> Outcome<Self, (Status, Self::Error), ()> { |
||||
Outcome::Success(Session { |
||||
id: request.local_cache(|| { |
||||
if let Some(cookie) = request.cookies().get(SESSION_ID) { |
||||
SessionID(cookie.value().to_string()) |
||||
} else { |
||||
SessionID(rand::thread_rng() |
||||
.sample_iter(&rand::distributions::Alphanumeric) |
||||
.take(16) |
||||
.collect()) |
||||
} |
||||
}), |
||||
store: request.guard().unwrap() |
||||
}) |
||||
} |
||||
} |
||||
|
||||
impl<'a> Session<'a> { |
||||
pub fn fairing() -> impl Fairing { |
||||
SessionFairing |
||||
} |
||||
|
||||
pub fn get<T : DeserializeOwned>(&self, path : &str) -> Option<T> { |
||||
let rg = self.store.inner.read(); |
||||
if let Some(ses) = rg.get(&self.id.0) { |
||||
ses.data.dot_get(path) |
||||
} else { |
||||
None |
||||
} |
||||
} |
||||
|
||||
pub fn get_or<T : DeserializeOwned>(&self, path : &str, def : T) -> T { |
||||
self.get(path).unwrap_or(def) |
||||
} |
||||
|
||||
pub fn get_or_else<T : DeserializeOwned, F : FnOnce() -> T>(&self, path : &str, def : F) -> T { |
||||
self.get(path).unwrap_or_else(def) |
||||
} |
||||
|
||||
pub fn get_or_default<T : DeserializeOwned + Default>(&self, path : &str) -> T { |
||||
self.get(path).unwrap_or_default() |
||||
} |
||||
|
||||
pub fn take<T : DeserializeOwned>(&self, path : &str) -> Option<T> { |
||||
let mut wg = self.store.inner.write(); |
||||
if let Some(ses) = wg.get_mut(&self.id.0) { |
||||
ses.data.dot_take(path) |
||||
} else { |
||||
None |
||||
} |
||||
} |
||||
|
||||
pub fn replace<O: DeserializeOwned, N: Serialize>(&self, path : &str, new : N) -> Option<O> { |
||||
let mut wg = self.store.inner.write(); |
||||
if let Some(ses) = wg.get_mut(&self.id.0) { |
||||
ses.data.dot_replace(path, new) |
||||
} else { |
||||
None |
||||
} |
||||
} |
||||
|
||||
pub fn set<T : Serialize>(&self, path : &str, value : T) { |
||||
let mut wg = self.store.inner.write(); |
||||
if let Some(ses) = wg.get_mut(&self.id.0) { |
||||
ses.data.dot_set(path, value); |
||||
} else { |
||||
let mut map = Map::new(); |
||||
map.dot_set(path, value); |
||||
wg.insert(self.id.0.clone(), SessionInstance { |
||||
data : map, |
||||
}); |
||||
} |
||||
} |
||||
|
||||
pub fn remove(&self, path : &str) { |
||||
let mut wg = self.store.inner.write(); |
||||
if let Some(ses) = wg.get_mut(&self.id.0) { |
||||
ses.data.dot_remove(path); |
||||
} |
||||
} |
||||
} |
||||
|
||||
/// Fairing struct
|
||||
struct SessionFairing; |
||||
|
||||
impl Fairing for SessionFairing { |
||||
fn info(&self) -> Info { |
||||
Info { |
||||
name: "Session Fairing", |
||||
kind: fairing::Kind::Attach | fairing::Kind::Response |
||||
} |
||||
} |
||||
|
||||
fn on_attach(&self, rocket: Rocket) -> Result<Rocket, Rocket> { |
||||
Ok(rocket.manage(SessionStore::default())) |
||||
} |
||||
|
||||
fn on_response<'r>(&self, request: &'r Request, response: &mut Response) { |
||||
let session = request.local_cache(|| { |
||||
SessionID("".to_string()) |
||||
}); |
||||
|
||||
if !session.0.is_empty() { |
||||
response.adjoin_header(Cookie::build(SESSION_ID, session.0.clone()).finish()); |
||||
} |
||||
} |
||||
} |
||||
Loading…
Reference in new issue