1
0
mirror of https://github.com/actix/actix-extras.git synced 2025-02-22 10:23:18 +01:00

feat: optional scopes for middleware

This commit is contained in:
GreeFine 2022-10-03 18:26:08 +02:00
parent da0a806e8d
commit 368ada2b75
2 changed files with 66 additions and 9 deletions

View File

@ -7,7 +7,7 @@
//! ```
//!
//! ```no_run
//! use std::{sync::Arc, time::Duration};
//! use std::{time::Duration};
//! use actix_web::{dev::ServiceRequest, get, web, App, HttpServer, Responder};
//! use actix_session::SessionExt as _;
//! use actix_limitation::{Limiter, RateLimiter};
@ -23,8 +23,8 @@
//! Limiter::builder("redis://127.0.0.1")
//! .key_by(|req: &ServiceRequest| {
//! req.get_session()
//! .get(&"session-id")
//! .unwrap_or_else(|_| req.cookie(&"rate-api-id").map(|c| c.to_string()))
//! .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
@ -176,4 +176,41 @@ mod tests {
assert_eq!(limiter.limit, 5000);
assert_eq!(limiter.period, Duration::from_secs(3600));
}
#[actix_web::test]
async fn test_create_scoped_limiter() {
todo!("finish tests")
// use actix_session::SessionExt as _;
// use actix_web::{dev::ServiceRequest, get, web, App, HttpServer, Responder};
// use std::time::Duration;
// #[get("/{id}/{name}")]
// async fn index(info: web::Path<(u32, String)>) -> impl Responder {
// format!("Hello {}! id:{}", info.1, info.0)
// }
// let limiter = web::Data::new(
// Limiter::builder("redis://127.0.0.1")
// .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()
// .unwrap(),
// );
// HttpServer::new(move || {
// App::new()
// .wrap(RateLimiter::default())
// .app_data(limiter.clone())
// .service(index)
// })
// .bind(("127.0.0.1", 8080))
// .expect("test")
// .run()
// .await;
}
}

View File

@ -1,4 +1,4 @@
use std::{future::Future, pin::Pin, rc::Rc};
use std::{collections::HashMap, future::Future, pin::Pin, rc::Rc};
use actix_utils::future::{ok, Ready};
use actix_web::{
@ -11,9 +11,16 @@ use actix_web::{
use crate::{Error as LimitationError, Limiter};
/// Rate limit middleware.
///
/// Use the `scope` variable to define multiple limiter
#[derive(Debug, Default)]
#[non_exhaustive]
pub struct RateLimiter;
pub struct RateLimiter {
/// Used to define multiple limiter, with different configurations
///
/// WARNING: When used (not None) the middleware will expect a `HashMap<Limiter>` in the actix-web `app_data`
pub scope: Option<&'static str>,
}
impl<S, B> Transform<S, ServiceRequest> for RateLimiter
where
@ -30,6 +37,7 @@ where
fn new_transform(&self, service: S) -> Self::Future {
ok(RateLimiterMiddleware {
service: Rc::new(service),
scope: self.scope,
})
}
}
@ -38,6 +46,7 @@ where
#[derive(Debug)]
pub struct RateLimiterMiddleware<S> {
service: Rc<S>,
scope: Option<&'static str>,
}
impl<S, B> Service<ServiceRequest> for RateLimiterMiddleware<S>
@ -55,10 +64,21 @@ where
fn call(&self, req: ServiceRequest) -> Self::Future {
// A mis-configuration of the Actix App will result in a **runtime** failure, so the expect
// method description is important context for the developer.
let limiter = req
.app_data::<web::Data<Limiter>>()
.expect("web::Data<Limiter> should be set in app data for RateLimiter middleware")
.clone();
let limiter = if let Some(scope) = self.scope {
let limiters = req.app_data::<web::Data<HashMap<&str, Limiter>>>().expect(
"web::Data<HashMap<Limiter>> should be set in app data for RateLimiter middleware",
);
limiters
.get(scope)
.unwrap_or_else(|| panic!("Unable to find defined limiter with scope: {}", scope))
.clone()
} else {
let a = req
.app_data::<web::Data<Limiter>>()
.expect("web::Data<Limiter> should be set in app data for RateLimiter middleware");
// Deref to get the Limiter
(***a).clone()
};
let key = (limiter.get_key_fn)(&req);
let service = Rc::clone(&self.service);