1use std::{collections::HashMap, convert::Infallible, rc::Rc};
19
20use cookie::{Cookie, CookieJar, Key, SameSite};
21use derive_more::{Display, From};
22use ntex::http::{HttpMessage, header::HeaderValue, header::SET_COOKIE};
23use ntex::service::{Middleware, Service, ServiceCtx};
24use ntex::web::{DefaultError, ErrorRenderer, WebRequest, WebResponse, WebResponseError};
25use serde_json::error::Error as JsonError;
26use time::{Duration, OffsetDateTime};
27
28use crate::{Session, SessionStatus};
29
30#[derive(Debug, From, Display)]
32pub enum CookieSessionError {
33 #[display("Size of the serialized session is greater than 4000 bytes.")]
35 Overflow,
36 #[display("Fail to serialize session")]
38 Serialize(JsonError),
39}
40
41impl WebResponseError<DefaultError> for CookieSessionError {}
42
43enum CookieSecurity {
44 Signed,
45 Private,
46}
47
48struct CookieSessionInner {
49 key: Key,
50 security: CookieSecurity,
51 name: String,
52 path: String,
53 domain: Option<String>,
54 secure: bool,
55 http_only: bool,
56 max_age: Option<Duration>,
57 expires_in: Option<Duration>,
58 same_site: Option<SameSite>,
59}
60
61impl CookieSessionInner {
62 fn new(key: &[u8], security: CookieSecurity) -> Self {
63 CookieSessionInner {
64 security,
65 key: Key::derive_from(key),
66 name: "ntex-session".to_owned(),
67 path: "/".to_owned(),
68 domain: None,
69 secure: true,
70 http_only: true,
71 max_age: None,
72 expires_in: None,
73 same_site: None,
74 }
75 }
76
77 fn set_cookie(
78 &self,
79 res: &mut WebResponse,
80 state: impl Iterator<Item = (String, String)>,
81 ) -> Result<(), CookieSessionError> {
82 let state: HashMap<String, String> = state.collect();
83 let value = serde_json::to_string(&state).map_err(CookieSessionError::Serialize)?;
84 if value.len() > 4064 {
85 return Err(CookieSessionError::Overflow);
86 }
87
88 let mut cookie = Cookie::new(self.name.clone(), value);
89 cookie.set_path(self.path.clone());
90 cookie.set_secure(self.secure);
91 cookie.set_http_only(self.http_only);
92
93 if let Some(ref domain) = self.domain {
94 cookie.set_domain(domain.clone());
95 }
96
97 if let Some(expires_in) = self.expires_in {
98 cookie.set_expires(OffsetDateTime::now_utc() + expires_in);
99 }
100
101 if let Some(max_age) = self.max_age {
102 cookie.set_max_age(max_age);
103 }
104
105 if let Some(same_site) = self.same_site {
106 cookie.set_same_site(same_site);
107 }
108
109 let mut jar = CookieJar::new();
110
111 match self.security {
112 CookieSecurity::Signed => jar.signed_mut(&self.key).add(cookie),
113 CookieSecurity::Private => jar.private_mut(&self.key).add(cookie),
114 }
115
116 for cookie in jar.delta() {
117 let val = HeaderValue::from_str(&cookie.encoded().to_string()).unwrap();
118 res.headers_mut().append(SET_COOKIE, val);
119 }
120
121 Ok(())
122 }
123
124 fn remove_cookie(&self, res: &mut WebResponse) -> Result<(), Infallible> {
126 let mut cookie = Cookie::from(self.name.clone());
127 cookie.set_value("");
128 cookie.set_max_age(Duration::ZERO);
129 cookie.set_expires(OffsetDateTime::now_utc() - Duration::days(365));
130
131 let val = HeaderValue::from_str(&cookie.to_string()).unwrap();
132 res.headers_mut().append(SET_COOKIE, val);
133
134 Ok(())
135 }
136
137 fn load<Err>(&self, req: &WebRequest<Err>) -> (bool, HashMap<String, String>) {
138 if let Ok(cookies) = req.cookies() {
139 for cookie in cookies.iter() {
140 if cookie.name() == self.name {
141 let mut jar = CookieJar::new();
142 jar.add_original(cookie.clone());
143
144 let cookie_opt = match self.security {
145 CookieSecurity::Signed => jar.signed(&self.key).get(&self.name),
146 CookieSecurity::Private => jar.private(&self.key).get(&self.name),
147 };
148 if let Some(cookie) = cookie_opt {
149 if let Ok(val) = serde_json::from_str(cookie.value()) {
150 return (false, val);
151 }
152 }
153 }
154 }
155 }
156 (true, HashMap::new())
157 }
158}
159
160pub struct CookieSession(Rc<CookieSessionInner>);
201
202impl CookieSession {
203 pub fn signed(key: &[u8]) -> Self {
207 CookieSession(Rc::new(CookieSessionInner::new(key, CookieSecurity::Signed)))
208 }
209
210 pub fn private(key: &[u8]) -> Self {
214 CookieSession(Rc::new(CookieSessionInner::new(key, CookieSecurity::Private)))
215 }
216
217 pub fn path<S: Into<String>>(mut self, value: S) -> Self {
219 Rc::get_mut(&mut self.0).unwrap().path = value.into();
220 self
221 }
222
223 pub fn name<S: Into<String>>(mut self, value: S) -> Self {
225 Rc::get_mut(&mut self.0).unwrap().name = value.into();
226 self
227 }
228
229 pub fn domain<S: Into<String>>(mut self, value: S) -> Self {
231 Rc::get_mut(&mut self.0).unwrap().domain = Some(value.into());
232 self
233 }
234
235 pub fn secure(mut self, value: bool) -> Self {
240 Rc::get_mut(&mut self.0).unwrap().secure = value;
241 self
242 }
243
244 pub fn http_only(mut self, value: bool) -> Self {
246 Rc::get_mut(&mut self.0).unwrap().http_only = value;
247 self
248 }
249
250 pub fn same_site(mut self, value: SameSite) -> Self {
252 Rc::get_mut(&mut self.0).unwrap().same_site = Some(value);
253 self
254 }
255
256 pub fn max_age(self, seconds: i64) -> Self {
258 self.max_age_time(Duration::seconds(seconds))
259 }
260
261 pub fn max_age_time(mut self, value: time::Duration) -> Self {
263 Rc::get_mut(&mut self.0).unwrap().max_age = Some(value);
264 self
265 }
266
267 pub fn expires_in(self, seconds: i64) -> Self {
269 self.expires_in_time(Duration::seconds(seconds))
270 }
271
272 pub fn expires_in_time(mut self, value: Duration) -> Self {
274 Rc::get_mut(&mut self.0).unwrap().expires_in = Some(value);
275 self
276 }
277}
278
279impl<S, C> Middleware<S, C> for CookieSession {
280 type Service = CookieSessionMiddleware<S>;
281
282 fn create(&self, service: S, _: C) -> Self::Service {
283 CookieSessionMiddleware { service, inner: self.0.clone() }
284 }
285}
286
287pub struct CookieSessionMiddleware<S> {
289 service: S,
290 inner: Rc<CookieSessionInner>,
291}
292
293impl<S, Err> Service<WebRequest<Err>> for CookieSessionMiddleware<S>
294where
295 S: Service<WebRequest<Err>, Response = WebResponse>,
296 S::Error: 'static,
297 Err: ErrorRenderer,
298 Err::Container: From<CookieSessionError>,
299{
300 type Response = WebResponse;
301 type Error = S::Error;
302
303 ntex::forward_ready!(service);
304 ntex::forward_shutdown!(service);
305
306 async fn call(
312 &self,
313 req: WebRequest<Err>,
314 ctx: ServiceCtx<'_, Self>,
315 ) -> Result<Self::Response, Self::Error> {
316 let inner = self.inner.clone();
317 let (is_new, state) = self.inner.load(&req);
318 let prolong_expiration = self.inner.expires_in.is_some();
319 Session::set_session(state.into_iter(), &req);
320
321 ctx.call(&self.service, req).await.map(|mut res| {
322 match Session::get_changes(&mut res) {
323 (SessionStatus::Changed, Some(state))
324 | (SessionStatus::Renewed, Some(state)) => {
325 res.checked_expr::<Err, _, _>(|res| inner.set_cookie(res, state))
326 }
327 (SessionStatus::Unchanged, Some(state)) if prolong_expiration => {
328 res.checked_expr::<Err, _, _>(|res| inner.set_cookie(res, state))
329 }
330 (SessionStatus::Unchanged, _) =>
331 {
333 if is_new {
334 let state: HashMap<String, String> = HashMap::new();
335 res.checked_expr::<Err, _, _>(|res| {
336 inner.set_cookie(res, state.into_iter())
337 })
338 } else {
339 res
340 }
341 }
342 (SessionStatus::Purged, _) => {
343 let _ = inner.remove_cookie(&mut res);
344 res
345 }
346 _ => res,
347 }
348 })
349 }
350}
351
352#[cfg(test)]
353mod tests {
354 use super::*;
355 use ntex::web::{self, App, test};
356 use ntex::{time, util::Bytes};
357
358 #[ntex::test]
359 async fn cookie_session() {
360 let app = test::init_service(
361 App::new().wrap(CookieSession::signed(&[0; 32]).secure(false)).service(
362 web::resource("/").to(|ses: Session| async move {
363 let _ = ses.set("counter", 100);
364 "test"
365 }),
366 ),
367 )
368 .await;
369
370 let request = test::TestRequest::get().to_request();
371 let response = app.call(request).await.unwrap();
372 assert!(response.response().cookies().any(|c| c.name() == "ntex-session"));
373 }
374
375 #[ntex::test]
376 async fn private_cookie() {
377 let app = test::init_service(
378 App::new().wrap(CookieSession::private(&[0; 32]).secure(false)).service(
379 web::resource("/").to(|ses: Session| async move {
380 let _ = ses.set("counter", 100);
381 "test"
382 }),
383 ),
384 )
385 .await;
386
387 let request = test::TestRequest::get().to_request();
388 let response = app.call(request).await.unwrap();
389 assert!(response.response().cookies().any(|c| c.name() == "ntex-session"));
390 }
391
392 #[ntex::test]
393 async fn cookie_session_extractor() {
394 let app = test::init_service(
395 App::new().wrap(CookieSession::signed(&[0; 32]).secure(false)).service(
396 web::resource("/").to(|ses: Session| async move {
397 let _ = ses.set("counter", 100);
398 "test"
399 }),
400 ),
401 )
402 .await;
403
404 let request = test::TestRequest::get().to_request();
405 let response = app.call(request).await.unwrap();
406 assert!(response.response().cookies().any(|c| c.name() == "ntex-session"));
407 }
408
409 #[ntex::test]
410 async fn basics() {
411 let app = test::init_service(
412 App::new()
413 .wrap(
414 CookieSession::signed(&[0; 32])
415 .path("/test/")
416 .name("ntex-test")
417 .domain("localhost")
418 .http_only(true)
419 .same_site(SameSite::Lax)
420 .max_age(100),
421 )
422 .service(web::resource("/").to(|ses: Session| async move {
423 let _ = ses.set("counter", 100);
424 "test"
425 }))
426 .service(web::resource("/test/").to(|ses: Session| async move {
427 let val: usize = ses.get("counter").unwrap().unwrap();
428 format!("counter: {}", val)
429 })),
430 )
431 .await;
432
433 let request = test::TestRequest::get().to_request();
434 let response = app.call(request).await.unwrap();
435 let cookie = response
436 .response()
437 .cookies()
438 .find(|c| c.name() == "ntex-test")
439 .unwrap()
440 .into_owned();
441 assert_eq!(cookie.path().unwrap(), "/test/");
442
443 let request = test::TestRequest::with_uri("/test/").cookie(cookie).to_request();
444 let body = test::read_response(&app, request).await;
445 assert_eq!(body, Bytes::from_static(b"counter: 100"));
446 }
447
448 #[ntex::test]
449 async fn prolong_expiration() {
450 let app = test::init_service(
451 App::new()
452 .wrap(CookieSession::signed(&[0; 32]).secure(false).expires_in(60))
453 .service(web::resource("/").to(|ses: Session| async move {
454 let _ = ses.set("counter", 100);
455 "test"
456 }))
457 .service(web::resource("/test/").to(|| async move { "no-changes-in-session" })),
458 )
459 .await;
460
461 let request = test::TestRequest::get().to_request();
462 let response = app.call(request).await.unwrap();
463 let expires_1 = response
464 .response()
465 .cookies()
466 .find(|c| c.name() == "ntex-session")
467 .expect("Cookie is set")
468 .expires()
469 .expect("Expiration is set");
470
471 time::sleep(time::Seconds::ONE).await;
472
473 let request = test::TestRequest::with_uri("/test/").to_request();
474 let response = app.call(request).await.unwrap();
475 let expires_2 = response
476 .response()
477 .cookies()
478 .find(|c| c.name() == "ntex-session")
479 .expect("Cookie is set")
480 .expires()
481 .expect("Expiration is set");
482
483 assert!(
484 expires_2.datetime().unwrap() - expires_1.datetime().unwrap()
485 >= Duration::seconds(1)
486 );
487 }
488}