1 use crate::{ApiError, IpAddr, LemmyError};
3 use std::{collections::HashMap, time::SystemTime};
4 use strum::IntoEnumIterator;
6 #[derive(Debug, Clone)]
7 struct RateLimitBucket {
8 last_checked: SystemTime,
12 #[derive(Eq, PartialEq, Hash, Debug, EnumIter, Copy, Clone, AsRefStr)]
13 pub(crate) enum RateLimitType {
20 /// Rate limiting based on rate type and IP addr
21 #[derive(Debug, Clone)]
22 pub struct RateLimiter {
23 buckets: HashMap<RateLimitType, HashMap<IpAddr, RateLimitBucket>>,
26 impl Default for RateLimiter {
27 fn default() -> Self {
29 buckets: HashMap::<RateLimitType, HashMap<IpAddr, RateLimitBucket>>::new(),
35 fn insert_ip(&mut self, ip: &IpAddr) {
36 for rate_limit_type in RateLimitType::iter() {
37 if self.buckets.get(&rate_limit_type).is_none() {
38 self.buckets.insert(rate_limit_type, HashMap::new());
41 if let Some(bucket) = self.buckets.get_mut(&rate_limit_type) {
42 if bucket.get(ip).is_none() {
46 last_checked: SystemTime::now(),
55 #[allow(clippy::float_cmp)]
56 pub(super) fn check_rate_limit_full(
63 ) -> Result<(), LemmyError> {
65 if let Some(bucket) = self.buckets.get_mut(&type_) {
66 if let Some(rate_limit) = bucket.get_mut(ip) {
67 let current = SystemTime::now();
68 let time_passed = current.duration_since(rate_limit.last_checked)?.as_secs() as f64;
71 if rate_limit.allowance == -2f64 {
72 rate_limit.allowance = rate as f64;
75 rate_limit.last_checked = current;
76 rate_limit.allowance += time_passed * (rate as f64 / per as f64);
77 if !check_only && rate_limit.allowance > rate as f64 {
78 rate_limit.allowance = rate as f64;
81 if rate_limit.allowance < 1.0 {
83 "Rate limited type: {}, IP: {}, time_passed: {}, allowance: {}",
92 "Too many requests. type: {}, IP: {}, {} per {} seconds",
103 rate_limit.allowance -= 1.0;