From 6e79465362605fa6bdbba2551f09a25bd7a7e2fe Mon Sep 17 00:00:00 2001
From: Rob Ede <robjtede@icloud.com>
Date: Sun, 31 Jul 2022 03:03:43 +0100
Subject: [PATCH] make RateLimiter non-exhaustive

---
 actix-limitation/CHANGES.md        | 16 +++++-----
 actix-limitation/src/lib.rs        |  2 +-
 actix-limitation/src/middleware.rs | 49 +++++++++++-------------------
 actix-session/tests/session.rs     |  2 ++
 4 files changed, 28 insertions(+), 41 deletions(-)

diff --git a/actix-limitation/CHANGES.md b/actix-limitation/CHANGES.md
index 83287babd..819c496cc 100644
--- a/actix-limitation/CHANGES.md
+++ b/actix-limitation/CHANGES.md
@@ -1,7 +1,8 @@
 # Changes
 
 ## Unreleased - 2022-xx-xx
-
+- Implement `Default` for `RateLimiter`.
+- `RateLimiter` can no longer be constructed without `::default()`.
 
 ## 0.3.0 - 2022-07-11
 - `Limiter::builder` now takes an `impl Into<String>`.
@@ -10,14 +11,11 @@
 
 
 ## 0.2.0 - 2022-03-22
-- Update Actix Web dependency to v4 ecosystem. [#229]
-- Update Tokio dependencies to v1 ecosystem. [#229]
-- Rename `Limiter::{build => builder}()`. [#232]
-- Rename `Builder::{finish => build}()`. [#232]
-- Exceeding the rate limit now returns a 429 Too Many Requests response. [#232]
-
-[#229]: https://github.com/actix/actix-extras/pull/229
-[#232]: https://github.com/actix/actix-extras/pull/232
+- Update Actix Web dependency to v4 ecosystem.
+- Update Tokio dependencies to v1 ecosystem.
+- Rename `Limiter::{build => builder}()`.
+- Rename `Builder::{finish => build}()`.
+- Exceeding the rate limit now returns a 429 Too Many Requests response.
 
 
 ## 0.1.4 - 2022-03-18
diff --git a/actix-limitation/src/lib.rs b/actix-limitation/src/lib.rs
index 7f160a309..8124fe632 100644
--- a/actix-limitation/src/lib.rs
+++ b/actix-limitation/src/lib.rs
@@ -30,7 +30,7 @@
 //!
 //!     HttpServer::new(move || {
 //!         App::new()
-//!             .wrap(RateLimiter)
+//!             .wrap(RateLimiter::default())
 //!             .app_data(limiter.clone())
 //!             .service(index)
 //!     })
diff --git a/actix-limitation/src/middleware.rs b/actix-limitation/src/middleware.rs
index 2acdb776a..ab1453ee3 100644
--- a/actix-limitation/src/middleware.rs
+++ b/actix-limitation/src/middleware.rs
@@ -4,16 +4,16 @@ use actix_session::SessionExt as _;
 use actix_utils::future::{ok, Ready};
 use actix_web::{
     body::EitherBody,
-    cookie::Cookie,
     dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
-    http::{header::COOKIE, StatusCode},
+    http::StatusCode,
     web, Error, HttpResponse,
 };
 
 use crate::Limiter;
 
 /// Rate limit middleware.
-#[derive(Debug)]
+#[derive(Debug, Default)]
+#[non_exhaustive]
 pub struct RateLimiter;
 
 impl<S, B> Transform<S, ServiceRequest> for RateLimiter
@@ -61,23 +61,26 @@ where
             .expect("web::Data<Limiter> should be set in app data for RateLimiter middleware")
             .clone();
 
-        let (key, fallback) = key(&req, limiter.clone());
-
+        let key = req.get_session().get(&limiter.session_key).unwrap_or(None);
         let service = Rc::clone(&self.service);
 
         let key = match key {
             Some(key) => key,
-            None => match fallback {
-                Some(key) => key,
-                None => {
-                    return Box::pin(async move {
-                        service
-                            .call(req)
-                            .await
-                            .map(ServiceResponse::map_into_left_body)
-                    });
+            None => {
+                let fallback = req.cookie(&limiter.cookie_name).map(|c| c.to_string());
+
+                match fallback {
+                    Some(key) => key,
+                    None => {
+                        return Box::pin(async move {
+                            service
+                                .call(req)
+                                .await
+                                .map(ServiceResponse::map_into_left_body)
+                        });
+                    }
                 }
-            },
+            }
         };
 
         Box::pin(async move {
@@ -98,19 +101,3 @@ where
         })
     }
 }
-
-fn key(req: &ServiceRequest, limiter: web::Data<Limiter>) -> (Option<String>, Option<String>) {
-    let session = req.get_session();
-    let result: Option<String> = session.get(&limiter.session_key).unwrap_or(None);
-    let cookies = req.headers().get_all(COOKIE);
-    let cookie = cookies
-        .filter_map(|i| i.to_str().ok())
-        .find(|i| i.contains(limiter.cookie_name.as_ref()));
-
-    let fallback = match cookie {
-        Some(value) => Cookie::parse(value).ok().map(|i| i.to_string()),
-        None => None,
-    };
-
-    (result, fallback)
-}
diff --git a/actix-session/tests/session.rs b/actix-session/tests/session.rs
index 32d24a4b6..53319e653 100644
--- a/actix-session/tests/session.rs
+++ b/actix-session/tests/session.rs
@@ -68,6 +68,7 @@ async fn session_entries() {
     map.contains_key("test_str");
     map.contains_key("test_num");
 }
+
 #[actix_web::test]
 async fn insert_session_after_renew() {
     let session = test::TestRequest::default().to_srv_request().get_session();
@@ -81,6 +82,7 @@ async fn insert_session_after_renew() {
     session.insert("test_val1", "val1").unwrap();
     assert_eq!(session.status(), SessionStatus::Renewed);
 }
+
 #[actix_web::test]
 async fn remove_session_after_renew() {
     let session = test::TestRequest::default().to_srv_request().get_session();