crates/utils/src/rate_limit/mod.rs

   1 use crate::{utils::get_ip, IpAddr};
   2 use actix_web::{
   3   dev::{Service, ServiceRequest, ServiceResponse, Transform},
   4   HttpResponse,
   5 };
   6 use futures::future::{ok, Ready};
   7 use rate_limiter::{RateLimitType, RateLimiter};
   8 use serde::{Deserialize, Serialize};
   9 use std::{
  10   future::Future,
  11   pin::Pin,
  12   rc::Rc,
  13   sync::{Arc, Mutex},
  14   task::{Context, Poll},
  15 };
  16 use typed_builder::TypedBuilder;
  17
  18 pub mod rate_limiter;
  19
  20 #[derive(Debug, Deserialize, Serialize, Clone, TypedBuilder)]
  21 pub struct RateLimitConfig {
  22   #[builder(default = 180)]
  23   /// Maximum number of messages created in interval
  24   pub message: i32,
  25   #[builder(default = 60)]
  26   /// Interval length for message limit, in seconds
  27   pub message_per_second: i32,
  28   #[builder(default = 6)]
  29   /// Maximum number of posts created in interval
  30   pub post: i32,
  31   #[builder(default = 300)]
  32   /// Interval length for post limit, in seconds
  33   pub post_per_second: i32,
  34   #[builder(default = 3)]
  35   /// Maximum number of registrations in interval
  36   pub register: i32,
  37   #[builder(default = 3600)]
  38   /// Interval length for registration limit, in seconds
  39   pub register_per_second: i32,
  40   #[builder(default = 6)]
  41   /// Maximum number of image uploads in interval
  42   pub image: i32,
  43   #[builder(default = 3600)]
  44   /// Interval length for image uploads, in seconds
  45   pub image_per_second: i32,
  46   #[builder(default = 6)]
  47   /// Maximum number of comments created in interval
  48   pub comment: i32,
  49   #[builder(default = 600)]
  50   /// Interval length for comment limit, in seconds
  51   pub comment_per_second: i32,
  52   #[builder(default = 60)]
  53   /// Maximum number of searches created in interval
  54   pub search: i32,
  55   #[builder(default = 600)]
  56   /// Interval length for search limit, in seconds
  57   pub search_per_second: i32,
  58 }
  59
  60 #[derive(Debug, Clone)]
  61 pub struct RateLimit {
  62   // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
  63   // across await points
  64   pub rate_limiter: Arc<Mutex<RateLimiter>>,
  65   pub rate_limit_config: RateLimitConfig,
  66 }
  67
  68 #[derive(Debug, Clone)]
  69 pub struct RateLimited {
  70   rate_limiter: Arc<Mutex<RateLimiter>>,
  71   rate_limit_config: RateLimitConfig,
  72   type_: RateLimitType,
  73 }
  74
  75 pub struct RateLimitedMiddleware<S> {
  76   rate_limited: RateLimited,
  77   service: Rc<S>,
  78 }
  79
  80 impl RateLimit {
  81   pub fn message(&self) -> RateLimited {
  82     self.kind(RateLimitType::Message)
  83   }
  84
  85   pub fn post(&self) -> RateLimited {
  86     self.kind(RateLimitType::Post)
  87   }
  88
  89   pub fn register(&self) -> RateLimited {
  90     self.kind(RateLimitType::Register)
  91   }
  92
  93   pub fn image(&self) -> RateLimited {
  94     self.kind(RateLimitType::Image)
  95   }
  96
  97   pub fn comment(&self) -> RateLimited {
  98     self.kind(RateLimitType::Comment)
  99   }
 100
 101   pub fn search(&self) -> RateLimited {
 102     self.kind(RateLimitType::Search)
 103   }
 104
 105   fn kind(&self, type_: RateLimitType) -> RateLimited {
 106     RateLimited {
 107       rate_limiter: self.rate_limiter.clone(),
 108       rate_limit_config: self.rate_limit_config.clone(),
 109       type_,
 110     }
 111   }
 112 }
 113
 114 impl RateLimited {
 115   /// Returns true if the request passed the rate limit, false if it failed and should be rejected.
 116   pub fn check(self, ip_addr: IpAddr) -> bool {
 117     // Does not need to be blocking because the RwLock in settings never held across await points,
 118     // and the operation here locks only long enough to clone
 119     let rate_limit = self.rate_limit_config;
 120
 121     let (kind, interval) = match self.type_ {
 122       RateLimitType::Message => (rate_limit.message, rate_limit.message_per_second),
 123       RateLimitType::Post => (rate_limit.post, rate_limit.post_per_second),
 124       RateLimitType::Register => (rate_limit.register, rate_limit.register_per_second),
 125       RateLimitType::Image => (rate_limit.image, rate_limit.image_per_second),
 126       RateLimitType::Comment => (rate_limit.comment, rate_limit.comment_per_second),
 127       RateLimitType::Search => (rate_limit.search, rate_limit.search_per_second),
 128     };
 129     let mut limiter = self.rate_limiter.lock().expect("mutex poison error");
 130
 131     limiter.check_rate_limit_full(self.type_, &ip_addr, kind, interval)
 132   }
 133 }
 134
 135 impl<S> Transform<S, ServiceRequest> for RateLimited
 136 where
 137   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
 138   S::Future: 'static,
 139 {
 140   type Response = S::Response;
 141   type Error = actix_web::Error;
 142   type InitError = ();
 143   type Transform = RateLimitedMiddleware<S>;
 144   type Future = Ready<Result<Self::Transform, Self::InitError>>;
 145
 146   fn new_transform(&self, service: S) -> Self::Future {
 147     ok(RateLimitedMiddleware {
 148       rate_limited: self.clone(),
 149       service: Rc::new(service),
 150     })
 151   }
 152 }
 153
 154 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
 155
 156 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
 157 where
 158   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
 159   S::Future: 'static,
 160 {
 161   type Response = S::Response;
 162   type Error = actix_web::Error;
 163   type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
 164
 165   fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
 166     self.service.poll_ready(cx)
 167   }
 168
 169   fn call(&self, req: ServiceRequest) -> Self::Future {
 170     let ip_addr = get_ip(&req.connection_info());
 171
 172     let rate_limited = self.rate_limited.clone();
 173     let service = self.service.clone();
 174
 175     Box::pin(async move {
 176       if rate_limited.check(ip_addr) {
 177         service.call(req).await
 178       } else {
 179         let (http_req, _) = req.into_parts();
 180         // if rate limit was hit, respond with http 400
 181         Ok(ServiceResponse::new(
 182           http_req,
 183           HttpResponse::BadRequest().finish(),
 184         ))
 185       }
 186     })
 187   }
 188 }