From 9e6db09c1241c55eac4f60b679813bcb6a059037 Mon Sep 17 00:00:00 2001
From: Tglman <tglman@tglman.com>
Date: Tue, 22 Mar 2022 11:45:26 +0000
Subject: [PATCH] add support for session access from guards

---
 actix-session/CHANGES.md         |  1 +
 actix-session/src/lib.rs         | 42 ++++++++++++++++++++++++++++----
 actix-session/src/session_ext.rs |  7 ++++++
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/actix-session/CHANGES.md b/actix-session/CHANGES.md
index aaafad6506..5a380f02b9 100644
--- a/actix-session/CHANGES.md
+++ b/actix-session/CHANGES.md
@@ -2,6 +2,7 @@
 
 ## Unreleased - 2021-xx-xx
 
+- Implement `SessionExt` for `GuardContext`.
 
 ## 0.6.1 - 2022-03-21
 - No significant changes since `0.6.0`.
diff --git a/actix-session/src/lib.rs b/actix-session/src/lib.rs
index a10e65ecb7..c4207bc8d4 100644
--- a/actix-session/src/lib.rs
+++ b/actix-session/src/lib.rs
@@ -199,7 +199,7 @@ pub mod test_helpers {
     mod acceptance_tests {
         use actix_web::{
             dev::Service,
-            middleware, test,
+            guard, middleware, test,
             web::{self, get, post, resource, Bytes},
             App, HttpResponse, Result,
         };
@@ -209,7 +209,7 @@ pub mod test_helpers {
 
         use crate::{
             middleware::SessionLength, storage::SessionStore, test_helpers::key,
-            CookieContentSecurity, Session, SessionMiddleware,
+            CookieContentSecurity, Session, SessionExt, SessionMiddleware,
         };
 
         pub(super) async fn basic_workflow<F, Store>(
@@ -332,6 +332,13 @@ pub mod test_helpers {
                     .service(resource("/do_something").route(post().to(do_something)))
                     .service(resource("/login").route(post().to(login)))
                     .service(resource("/logout").route(post().to(logout)))
+                    .service(
+                        web::scope("/protected")
+                            .guard(guard::fn_guard(|g| {
+                                g.get_session().get::<String>("user_id").unwrap().is_some()
+                            }))
+                            .service(resource("/").route(get().to(count))),
+                    )
             });
 
             // Step 1:  GET index
@@ -462,8 +469,21 @@ pub mod test_helpers {
                     counter: 3
                 }
             );
+            // Step 8: POST again to do_something, including session cookie #3 in request
+            //   - updates session state: {"counter": 3, "user_id": "ferris"}
+            //   - response should be: {"counter": 3, "user_id": "ferris"}
+            let req_7 = srv.post("/protected/count").cookie(cookie_3.clone()).send();
+            let mut resp_7 = req_7.await.unwrap();
+            let result_7 = resp_7.json::<IndexResponse>().await.unwrap();
+            assert_eq!(
+                result_7,
+                IndexResponse {
+                    user_id: Some("ferris".into()),
+                    counter: 3
+                }
+            );
 
-            // Step 8: GET index, including session cookie #2 in request
+            // Step 9: GET index, including session cookie #2 in request
             // If invalidation is supported, no state will be found associated to this session.
             // If invalidation is not supported, the old state will still be retrieved.
             let req_8 = srv.get("/").cookie(cookie_2.clone()).send();
@@ -489,7 +509,7 @@ pub mod test_helpers {
                 );
             }
 
-            // Step 9: POST to logout, including session cookie #3
+            // Step 10: POST to logout, including session cookie #3
             //   - set-cookie actix-session will be in response with session cookie #3
             //     invalidation logic
             let req_9 = srv.post("/logout").cookie(cookie_3.clone()).send();
@@ -504,7 +524,13 @@ pub mod test_helpers {
             assert_eq!(0, cookie_3.max_age().map(|t| t.whole_seconds()).unwrap());
             assert_eq!("/", cookie_3.path().unwrap());
 
-            // Step 10: GET index, including session cookie #3 in request
+            // Step 11: GET protected count without session cookie
+            //   - response should be not a success
+            let req_9 = srv.get("/protected/count").send();
+            let resp_9 = req_9.await.unwrap();
+            assert!(!resp_9.status().is_success());
+
+            // Step 12: GET index, including session cookie #3 in request
             //   - set-cookie actix-session should NOT be in response if invalidation is supported
             //   - response should be: {"counter": 0, "user_id": None}
             let req_10 = srv.get("/").cookie(cookie_3.clone()).send();
@@ -548,6 +574,12 @@ pub mod test_helpers {
 
             Ok(HttpResponse::Ok().json(&IndexResponse { user_id, counter }))
         }
+        async fn count(session: Session) -> Result<HttpResponse> {
+            let user_id: Option<String> = session.get::<String>("user_id").unwrap();
+            let counter: i32 = session.get::<i32>("counter").unwrap().unwrap();
+
+            Ok(HttpResponse::Ok().json(&IndexResponse { user_id, counter }))
+        }
 
         #[derive(Deserialize)]
         struct Identity {
diff --git a/actix-session/src/session_ext.rs b/actix-session/src/session_ext.rs
index 00d799c115..1cff2b4725 100644
--- a/actix-session/src/session_ext.rs
+++ b/actix-session/src/session_ext.rs
@@ -1,5 +1,6 @@
 use actix_web::{
     dev::{ServiceRequest, ServiceResponse},
+    guard::GuardContext,
     HttpMessage, HttpRequest,
 };
 
@@ -29,3 +30,9 @@ impl SessionExt for ServiceResponse {
         self.request().get_session()
     }
 }
+
+impl<'a> SessionExt for GuardContext<'a> {
+    fn get_session(&self) -> Session {
+        Session::get_session(&mut *self.req_data_mut())
+    }
+}