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

make RateLimiter non-exhaustive

Browse Source
This commit is contained in:
Rob Ede
2022-07-31 03:03:43 +01:00
parent cd9dc163e5
commit 6e79465362
4 changed files with 28 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -30,7 +30,7 @@
//!
//! HttpServer::new(move || {
//! App::new()
//! .wrap(RateLimiter)
//! .wrap(RateLimiter::default())
//! .app_data(limiter.clone())
//! .service(index)
//! })

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

@ -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)
}