crates/utils/src/rate_limit/rate_limiter.rs

   1 use crate::IpAddr;
   2 use std::{collections::HashMap, time::Instant};
   3 use strum::IntoEnumIterator;
   4 use tracing::debug;
   5
   6 #[derive(Debug, Clone)]
   7 struct RateLimitBucket {
   8   last_checked: Instant,
   9   allowance: f64,
  10 }
  11
  12 #[derive(Eq, PartialEq, Hash, Debug, EnumIter, Copy, Clone, AsRefStr)]
  13 pub(crate) enum RateLimitType {
  14   Message,
  15   Register,
  16   Post,
  17   Image,
  18   Comment,
  19 }
  20
  21 /// Rate limiting based on rate type and IP addr
  22 #[derive(Debug, Clone, Default)]
  23 pub struct RateLimiter {
  24   buckets: HashMap<RateLimitType, HashMap<IpAddr, RateLimitBucket>>,
  25 }
  26
  27 impl RateLimiter {
  28   fn insert_ip(&mut self, ip: &IpAddr) {
  29     for rate_limit_type in RateLimitType::iter() {
  30       if self.buckets.get(&rate_limit_type).is_none() {
  31         self.buckets.insert(rate_limit_type, HashMap::new());
  32       }
  33
  34       if let Some(bucket) = self.buckets.get_mut(&rate_limit_type) {
  35         if bucket.get(ip).is_none() {
  36           bucket.insert(
  37             ip.clone(),
  38             RateLimitBucket {
  39               last_checked: Instant::now(),
  40               allowance: -2f64,
  41             },
  42           );
  43         }
  44       }
  45     }
  46   }
  47
  48   /// Rate limiting Algorithm described here: https://stackoverflow.com/a/668327/1655478
  49   ///
  50   /// Returns true if the request passed the rate limit, false if it failed and should be rejected.
  51   #[allow(clippy::float_cmp)]
  52   pub(super) fn check_rate_limit_full(
  53     &mut self,
  54     type_: RateLimitType,
  55     ip: &IpAddr,
  56     rate: i32,
  57     per: i32,
  58   ) -> bool {
  59     self.insert_ip(ip);
  60     if let Some(bucket) = self.buckets.get_mut(&type_) {
  61       if let Some(rate_limit) = bucket.get_mut(ip) {
  62         let current = Instant::now();
  63         let time_passed = current.duration_since(rate_limit.last_checked).as_secs() as f64;
  64
  65         // The initial value
  66         if rate_limit.allowance == -2f64 {
  67           rate_limit.allowance = rate as f64;
  68         };
  69
  70         rate_limit.last_checked = current;
  71         rate_limit.allowance += time_passed * (rate as f64 / per as f64);
  72         if rate_limit.allowance > rate as f64 {
  73           rate_limit.allowance = rate as f64;
  74         }
  75
  76         if rate_limit.allowance < 1.0 {
  77           debug!(
  78             "Rate limited type: {}, IP: {}, time_passed: {}, allowance: {}",
  79             type_.as_ref(),
  80             ip,
  81             time_passed,
  82             rate_limit.allowance
  83           );
  84           false
  85         } else {
  86           rate_limit.allowance -= 1.0;
  87           true
  88         }
  89       } else {
  90         true
  91       }
  92     } else {
  93       true
  94     }
  95   }
  96 }