pg_ntex_session/
lib.rs

1//! User sessions.
2//!
3//!
4//! By default, only cookie session backend is implemented. This crate
5//! provides PostgreSQL back end with Diesel's engine
6//!
7//! In general, you insert a *session* middleware and initialize it
8//! , such as a `PgNtexSession`. To access session data,
9//! [*Session*](https://docs.rs/ntex-session/latest/ntex_session/struct.Session.html) extractor must be used. Session
10//! extractor allows us to get or set session data.
11//!
12//! ```rust,no_run
13//! use ntex::web::{self, App, HttpResponse, Error};
14//! use pg_ntex_session::{encryption::Simple, PgNtexSession};
15//! use diesel::{r2d2::{ConnectionManager, Pool}, PgConnection};
16//! use std::sync::Arc;
17//!
18//! fn index(session: Session) -> Result<&'static str, Error> {
19//!     // access session data
20//!     if let Some(count) = session.get::<i32>("counter")? {
21//!         println!("SESSION value: {}", count);
22//!         session.set("counter", count+1)?;
23//!     } else {
24//!         session.set("counter", 1)?;
25//!     }
26//!
27//!     Ok("Welcome!")
28//! }
29//!
30//! //DB connection sting
31//! const DATABASE_URL:&'static str="postgres://postgres:password@localhost/my_db";
32//! //32 characters
33//! const ENC_KEY:&'static str="29dba93e4ce64609bb5d592dab92ec00";
34//! //Pool of connection manager
35//! const :Arc<Pool<ConnectionManager<PgConnection>>>=Arc::new(
36//!     Pool::builder().max_size(16)
37//!        .build(ConnectionManager::<PgConnection>::new(DATABASE_URL))
38//!        .unwrap()
39//! );
40//! 
41//! #[ntex::main]
42//! async fn main() -> std::io::Result<()> {
43//!     web::server(
44//!         || App::new().wrap(
45//!               <PgNtexSession<Simple>>::new(ENC_KEY.as_bytes(), Some(Simple::new(ENC_KEY.as_str())), CONNECTION.clone()))
46//!              )
47//!             .service(web::resource("/").to(|| async { HttpResponse::Ok() })))
48//!         .bind("127.0.0.1:59880")?
49//!         .run()
50//!         .await
51//! }
52//! ```
53pub mod encryption;
54
55use chrono::Local;
56use cookie::{Cookie, CookieJar, Key, time::{OffsetDateTime, Duration as TimeDuration}, SameSite};
57use ntex::{http::{header::{HeaderValue, COOKIE, SET_COOKIE, USER_AGENT}, HttpMessage}, web::{DefaultError, ErrorRenderer, HttpRequest, WebRequest, WebResponse, WebResponseError}};
58use ntex_session::{Session, SessionStatus};
59use uuid::Uuid;
60use std::{collections::HashMap, fmt::Display, rc::Rc, sync::Arc, time::Duration};
61use once_cell::sync::OnceCell;
62use diesel::{prelude::*, r2d2::{ConnectionManager, Pool}, sql_query, PgConnection, RunQueryDsl};
63use ntex::service::{Middleware, Service, ServiceCtx};
64use encryption::*;
65use base64::prelude::*;
66
67#[derive(Debug)]
68pub struct PgSessError(anyhow::Error);
69
70impl PgSessError {
71    pub fn new(error:anyhow::Error) -> Self
72    {
73        Self(error)
74    }
75}
76
77impl Display for PgSessError
78{
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        write!(f,"{}",self.0)
81    }
82}
83
84impl std::error::Error for PgSessError {}
85
86impl WebResponseError<DefaultError> for PgSessError {}
87
88// impl ResponseError for PgSessError
89// {
90//     fn error_response(&self) -> ntex::http::Response {
91//         ErrorInternalServerError(&self.0).error_response()
92//     }
93// }
94
95impl From<anyhow::Error> for PgSessError {
96    fn from(value: anyhow::Error) -> Self {
97        Self(value)
98    }
99}
100
101#[derive(Debug, QueryableByName)]
102struct PgSessionRow
103{
104    #[diesel(sql_type = diesel::sql_types::Text)]
105    sess_state:String
106}
107
108#[derive(QueryableByName)]
109struct ExistsResult {
110    #[diesel(sql_type = diesel::sql_types::Bool)]
111    exists: bool,
112}
113
114struct PgSessionInner<SE:StateEncryption + Clone> {
115    encryption_engine:Option<Box<SE>>,
116    key: Key,
117    name: String,
118    path: String,
119    domain: Option<String>,
120    secure: bool,
121    http_only: bool,
122    max_age: Option<Duration>,
123    expires_in: Option<Duration>,
124    same_site: Option<SameSite>,
125}
126
127impl <SE:StateEncryption + Clone> PgSessionInner<SE>
128{
129    fn new(key: &[u8],ee:Option<SE>) -> Self {
130        Self {
131            encryption_engine:ee.clone().map(|v| Box::new(v)),
132            key: Key::derive_from(key),
133            name: "nx-sess".to_owned(),
134            path: "/".to_owned(),
135            domain: None,
136            secure: true,
137            http_only: true,
138            max_age: None,
139            expires_in: None,
140            same_site: None,
141        }
142    }
143
144    fn check_table(&self)->anyhow::Result<()>
145    {
146        if let Some(pg)=PG_CONNECTION.get()
147        {
148            if let Ok(mut conn)=pg.get() {
149                if let Ok(res,)=sql_query("SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'web_sessions')").get_result::<ExistsResult>(&mut conn)
150                {
151                    if !res.exists {
152                        let query=sql_query("CREATE TABLE web_sessions(id VARCHAR(32) NOT NULL, sess_state TEXT DEFAULT NULL, user_agent TEXT NOT NULL, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, expired_at TIMESTAMP NOT NULL, CONSTRAINT pk_web_sessions PRIMARY KEY(id))");
153                        query.execute(&mut conn)?;
154                        let query=sql_query("CREATE INDEX IF NOT EXISTS idx_web_sessions_exp ON web_sessions(expired_at)");
155                        query.execute(&mut conn)?;
156                    }
157                }
158            }
159        }
160
161        Ok(())
162    }
163
164    fn fetch_session(&self,session_id:impl AsRef<str>) -> anyhow::Result<Option<PgSessionRow>>
165    {
166        let session_id=session_id.as_ref();
167
168        if let Some(pg)=PG_CONNECTION.get()
169        {
170            let mut conn=pg.get()?;
171            Ok(sql_query(format!("SELECT sess_state FROM web_sessions WHERE id='{}'",session_id)).get_result::<PgSessionRow>(&mut conn).optional()?)
172        } else {
173            Err(anyhow::Error::msg("OnceCell not initialized"))
174        }
175    }
176
177    fn load<Err>(&self, req: &WebRequest<Err>) -> anyhow::Result<(bool, String, HashMap<String, String>)>
178    {
179        self.check_table()?;
180        let mut session_id=Uuid::new_v4().simple().to_string();
181        
182        if let Ok(cookies) = req.cookies() {
183            for cookie in cookies.iter() {
184                if cookie.name() == self.name {
185                    let mut jar = CookieJar::new();
186                    jar.add_original(cookie.clone());
187
188                    let cookie_opt = jar.signed(&self.key).get(&self.name);
189                    if let Some(cookie) = cookie_opt {
190                        if let Ok(val) = serde_json::from_str::<HashMap<String, String>>(cookie.value()) {
191                            if let Some(sid)=val.get("sid") {
192                                session_id=sid.clone();
193
194                                if let Some(row)=self.fetch_session(&session_id).expect("Unable to fetch session row") {
195                                    let mut s=row.sess_state.clone();
196                                    
197                                    if let Some(en)=self.encryption_engine.as_ref() {
198                                        s=en.decrypt(BASE64_STANDARD.decode(&s)?)?;
199                                    }
200
201                                    let val=serde_json::from_str::<HashMap<String,String>>(&s)?;
202                                    return Ok((false, session_id,val));
203                                }
204                            }
205                        }
206                    }
207                }
208            }
209        }
210
211        Ok((true, session_id, HashMap::new()))
212    }
213
214    fn remove(&self, res: &mut WebResponse, session_id:impl AsRef<str>)->Result<(),PgSessError>
215    {
216        let session_id=session_id.as_ref();
217
218        if let Some(db)=PG_CONNECTION.get() {
219            let mut conn=db.get().map_err(|err|PgSessError::new(anyhow::Error::new(err)))?;
220            sql_query(format!("DELETE FROM web_sessions WHERE id='{}'",session_id)).execute(&mut conn).map_err(|err|PgSessError::new(anyhow::Error::new(err)))?;
221        }
222
223        let mut cookie = Cookie::from(self.name.clone());
224        cookie.set_value("");
225        cookie.set_max_age(TimeDuration::ZERO);
226        cookie.set_expires(OffsetDateTime::now_utc() - TimeDuration::days(365));
227
228        let val = HeaderValue::from_str(&cookie.to_string()).unwrap();
229        res.headers_mut().append(SET_COOKIE, val);
230
231        Ok(())
232    }
233
234    fn save(&self,session_id:impl Display,wres:&mut WebResponse,states:impl Iterator<Item = (String,String)>)->Result<(),PgSessError>
235    {
236        let session_id=session_id.to_string();
237        let states=<HashMap<String,String>>::from_iter(states);
238        let mut state=serde_json::to_string(&states).map_err(|err|PgSessError::new(anyhow::Error::new(err)))?;
239
240        if let Some(en)=self.encryption_engine.as_ref() {
241            let val=en.encrypt(&state)?;
242            state=BASE64_STANDARD.encode(&val);
243        }
244
245        let mut cookie = Cookie::new(self.name.clone(), format!(r#"{{"sid":"{}"}}"#,&session_id));
246        cookie.set_path(self.path.clone());
247        cookie.set_secure(self.secure);
248        cookie.set_http_only(self.http_only);
249
250        if let Some(ref domain) = self.domain {
251            cookie.set_domain(domain.clone());
252        }
253
254        let expired_at=if let Some(expires_in) = self.expires_in {
255            cookie.set_expires(OffsetDateTime::now_utc() + expires_in);
256            Local::now() + expires_in
257        } else {
258            Local::now() + Duration::from_secs(3600)
259        }.naive_local();
260        
261        if let Some(db)=PG_CONNECTION.get() {
262            let now=Local::now();
263            let dt_format="%Y-%m-%d %H:%M:%S%.9f";
264            let query=sql_query(format!("INSERT INTO web_sessions (id, sess_state, user_agent, created_at, expired_at) VALUES('{}', '{}', '{}', '{}', '{}') ON CONFLICT(id) DO UPDATE SET sess_state=EXCLUDED.sess_state, expired_at=EXCLUDED.expired_at",&session_id,&state,wres.headers().get(USER_AGENT).and_then(|v|v.to_str().ok().map(|s|s.to_string())).unwrap_or("ntex::web".to_string()),now.format(dt_format).to_string(),expired_at.format(dt_format).to_string()));
265            let mut conn=db.get().map_err(|err|PgSessError::new(anyhow::Error::new(err)))?;
266            query.execute(&mut conn).map_err(|err|PgSessError::new(anyhow::Error::new(err)))?;
267        }
268
269        let mut jar = CookieJar::new();
270        jar.signed_mut(&self.key).add(cookie);
271
272        for cookie in jar.delta() {
273            let val = HeaderValue::from_str(&cookie.encoded().to_string()).unwrap();
274            wres.headers_mut().append(SET_COOKIE, val);
275        }
276
277        Ok(())
278    }
279}
280
281pub struct PgNtexSession<SE:StateEncryption + Clone>(Rc<PgSessionInner<SE>>);
282
283impl <SE:StateEncryption + Clone> PgNtexSession<SE>
284{
285/// Construct new *signed* `PgNtexSession` instance.
286///
287/// You can use UUID v4 as encryption key.
288/// Panics if key length is less than 32 bytes.
289/// ```pooled_connection``` is a Pool of Postgres Connection Manager. Check [*r2d2 Pool sample*](https://docs.rs/r2d2/latest/r2d2/) and [*ConnectionManager sample*](https://docs.rs/diesel/latest/diesel/r2d2/struct.ConnectionManager.html#method.new)
290    pub fn new(key:impl AsRef<[u8]>,encryption_engine:Option<SE>,pooled_connection:Arc<Pool<ConnectionManager<PgConnection>>>)->Self
291    {
292        let key=key.as_ref();
293        PG_CONNECTION.get_or_init(||pooled_connection);
294        COOKIE_ENC_KEY.get_or_init(||key.to_vec());
295        Self(Rc::new(PgSessionInner::new(key,encryption_engine)))
296    }
297
298    /// Sets the `path` field in the session cookie being built.
299    pub fn path<S: Into<String>>(mut self, value: S) -> Self {
300        Rc::get_mut(&mut self.0).unwrap().path = value.into();
301        self
302    }
303
304    /// Sets the `name` field in the session cookie being built.
305    pub fn name<S: Into<String>>(mut self, value: S) -> Self {
306        Rc::get_mut(&mut self.0).unwrap().name = value.into();
307        self
308    }
309
310    /// Sets the `domain` field in the session cookie being built.
311    pub fn domain<S: Into<String>>(mut self, value: S) -> Self {
312        Rc::get_mut(&mut self.0).unwrap().domain = Some(value.into());
313        self
314    }
315
316    /// Sets the `secure` field in the session cookie being built.
317    ///
318    /// If the `secure` field is set, a cookie will only be transmitted when the
319    /// connection is secure - i.e. `https`
320    pub fn secure(mut self, value: bool) -> Self {
321        Rc::get_mut(&mut self.0).unwrap().secure = value;
322        self
323    }
324
325    /// Sets the `http_only` field in the session cookie being built.
326    pub fn http_only(mut self, value: bool) -> Self {
327        Rc::get_mut(&mut self.0).unwrap().http_only = value;
328        self
329    }
330
331    /// Sets the `same_site` field in the session cookie being built.
332    pub fn same_site(mut self, value: SameSite) -> Self {
333        Rc::get_mut(&mut self.0).unwrap().same_site = Some(value);
334        self
335    }
336
337    /// Sets the `max-age` field in the session cookie being built.
338    pub fn max_age(self, seconds: u64) -> Self {
339        self.max_age_time(Duration::from_secs(seconds))
340    }
341
342    /// Sets the `max-age` field in the session cookie being built.
343    pub fn max_age_time(mut self, value: Duration) -> Self {
344        Rc::get_mut(&mut self.0).unwrap().max_age = Some(value);
345        self
346    }
347
348    /// Sets the `expires` field in the session cookie being built.
349    pub fn expires_in(self, seconds: u64) -> Self {
350        self.expires_in_time(Duration::from_secs(seconds))
351    }
352
353    /// Sets the `expires` field in the session cookie being built.
354    pub fn expires_in_time(mut self, value: Duration) -> Self {
355        Rc::get_mut(&mut self.0).unwrap().expires_in = Some(value);
356        self
357    }
358}
359
360impl<S,SE:StateEncryption + Clone> Middleware<S> for PgNtexSession<SE> {
361    type Service = PgNtexSessionMiddleware<S,SE>;
362
363    fn create(&self, service: S) -> Self::Service {
364        PgNtexSessionMiddleware { service, inner: self.0.clone() }
365    }
366}
367
368/// Session middleware with Postgres backend for Ntex web framework
369pub struct PgNtexSessionMiddleware<S,SE: StateEncryption + Clone> {
370    service: S,
371    inner: Rc<PgSessionInner<SE>>,
372}
373
374impl<S, Err,SE:StateEncryption + Clone> Service<WebRequest<Err>> for PgNtexSessionMiddleware<S,SE>
375where
376    S: Service<WebRequest<Err>, Response = WebResponse>,
377    S::Error: 'static,
378    Err: ErrorRenderer,
379    Err::Container: From<PgSessError>,
380{
381    type Response = WebResponse;
382    type Error = S::Error;
383
384    ntex::forward_poll!(service);
385    ntex::forward_ready!(service);
386    ntex::forward_shutdown!(service);
387
388    /// On first request, a new session cookie is returned in response, regardless
389    /// of whether any session state is set.  This cookie only stores one field:
390    /// ```sid``` which is a session ID, that points to primary key to PostgreSQL
391    /// table. With subsequent requests, if the session state changes, then
392    /// set-cookie is returned in response.  As a user logs out, call
393    /// session.purge() to set SessionStatus accordingly and this will trigger
394    /// removal of the session cookie in the response.
395    async fn call(&self,req: WebRequest<Err>, ctx: ServiceCtx<'_, Self>) -> Result<Self::Response, Self::Error>
396    {
397        let inner = self.inner.clone();
398        let (is_new, session_id, state) = self.inner.load(&req).expect("Error loading session");
399
400        let prolong_expiration = self.inner.expires_in.is_some();
401        Session::set_session(state.into_iter(), &req);
402        clean_expired_session();
403        
404        ctx.call(&self.service, req).await.map(|mut res| {
405            match Session::get_changes(&mut res) {
406                (SessionStatus::Changed, Some(state))
407                |(SessionStatus::Renewed, Some(state))=>res.checked_expr::<Err, _, _>(|res| {
408                    inner.save(&session_id, res, state)
409                }),
410                (SessionStatus::Unchanged, Some(state)) if prolong_expiration =>  {
411                    res.checked_expr::<Err, _, _>(|res| {
412                        inner.save(&session_id, res, state)
413                    })
414                },
415                (SessionStatus::Unchanged, _) => if is_new  {
416                    let state: HashMap<String, String> = HashMap::new();
417                    res.checked_expr::<Err, _, _>(|res| {
418                        inner.save(&session_id, res, state.iter().map(|(k,v)|(k.clone(),v.clone())))
419                    })
420                } else {
421                    res
422                },
423                (SessionStatus::Purged, _) => {
424                    let _ = inner.remove(&mut res, &session_id);
425                    res
426                },
427                _ => res
428            }
429        })
430    }
431}
432
433/// Returns saved session ID, in form of 32 hex characters, or returns
434/// ```None``` instead, if you haven't saved your states in this 
435/// current session
436pub fn get_session_id<'a>(http_req:&'a HttpRequest)->Option<String>
437{
438    if let Some(cookie_header)=http_req.headers().get(COOKIE)
439    {
440        if let Ok(cookie_str) = cookie_header.to_str().map(|v|v.to_string()) {
441            for cookie in cookie_str.split(';').map(|s|s.trim().to_string()) {
442                if let Ok(parsed_cookie) = Cookie::parse(cookie) {
443                    if parsed_cookie.name() == "nx-sess" {
444                        let mut jar = CookieJar::new();
445                        jar.add_original(parsed_cookie.clone());
446                        
447                        if let Some(key)=COOKIE_ENC_KEY.get() {
448                            let cookie_key=Key::derive_from(key.as_slice());
449                            if let Some(cookie)=jar.signed(&cookie_key).get("nx-sess") {
450                                if let Ok(val) = serde_json::from_str::<HashMap<String, String>>(cookie.value()) {
451                                    return val.get("sid").cloned();
452                                }
453                            }
454                        }
455                    }
456                }
457            }
458        }
459    }
460
461    None
462}
463
464/// Clean expired sessions in database, which tolerates +3 minutes
465/// from actual ```expired_at``` value.
466/// You can use it with timer to clean it periodically.
467pub fn clean_expired_session()
468{
469    let now_minus=Local::now() - Duration::from_secs(180);
470    
471    if let Some(pool)=PG_CONNECTION.get() {
472        if let Ok(mut conn)=pool.get() {
473            let _=sql_query(format!("DELETE FROM web_sessions WHERE expiired_at <= '{}'",now_minus.format("%Y-%m-%d %H:%M:%S%.9f").to_string())).execute(&mut conn);
474        }
475    }
476}
477
478static PG_CONNECTION: OnceCell<Arc<Pool<ConnectionManager<PgConnection>>>>= OnceCell::new();
479static COOKIE_ENC_KEY:OnceCell<Vec<u8>>=OnceCell::new();
pg_ntex_session/lib.rs

pg_ntex_session/
lib.rs
