Limitation: custom key from closure (#281)

Co-authored-by: Rob Ede <robjtede@icloud.com>

actix-limitation/CHANGES.md

@@ -1,6 +1,7 @@
 # Changes
 
 ## Unreleased - 2022-xx-xx
+- Add `Builder::key_by` for setting a custom rate limit key function.
 - Implement `Default` for `RateLimiter`.
 - `RateLimiter` is marked `#[non_exhaustive]`; use `RateLimiter::default()` instead.
 - In the middleware errors from the count function are matched and respond with `INTERNAL_SERVER_ERROR` if it's an unexpected error, instead of the default `TOO_MANY_REQUESTS`.

actix-limitation/Cargo.toml

@@ -12,10 +12,13 @@ repository = "https://github.com/actix/actix-extras.git"
 license = "MIT OR Apache-2.0"
 edition = "2018"
 
+[features]
+default = ["session"]
+session = ["actix-session"]
+
 [dependencies]
-actix-session = "0.7"
 actix-utils = "3"
-actix-web = { version = "4", default-features = false }
+actix-web = { version = "4", features = ["cookies"] }
 
 chrono = "0.4"
 derive_more = "0.99.5"
@@ -23,6 +26,9 @@ log = "0.4"
 redis = { version = "0.21", default-features = false, features = ["tokio-comp"] }
 time = "0.3"
 
+# session
+actix-session = { version = "0.7", optional = true }
+
 [dev-dependencies]
 actix-web = "4"
 static_assertions = "1"

actix-limitation/README.md

@@ -17,9 +17,10 @@ actix-limitation = "0.3"
 ```
 
 ```rust
-use std::time::Duration;
-use actix_web::{get, web, App, HttpServer, Responder};
 use actix_limitation::{Limiter, RateLimiter};
+use actix_session::SessionExt as _;
+use actix_web::{dev::ServiceRequest, get, web, App, HttpServer, Responder};
+use std::{sync::Arc, time::Duration};
 
 #[get("/{id}/{name}")]
 async fn index(info: web::Path<(u32, String)>) -> impl Responder {
@@ -29,22 +30,24 @@ async fn index(info: web::Path<(u32, String)>) -> impl Responder {
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
     let limiter = web::Data::new(
-        Limiter::build("redis://127.0.0.1")
+        Limiter::builder("redis://127.0.0.1")
-            .cookie_name("session-id".to_owned())
+            .key_by(|req: &ServiceRequest| {
-            .session_key("rate-api-id".to_owned())
+                req.get_session()
+                    .get(&"session-id")
+                    .unwrap_or_else(|_| req.cookie(&"rate-api-id").map(|c| c.to_string()))
+            })
             .limit(5000)
             .period(Duration::from_secs(3600)) // 60 minutes
-            .finish()
+            .build()
-            .expect("Can't build actix-limiter"),
+            .unwrap(),
     );
-
     HttpServer::new(move || {
         App::new()
-            .wrap(RateLimiter)
+            .wrap(RateLimiter::default())
             .app_data(limiter.clone())
             .service(index)
     })
-    .bind("127.0.0.1:8080")?
+    .bind(("127.0.0.1", 8080))?
     .run()
     .await
 }

actix-limitation/src/builder.rs

@@ -1,8 +1,11 @@
-use std::{borrow::Cow, time::Duration};
+use std::{borrow::Cow, sync::Arc, time::Duration};
 
+#[cfg(feature = "session")]
+use actix_session::SessionExt as _;
+use actix_web::dev::ServiceRequest;
 use redis::Client;
 
-use crate::{errors::Error, Limiter};
+use crate::{errors::Error, GetArcBoxKeyFn, Limiter};
 
 /// Rate limiter builder.
 #[derive(Debug)]
@@ -10,7 +13,9 @@ pub struct Builder {
     pub(crate) redis_url: String,
     pub(crate) limit: usize,
     pub(crate) period: Duration,
+    pub(crate) get_key_fn: Option<GetArcBoxKeyFn>,
     pub(crate) cookie_name: Cow<'static, str>,
+    #[cfg(feature = "session")]
     pub(crate) session_key: Cow<'static, str>,
 }
 
@@ -27,14 +32,38 @@ impl Builder {
         self
     }
 
-    /// Set name of cookie to be sent.
+    /// Sets rate limit key derivation function.
+    ///
+    /// Should not be used in combination with `cookie_name` or `session_key` as they conflict.
+    pub fn key_by<F>(&mut self, resolver: F) -> &mut Self
+    where
+        F: Fn(&ServiceRequest) -> Option<String> + Send + Sync + 'static,
+    {
+        self.get_key_fn = Some(Arc::new(resolver));
+        self
+    }
+
+    /// Sets name of cookie to be sent.
+    ///
+    /// This method should not be used in combination of `key_by` as they conflict.
+    #[deprecated = "Prefer `key_by`."]
     pub fn cookie_name(&mut self, cookie_name: impl Into<Cow<'static, str>>) -> &mut Self {
+        if self.get_key_fn.is_some() {
+            panic!("This method should not be used in combination of get_key as they overwrite each other")
+        }
         self.cookie_name = cookie_name.into();
         self
     }
 
-    /// Set session key to be used in backend.
+    /// Sets session key to be used in backend.
+    ///
+    /// This method should not be used in combination of `key_by` as they conflict.
+    #[deprecated = "Prefer `key_by`."]
+    #[cfg(feature = "session")]
     pub fn session_key(&mut self, session_key: impl Into<Cow<'static, str>>) -> &mut Self {
+        if self.get_key_fn.is_some() {
+            panic!("This method should not be used in combination of get_key as they overwrite each other")
+        }
         self.session_key = session_key.into();
         self
     }
@@ -43,13 +72,35 @@ impl Builder {
     ///
     /// Note that this method will connect to the Redis server to test its connection which is a
     /// **synchronous** operation.
-    pub fn build(&self) -> Result<Limiter, Error> {
+    pub fn build(&mut self) -> Result<Limiter, Error> {
+        let get_key = if let Some(resolver) = self.get_key_fn.clone() {
+            resolver
+        } else {
+            let cookie_name = self.cookie_name.clone();
+
+            #[cfg(feature = "session")]
+            let session_key = self.session_key.clone();
+
+            let closure: GetArcBoxKeyFn = Arc::new(Box::new(move |req: &ServiceRequest| {
+                #[cfg(feature = "session")]
+                let res = req
+                    .get_session()
+                    .get(&session_key)
+                    .unwrap_or_else(|_| req.cookie(&cookie_name).map(|c| c.to_string()));
+
+                #[cfg(not(feature = "session"))]
+                let res = req.cookie(&cookie_name).map(|c| c.to_string());
+
+                res
+            }));
+            closure
+        };
+
         Ok(Limiter {
             client: Client::open(self.redis_url.as_str())?,
             limit: self.limit,
             period: self.period,
-            cookie_name: self.cookie_name.clone(),
+            get_key_fn: get_key,
-            session_key: self.session_key.clone(),
         })
     }
 }
@@ -66,13 +117,16 @@ mod tests {
             redis_url: redis_url.to_owned(),
             limit: 100,
             period,
+            get_key_fn: Some(Arc::new(|_| None)),
             cookie_name: Cow::Owned("session".to_string()),
+            #[cfg(feature = "session")]
             session_key: Cow::Owned("rate-api".to_string()),
         };
 
         assert_eq!(builder.redis_url, redis_url);
         assert_eq!(builder.limit, 100);
         assert_eq!(builder.period, period);
+        #[cfg(feature = "session")]
         assert_eq!(builder.session_key, "rate-api");
         assert_eq!(builder.cookie_name, "session");
     }
@@ -85,22 +139,16 @@ mod tests {
             redis_url: redis_url.to_owned(),
             limit: 100,
             period: Duration::from_secs(10),
-            session_key: Cow::Borrowed("key"),
+            get_key_fn: Some(Arc::new(|_| None)),
             cookie_name: Cow::Borrowed("sid"),
+            #[cfg(feature = "session")]
+            session_key: Cow::Borrowed("key"),
         };
 
-        let limiter = builder
+        let limiter = builder.limit(200).period(period).build().unwrap();
-            .limit(200)
-            .period(period)
-            .cookie_name("session".to_string())
-            .session_key("rate-api".to_string())
-            .build()
-            .unwrap();
 
         assert_eq!(limiter.limit, 200);
         assert_eq!(limiter.period, period);
-        assert_eq!(limiter.session_key, "rate-api");
-        assert_eq!(limiter.cookie_name, "session");
     }
 
     #[test]
@@ -112,8 +160,10 @@ mod tests {
             redis_url: redis_url.to_owned(),
             limit: 100,
             period: Duration::from_secs(10),
-            session_key: Cow::Borrowed("key"),
+            get_key_fn: Some(Arc::new(|_| None)),
             cookie_name: Cow::Borrowed("sid"),
+            #[cfg(feature = "session")]
+            session_key: Cow::Borrowed("key"),
         };
 
         builder.limit(200).period(period).build().unwrap();

actix-limitation/src/lib.rs

@@ -7,8 +7,9 @@
 //! ```
 //!
 //! ```no_run
-//! use std::time::Duration;
+//! use std::{sync::Arc, time::Duration};
-//! use actix_web::{get, web, App, HttpServer, Responder};
+//! use actix_web::{dev::ServiceRequest, get, web, App, HttpServer, Responder};
+//! use actix_session::SessionExt as _;
 //! use actix_limitation::{Limiter, RateLimiter};
 //!
 //! #[get("/{id}/{name}")]
@@ -20,8 +21,11 @@
 //! async fn main() -> std::io::Result<()> {
 //!     let limiter = web::Data::new(
 //!         Limiter::builder("redis://127.0.0.1")
-//!             .cookie_name("session-id".to_owned())
+//!             .key_by(|req: &ServiceRequest| {
-//!             .session_key("rate-api-id".to_owned())
+//!                 req.get_session()
+//!                     .get(&"session-id")
+//!                     .unwrap_or_else(|_| req.cookie(&"rate-api-id").map(|c| c.to_string()))
+//!             })
 //!             .limit(5000)
 //!             .period(Duration::from_secs(3600)) // 60 minutes
 //!             .build()
@@ -46,8 +50,9 @@
 #![doc(html_logo_url = "https://actix.rs/img/logo.png")]
 #![doc(html_favicon_url = "https://actix.rs/favicon.ico")]
 
-use std::{borrow::Cow, time::Duration};
+use std::{borrow::Cow, fmt, sync::Arc, time::Duration};
 
+use actix_web::dev::ServiceRequest;
 use redis::Client;
 
 mod builder;
@@ -70,16 +75,34 @@ pub const DEFAULT_PERIOD_SECS: u64 = 3600;
 pub const DEFAULT_COOKIE_NAME: &str = "sid";
 
 /// Default session key.
+#[cfg(feature = "session")]
 pub const DEFAULT_SESSION_KEY: &str = "rate-api-id";
 
+/// Helper trait to impl Debug on GetKeyFn type
+trait GetKeyFnT: Fn(&ServiceRequest) -> Option<String> {}
+
+impl<T> GetKeyFnT for T where T: Fn(&ServiceRequest) -> Option<String> {}
+
+/// Get key function type with auto traits
+type GetKeyFn = dyn GetKeyFnT + Send + Sync;
+
+/// Get key resolver function type
+impl fmt::Debug for GetKeyFn {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "GetKeyFn")
+    }
+}
+
+/// Wrapped Get key function Trait
+type GetArcBoxKeyFn = Arc<GetKeyFn>;
+
 /// Rate limiter.
 #[derive(Debug, Clone)]
 pub struct Limiter {
     client: Client,
     limit: usize,
     period: Duration,
-    cookie_name: Cow<'static, str>,
+    get_key_fn: GetArcBoxKeyFn,
-    session_key: Cow<'static, str>,
 }
 
 impl Limiter {
@@ -93,7 +116,9 @@ impl Limiter {
             redis_url: redis_url.into(),
             limit: DEFAULT_REQUEST_LIMIT,
             period: Duration::from_secs(DEFAULT_PERIOD_SECS),
+            get_key_fn: None,
             cookie_name: Cow::Borrowed(DEFAULT_COOKIE_NAME),
+            #[cfg(feature = "session")]
             session_key: Cow::Borrowed(DEFAULT_SESSION_KEY),
         }
     }
