1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
// Copyright (c) 2026 Kirky.X
// SPDX-License-Identifier: MIT
//! Quota Limiter
//!
//! Implements a simple quota-based limiter that tracks usage per key
//! with configurable limits and time windows.
use crate::error::LimiteronError;
#[cfg(feature = "quota-control")]
use crate::quota::QuotaConfig;
use async_trait::async_trait;
use dashmap::DashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
/// Quota usage record for a single key
#[derive(Debug, Clone)]
struct QuotaRecord {
/// Current usage count
usage: u64,
/// Window start time
window_start: Instant,
}
/// QuotaLimiter - A simple quota-based rate limiter
///
/// Tracks usage per identifier key within a time window.
/// When a key exceeds its quota limit, requests are rejected.
pub struct QuotaLimiter {
/// Quota configuration
config: QuotaConfig,
/// Per-key usage tracking (key -> usage, window_start)
usage: Arc<DashMap<String, QuotaRecord>>,
}
impl QuotaLimiter {
/// Creates a new QuotaLimiter with the given configuration.
///
/// # Arguments
/// * `config` - Quota configuration including limit, window size, etc.
///
/// # Panic
///
/// `config.window_size == 0` 时 panic(audit-L-003)。
/// `window_size = 0` 会导致 `Duration::from_secs(0)` 窗口立即过期,
/// 每次请求都重置 usage,配额限制形同虚设——这是配置 bug,应在开发阶段发现。
/// 用 `assert!` 而非 `Result` 以保持 API 兼容性(Rule 12:失败必须显性化)。
///
/// # Examples
/// ```rust
/// use limiteron::limiters::QuotaLimiter;
/// use limiteron::quota::QuotaConfig;
/// use limiteron::quota::QuotaType;
///
/// let config = QuotaConfig {
/// quota_type: QuotaType::Count,
/// limit: 1000,
/// window_size: 3600,
/// allow_overdraft: false,
/// overdraft_limit_percent: 20,
/// alert_config: Default::default(),
/// };
/// let limiter = QuotaLimiter::new(config);
/// ```
pub fn new(config: QuotaConfig) -> Self {
assert!(
config.window_size > 0,
"QuotaConfig.window_size must be greater than 0 (audit-L-003); \
window_size=0 would cause immediate window expiry, making quota useless"
);
Self {
config,
usage: Arc::new(DashMap::new()),
}
}
/// 获取配额上限
///
/// # 注意
///
/// 仅供 `LimiterManager` 参数一致性校验使用(audit-M-002),不应在业务代码中调用。
/// 业务代码应通过 `Limiter::check` 接口与 limiter 交互,而非直接读取配置。
pub fn max(&self) -> u64 {
self.config.limit
}
/// 获取配额周期
///
/// # 注意
///
/// 仅供 `LimiterManager` 参数一致性校验使用(audit-M-002),不应在业务代码中调用。
/// 业务代码应通过 `Limiter::check` 接口与 limiter 交互,而非直接读取配置。
pub fn period(&self) -> std::time::Duration {
std::time::Duration::from_secs(self.config.window_size)
}
/// Checks and consumes quota for the given key.
///
/// # Arguments
/// * `key` - The identifier key (user ID, API key, etc.)
///
/// # Returns
/// * `Ok(())` - Quota available, consumption successful
/// * `Err(LimiteronError)` - Quota exceeded or error
async fn check_and_consume(&self, key: &str) -> Result<bool, LimiteronError> {
let now = Instant::now();
let window_duration = Duration::from_secs(self.config.window_size);
let mut record = self
.usage
.entry(key.to_string())
.or_insert_with(|| QuotaRecord {
usage: 0,
window_start: now,
});
// Check if window has expired
if now.duration_since(record.window_start) >= window_duration {
// Reset for new window
record.usage = 0;
record.window_start = now;
}
// Check if quota allows overdraft
let max_usage = if self.config.allow_overdraft {
let overdraft_limit =
self.config.limit * self.config.overdraft_limit_percent as u64 / 100;
self.config.limit + overdraft_limit
} else {
self.config.limit
};
if record.usage >= max_usage {
return Err(LimiteronError::QuotaExceeded(format!(
"Quota exceeded for key '{}': used {}/{}",
key, record.usage, max_usage
)));
}
record.usage += 1;
Ok(true)
}
}
#[async_trait]
impl crate::limiters::Limiter for QuotaLimiter {
async fn allow(&self, _cost: u64) -> Result<bool, LimiteronError> {
// For quota limiter, we need a key to track, but the allow() method doesn't provide one
// This is a limitation - quota tracking requires a key
// Return true to allow the request (quota enforcement happens via check() with key)
Ok(true)
}
async fn check(&self, key: &str) -> Result<(), LimiteronError> {
self.check_and_consume(key).await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::limiters::Limiter;
use crate::quota::QuotaType;
fn create_test_config() -> QuotaConfig {
QuotaConfig {
quota_type: QuotaType::Count,
limit: 10,
window_size: 60,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: Default::default(),
}
}
#[tokio::test]
async fn test_quota_limiter_allows_within_limit() {
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
// Should allow 10 requests
for i in 0..10 {
let result = limiter.check("user1").await;
assert!(result.is_ok(), "Request {} should be allowed", i);
}
}
#[tokio::test]
async fn test_quota_limiter_rejects_over_limit() {
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
// Use up the quota
for _ in 0..10 {
let _ = limiter.check("user1").await;
}
// Next request should be rejected
let result = limiter.check("user1").await;
assert!(result.is_err());
assert!(matches!(result, Err(LimiteronError::QuotaExceeded(_))));
}
#[tokio::test]
async fn test_quota_limiter_independent_keys() {
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
// user1 uses 10 requests
for _ in 0..10 {
let _ = limiter.check("user1").await;
}
// user2 should still be able to make requests
let result = limiter.check("user2").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_quota_limiter_with_overdraft() {
let mut config = create_test_config();
config.allow_overdraft = true;
config.overdraft_limit_percent = 20; // 20% overdraft
let limiter = QuotaLimiter::new(config);
// Should allow 10 + 2 = 12 requests (10 limit + 20% overdraft)
for i in 0..12 {
let result = limiter.check("user1").await;
assert!(result.is_ok(), "Request {} should be allowed", i);
}
// Next request should be rejected
let result = limiter.check("user1").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_quota_limiter_allow_method() {
// allow() 方法不跟踪配额,总是返回 Ok(true)
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
let result = limiter.allow(1).await;
assert!(result.is_ok());
assert!(result.unwrap());
// 即使多次调用也应返回 Ok(true)
for _ in 0..100 {
let result = limiter.allow(1).await;
assert!(result.is_ok());
assert!(result.unwrap());
}
}
#[tokio::test]
async fn test_quota_limiter_allow_with_zero_cost() {
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
// allow() 不验证 cost,总是返回 Ok(true)
let result = limiter.allow(0).await;
assert!(result.is_ok());
assert!(result.unwrap());
}
#[tokio::test]
async fn test_quota_limiter_window_expiry_reset() {
// 使用 1 秒窗口,等待过期后验证重置
let mut config = create_test_config();
config.window_size = 1; // 1 秒窗口
config.limit = 3;
let limiter = QuotaLimiter::new(config);
// 用完配额
for _ in 0..3 {
assert!(limiter.check("user1").await.is_ok());
}
// 此时应被拒绝
assert!(limiter.check("user1").await.is_err());
// 等待窗口过期
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
// 窗口重置后应再次允许
let result = limiter.check("user1").await;
assert!(
result.is_ok(),
"Request after window reset should be allowed"
);
}
#[tokio::test]
async fn test_quota_limiter_overdraft_boundary() {
// 测试透支边界:limit + overdraft_limit 恰好用完
let mut config = create_test_config();
config.limit = 10;
config.allow_overdraft = true;
config.overdraft_limit_percent = 50; // 50% overdraft = 5 extra
let limiter = QuotaLimiter::new(config);
// Should allow 10 + 5 = 15 requests
for i in 0..15 {
let result = limiter.check("user1").await;
assert!(result.is_ok(), "Request {} should be allowed", i);
}
// 16th should be rejected
let result = limiter.check("user1").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_quota_limiter_check_propagates_error() {
// 当 check_and_consume 返回 Err 时,check 应传播错误
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
// 用完配额
for _ in 0..10 {
let _ = limiter.check("user1").await;
}
// 下一个 check 应返回 QuotaExceeded 错误
let result = limiter.check("user1").await;
assert!(result.is_err());
assert!(matches!(result, Err(LimiteronError::QuotaExceeded(_))));
}
// ========================================================================
// audit-macro-followup 修复20 (L-003): window_size=0 panic 测试
// ========================================================================
#[test]
#[should_panic(expected = "QuotaConfig.window_size must be greater than 0")]
fn test_quota_limiter_window_size_zero_panics() {
// audit-L-003: window_size=0 会导致窗口立即过期,配额限制失效
// 应在 new() 阶段 panic 而非静默接受错误配置(Rule 12)
let mut config = create_test_config();
config.window_size = 0;
let _ = QuotaLimiter::new(config);
}
#[tokio::test]
async fn test_quota_limiter_window_size_one_works() {
// 边界:window_size=1(最小合法值)应正常工作
let mut config = create_test_config();
config.window_size = 1;
config.limit = 3;
let limiter = QuotaLimiter::new(config);
// window_size=1s,前 3 个请求应成功
for _ in 0..3 {
assert!(limiter.check("boundary_user").await.is_ok());
}
// 第 4 个应失败
assert!(limiter.check("boundary_user").await.is_err());
}
}