lazy_limit/
types.rs

1/* src/types.rs */
2
3use std::time::{SystemTime, UNIX_EPOCH};
4
5/// HTTP Method enum
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
7pub enum HttpMethod {
8    GET,
9    POST,
10    PUT,
11    DELETE,
12    PATCH,
13    HEAD,
14    OPTIONS,
15    CONNECT,
16    TRACE,
17    OTHER,
18}
19
20impl HttpMethod {
21    pub fn as_str(&self) -> &str {
22        match self {
23            Self::GET => "GET",
24            Self::POST => "POST",
25            Self::PUT => "PUT",
26            Self::DELETE => "DELETE",
27            Self::PATCH => "PATCH",
28            Self::HEAD => "HEAD",
29            Self::OPTIONS => "OPTIONS",
30            Self::CONNECT => "CONNECT",
31            Self::TRACE => "TRACE",
32            Self::OTHER => "OTHER",
33        }
34    }
35
36    pub fn from_str(s: &str) -> Option<Self> {
37        match s.to_uppercase().as_str() {
38            "GET" => Some(Self::GET),
39            "POST" => Some(Self::POST),
40            "PUT" => Some(Self::PUT),
41            "DELETE" => Some(Self::DELETE),
42            "PATCH" => Some(Self::PATCH),
43            "HEAD" => Some(Self::HEAD),
44            "OPTIONS" => Some(Self::OPTIONS),
45            "CONNECT" => Some(Self::CONNECT),
46            "TRACE" => Some(Self::TRACE),
47            _ => Some(Self::OTHER),
48        }
49    }
50}
51
52impl std::fmt::Display for HttpMethod {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        write!(f, "{}", self.as_str())
55    }
56}
57
58impl From<&str> for HttpMethod {
59    fn from(s: &str) -> Self {
60        HttpMethod::from_str(s).unwrap_or_else(|| {
61            panic!("Invalid HTTP method: {}. Expected one of: GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS, CONNECT, TRACE", s)
62        })
63    }
64}
65
66impl From<&String> for HttpMethod {
67    fn from(s: &String) -> Self {
68        HttpMethod::from(s.as_str())
69    }
70}
71
72impl From<String> for HttpMethod {
73    fn from(s: String) -> Self {
74        HttpMethod::from(s.as_str())
75    }
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
79pub enum Duration {
80    Seconds(u64),
81    Minutes(u64),
82    Hours(u64),
83    Days(u64),
84}
85
86impl Duration {
87    pub fn seconds(n: u64) -> Self {
88        Duration::Seconds(n)
89    }
90    pub fn minutes(n: u64) -> Self {
91        Duration::Minutes(n)
92    }
93    pub fn hours(n: u64) -> Self {
94        Duration::Hours(n)
95    }
96    pub fn days(n: u64) -> Self {
97        Duration::Days(n)
98    }
99
100    pub fn as_seconds(&self) -> u64 {
101        match self {
102            Duration::Seconds(n) => *n,
103            Duration::Minutes(n) => n * 60,
104            Duration::Hours(n) => n * 3600,
105            Duration::Days(n) => n * 86400,
106        }
107    }
108
109    pub fn is_short_interval(&self) -> bool {
110        self.as_seconds() <= 300 // 5 minutes
111    }
112}
113
114#[derive(Debug, Clone)]
115pub struct RuleConfig {
116    pub interval: Duration,
117    pub limit: u32,
118    pub is_prefix: bool,
119    pub methods: Option<Vec<HttpMethod>>,
120}
121
122impl RuleConfig {
123    pub fn new(interval: Duration, limit: u32) -> Self {
124        Self {
125            interval,
126            limit,
127            is_prefix: false,
128            methods: None,
129        }
130    }
131
132    pub fn match_prefix(mut self, is_prefix: bool) -> Self {
133        self.is_prefix = is_prefix;
134        self
135    }
136
137    pub fn for_methods(mut self, methods: Vec<HttpMethod>) -> Self {
138        self.methods = Some(methods);
139        self
140    }
141
142    /// Check if this rule applies to the given HTTP method
143    pub fn matches_method(&self, method: &Option<HttpMethod>) -> bool {
144        match (&self.methods, method) {
145            // If no methods specified, rule applies to all methods
146            (None, _) => true,
147            // If methods specified but no method provided, no match
148            (Some(_), None) => false,
149            // If both specified, check if method is in the list
150            (Some(allowed), Some(m)) => allowed.contains(m),
151        }
152    }
153}
154
155#[derive(Debug, Clone)]
156pub struct RequestRecord {
157    pub count: u32,
158    pub window_start: u64,
159    pub timestamps: Vec<u64>,
160}
161
162impl RequestRecord {
163    pub fn new(is_short_interval: bool) -> Self {
164        Self {
165            count: 0,
166            window_start: current_timestamp(),
167            timestamps: if is_short_interval {
168                Vec::new()
169            } else {
170                Vec::with_capacity(16)
171            },
172        }
173    }
174
175    pub fn add_request(&mut self, is_short_interval: bool, window_size: u64) {
176        let now = current_timestamp();
177
178        if is_short_interval {
179            if now.saturating_sub(self.window_start) >= window_size {
180                self.window_start = now;
181                self.count = 1;
182            } else {
183                self.count += 1;
184            }
185        } else {
186            self.timestamps.push(now);
187            let cutoff = now.saturating_sub(window_size);
188            self.timestamps.retain(|&t| t > cutoff);
189            self.count = self.timestamps.len() as u32;
190        }
191    }
192
193    pub fn is_limit_exceeded(&self, limit: u32, is_short_interval: bool, window_size: u64) -> bool {
194        let now = current_timestamp();
195        if is_short_interval {
196            if now.saturating_sub(self.window_start) >= window_size {
197                false
198            } else {
199                self.count >= limit
200            }
201        } else {
202            let cutoff = now.saturating_sub(window_size);
203            let valid_requests = self.timestamps.iter().filter(|&&t| t > cutoff).count() as u32;
204            valid_requests >= limit
205        }
206    }
207
208    pub fn memory_usage(&self) -> usize {
209        std::mem::size_of::<Self>() + self.timestamps.capacity() * std::mem::size_of::<u64>()
210    }
211
212    pub fn should_cleanup(&self, max_age_seconds: u64) -> bool {
213        let now = current_timestamp();
214        let last_activity = if !self.timestamps.is_empty() {
215            *self.timestamps.last().unwrap_or(&self.window_start)
216        } else {
217            self.window_start
218        };
219        now.saturating_sub(last_activity) > max_age_seconds
220    }
221}
222
223pub fn current_timestamp() -> u64 {
224    SystemTime::now()
225        .duration_since(UNIX_EPOCH)
226        .expect("Time went backwards")
227        .as_secs()
228}
lazy_limit/types.rs

lazy_limit/
types.rs