@@ -146,14 +171,12 @@ mod tests {
 
     #[test]
     fn test_create_limiter() {
-        let builder = Limiter::builder("redis://127.0.0.1:6379/1");
+        let mut builder = Limiter::builder("redis://127.0.0.1:6379/1");
         let limiter = builder.build();
         assert!(limiter.is_ok());
 
         let limiter = limiter.unwrap();
         assert_eq!(limiter.limit, 5000);
         assert_eq!(limiter.period, Duration::from_secs(3600));
-        assert_eq!(limiter.cookie_name, DEFAULT_COOKIE_NAME);
-        assert_eq!(limiter.session_key, DEFAULT_SESSION_KEY);
     }
 }

actix-limitation/src/middleware.rs

@@ -1,6 +1,5 @@
 use std::{future::Future, pin::Pin, rc::Rc};
 
-use actix_session::SessionExt as _;
 use actix_utils::future::{ok, Ready};
 use actix_web::{
     body::EitherBody,
@@ -61,25 +60,18 @@ where
             .expect("web::Data<Limiter> should be set in app data for RateLimiter middleware")
             .clone();
 
-        let key = req.get_session().get(&limiter.session_key).unwrap_or(None);
+        let key = (limiter.get_key_fn)(&req);
         let service = Rc::clone(&self.service);
 
         let key = match key {
             Some(key) => key,
             None => {
-                let fallback = req.cookie(&limiter.cookie_name).map(|c| c.to_string());
+                return Box::pin(async move {
-
+                    service
-                match fallback {
+                        .call(req)
-                    Some(key) => key,
+                        .await
-                    None => {
+                        .map(ServiceResponse::map_into_left_body)
-                        return Box::pin(async move {
+                });
-                            service
-                                .call(req)
-                                .await
-                                .map(ServiceResponse::map_into_left_body)
-                        });
-                    }
-                }
             }
         };
 

actix-limitation/tests/tests.rs

@@ -1,9 +1,12 @@
-use actix_limitation::{Error, Limiter};
+use std::time::Duration;
+
+use actix_limitation::{Error, Limiter, RateLimiter};
+use actix_web::{dev::ServiceRequest, http::StatusCode, test, web, App, HttpRequest, HttpResponse};
 use uuid::Uuid;
 
 #[test]
 #[should_panic = "Redis URL did not parse"]
-fn test_create_limiter_error() {
+async fn test_create_limiter_error() {
     Limiter::builder("127.0.0.1").build().unwrap();
 }
 
@@ -51,3 +54,39 @@ async fn test_limiter_count_error() -> Result<(), Error> {
 
     Ok(())
 }
+
+#[actix_web::test]
+async fn test_limiter_key_by() -> Result<(), Error> {
+    let cooldown_period = Duration::from_secs(1);
+    let limiter = Limiter::builder("redis://127.0.0.1:6379/3")
+        .limit(2)
+        .period(cooldown_period)
+        .key_by(|_: &ServiceRequest| Some("fix_key".to_string()))
+        .build()
+        .unwrap();
+
+    let app = test::init_service(
+        App::new()
+            .wrap(RateLimiter::default())
+            .app_data(web::Data::new(limiter))
+            .route(
+                "/",
+                web::get().to(|_: HttpRequest| async { HttpResponse::Ok().body("ok") }),
+            ),
+    )
+    .await;
+    for _ in 1..2 {
+        for index in 1..4 {
+            let req = test::TestRequest::default().to_request();
+            let resp = test::call_service(&app, req).await;
+            if index <= 2 {
+                assert!(resp.status().is_success());
+            } else {
+                assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
+            }
+        }
+        std::thread::sleep(cooldown_period);
+    }
+
+    Ok(())
+}