2 use crate::{api::APIError, LemmyError};
4 use std::{collections::HashMap, time::SystemTime};
5 use strum::IntoEnumIterator;
7 #[derive(Debug, Clone)]
8 pub struct RateLimitBucket {
9 last_checked: SystemTime,
13 #[derive(Eq, PartialEq, Hash, Debug, EnumIter, Copy, Clone, AsRefStr)]
14 pub enum RateLimitType {
21 /// Rate limiting based on rate type and IP addr
22 #[derive(Debug, Clone)]
23 pub struct RateLimiter {
24 pub buckets: HashMap<RateLimitType, HashMap<IPAddr, RateLimitBucket>>,
27 impl Default for RateLimiter {
28 fn default() -> Self {
30 buckets: HashMap::new(),
36 fn insert_ip(&mut self, ip: &str) {
37 for rate_limit_type in RateLimitType::iter() {
38 if self.buckets.get(&rate_limit_type).is_none() {
39 self.buckets.insert(rate_limit_type, HashMap::new());
42 if let Some(bucket) = self.buckets.get_mut(&rate_limit_type) {
43 if bucket.get(ip).is_none() {
47 last_checked: SystemTime::now(),
56 #[allow(clippy::float_cmp)]
57 pub(super) fn check_rate_limit_full(
64 ) -> Result<(), LemmyError> {
66 if let Some(bucket) = self.buckets.get_mut(&type_) {
67 if let Some(rate_limit) = bucket.get_mut(ip) {
68 let current = SystemTime::now();
69 let time_passed = current.duration_since(rate_limit.last_checked)?.as_secs() as f64;
72 if rate_limit.allowance == -2f64 {
73 rate_limit.allowance = rate as f64;
76 rate_limit.last_checked = current;
77 rate_limit.allowance += time_passed * (rate as f64 / per as f64);
78 if !check_only && rate_limit.allowance > rate as f64 {
79 rate_limit.allowance = rate as f64;
82 if rate_limit.allowance < 1.0 {
84 "Rate limited type: {}, IP: {}, time_passed: {}, allowance: {}",
93 "Too many requests. type: {}, IP: {}, {} per {} seconds",
104 rate_limit.allowance -= 1.0;