1
0
mirror of https://github.com/actix/actix-extras.git synced 2025-02-23 02:43:16 +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 //! ```no_run
//! use std::{sync::Arc, time::Duration}; //! use std::{time::Duration};
//! use actix_web::{dev::ServiceRequest, get, web, App, HttpServer, Responder}; //! use actix_web::{dev::ServiceRequest, get, web, App, HttpServer, Responder};
//! use actix_session::SessionExt as _; //! use actix_session::SessionExt as _;
//! use actix_limitation::{Limiter, RateLimiter}; //! use actix_limitation::{Limiter, RateLimiter};
@ -23,8 +23,8 @@
//! Limiter::builder("redis://127.0.0.1") //! Limiter::builder("redis://127.0.0.1")
//! .key_by(|req: &ServiceRequest| { //! .key_by(|req: &ServiceRequest| {
//! req.get_session() //! req.get_session()
//! .get(&"session-id") //! .get("session-id")
//! .unwrap_or_else(|_| req.cookie(&"rate-api-id").map(|c| c.to_string())) //! .unwrap_or_else(|_| req.cookie("rate-api-id").map(|c| c.to_string()))
//! }) //! })
//! .limit(5000) //! .limit(5000)
//! .period(Duration::from_secs(3600)) // 60 minutes //! .period(Duration::from_secs(3600)) // 60 minutes
@ -176,4 +176,41 @@ mod tests {
assert_eq!(limiter.limit, 5000); assert_eq!(limiter.limit, 5000);
assert_eq!(limiter.period, Duration::from_secs(3600)); 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_utils::future::{ok, Ready};
use actix_web::{ use actix_web::{
@ -11,9 +11,16 @@ use actix_web::{
use crate::{Error as LimitationError, Limiter}; use crate::{Error as LimitationError, Limiter};
/// Rate limit middleware. /// Rate limit middleware.
///
/// Use the `scope` variable to define multiple limiter
#[derive(Debug, Default)] #[derive(Debug, Default)]
#[non_exhaustive] #[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 impl<S, B> Transform<S, ServiceRequest> for RateLimiter
where where
@ -30,6 +37,7 @@ where
fn new_transform(&self, service: S) -> Self::Future { fn new_transform(&self, service: S) -> Self::Future {
ok(RateLimiterMiddleware { ok(RateLimiterMiddleware {
service: Rc::new(service), service: Rc::new(service),
scope: self.scope,
}) })
} }
} }
@ -38,6 +46,7 @@ where
#[derive(Debug)] #[derive(Debug)]
pub struct RateLimiterMiddleware<S> { pub struct RateLimiterMiddleware<S> {
service: Rc<S>, service: Rc<S>,
scope: Option<&'static str>,
} }
impl<S, B> Service<ServiceRequest> for RateLimiterMiddleware<S> impl<S, B> Service<ServiceRequest> for RateLimiterMiddleware<S>
@ -55,10 +64,21 @@ where
fn call(&self, req: ServiceRequest) -> Self::Future { fn call(&self, req: ServiceRequest) -> Self::Future {
// A mis-configuration of the Actix App will result in a **runtime** failure, so the expect // A mis-configuration of the Actix App will result in a **runtime** failure, so the expect
// method description is important context for the developer. // method description is important context for the developer.
let limiter = req 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>>() .app_data::<web::Data<Limiter>>()
.expect("web::Data<Limiter> should be set in app data for RateLimiter middleware") .expect("web::Data<Limiter> should be set in app data for RateLimiter middleware");
.clone(); // Deref to get the Limiter
(***a).clone()
};
let key = (limiter.get_key_fn)(&req); let key = (limiter.get_key_fn)(&req);
let service = Rc::clone(&self.service); let service = Rc::clone(&self.service);