actix/actix-extras
1
0
Fork 0
mirror of https://github.com/actix/actix-extras.git synced 2025-06-26 18:37:41 +02:00
Code Issues Releases Wiki Activity

Limitation: custom key from closure (#281)

Browse Source 
Co-authored-by: Rob Ede <robjtede@icloud.com>
This commit is contained in:
Raphael C
2022-09-11 01:02:54 +02:00
committed by GitHub
parent a623c50e9c
commit 32313c0af6
7 changed files with 171 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

86
actix-limitation/src/builder.rs
View File

@ -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(),
session_key: self.session_key.clone(),
get_key_fn: get_key,
})
}
}
@ -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
.limit(200)
.period(period)
.cookie_name("session".to_string())
.session_key("rate-api".to_string())
.build()
.unwrap();
let limiter = builder.limit(200).period(period).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();

43
actix-limitation/src/lib.rs
View File

@ -7,8 +7,9 @@
//! ```
//!
//! ```no_run
//! use std::time::Duration;
//! use actix_web::{get, web, App, HttpServer, Responder};
//! use std::{sync::Arc, time::Duration};
//! 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())
//! .session_key("rate-api-id".to_owned())
//! .key_by(|req: &ServiceRequest| {
//! 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>,
session_key: Cow<'static, str>,
get_key_fn: GetArcBoxKeyFn,
}
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);
}
}

22
actix-limitation/src/middleware.rs
View File

@ -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());
match fallback {
Some(key) => key,
None => {
return Box::pin(async move {
service
.call(req)
.await
.map(ServiceResponse::map_into_left_body)
});
}
}
return Box::pin(async move {
service
.call(req)
.await
.map(ServiceResponse::map_into_left_body)
});
}
};