2 settings::structs::{RateLimitConfig, Settings},
7 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
8 use futures::future::{ok, Ready};
9 use rate_limiter::{RateLimitType, RateLimiter};
14 task::{Context, Poll},
16 use tokio::sync::Mutex;
20 #[derive(Debug, Clone)]
21 pub struct RateLimit {
22 // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
23 // across await points
24 pub rate_limiter: Arc<Mutex<RateLimiter>>,
27 #[derive(Debug, Clone)]
28 pub struct RateLimited {
29 rate_limiter: Arc<Mutex<RateLimiter>>,
33 pub struct RateLimitedMiddleware<S> {
34 rate_limited: RateLimited,
39 pub fn message(&self) -> RateLimited {
40 self.kind(RateLimitType::Message)
43 pub fn post(&self) -> RateLimited {
44 self.kind(RateLimitType::Post)
47 pub fn register(&self) -> RateLimited {
48 self.kind(RateLimitType::Register)
51 pub fn image(&self) -> RateLimited {
52 self.kind(RateLimitType::Image)
55 fn kind(&self, type_: RateLimitType) -> RateLimited {
57 rate_limiter: self.rate_limiter.clone(),
64 pub async fn wrap<T, E>(
67 fut: impl Future<Output = Result<T, E>>,
72 // Does not need to be blocking because the RwLock in settings never held across await points,
73 // and the operation here locks only long enough to clone
74 let rate_limit: RateLimitConfig = Settings::get().rate_limit.unwrap_or_default();
78 let mut limiter = self.rate_limiter.lock().await;
81 RateLimitType::Message => {
82 limiter.check_rate_limit_full(
86 rate_limit.message_per_second,
93 RateLimitType::Post => {
94 limiter.check_rate_limit_full(
98 rate_limit.post_per_second,
102 RateLimitType::Register => {
103 limiter.check_rate_limit_full(
107 rate_limit.register_per_second,
111 RateLimitType::Image => {
112 limiter.check_rate_limit_full(
116 rate_limit.image_per_second,
127 let mut limiter = self.rate_limiter.lock().await;
130 RateLimitType::Post => {
131 limiter.check_rate_limit_full(
135 rate_limit.post_per_second,
139 RateLimitType::Register => {
140 limiter.check_rate_limit_full(
144 rate_limit.register_per_second,
157 impl<S> Transform<S, ServiceRequest> for RateLimited
159 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
162 type Response = S::Response;
163 type Error = actix_web::Error;
165 type Transform = RateLimitedMiddleware<S>;
166 type Future = Ready<Result<Self::Transform, Self::InitError>>;
168 fn new_transform(&self, service: S) -> Self::Future {
169 ok(RateLimitedMiddleware {
170 rate_limited: self.clone(),
176 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
178 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
180 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
183 type Response = S::Response;
184 type Error = actix_web::Error;
185 type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
187 fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
188 self.service.poll_ready(cx)
191 fn call(&self, req: ServiceRequest) -> Self::Future {
192 let ip_addr = get_ip(&req.connection_info());
197 .wrap(ip_addr, self.service.call(req));
199 Box::pin(async move { fut.await.map_err(actix_web::Error::from) })