68 lines
1.6 KiB
Rust
68 lines
1.6 KiB
Rust
use std::iter;
|
|
use std::num::NonZeroUsize;
|
|
use std::sync::Mutex;
|
|
use std::time::SystemTime;
|
|
|
|
use lru::LruCache;
|
|
|
|
pub struct Limiter {
|
|
record: Mutex<LruCache<i32, Vec<u64>>>,
|
|
amount: u64,
|
|
interval: u64,
|
|
}
|
|
|
|
impl Limiter {
|
|
pub fn init(amount: u64, interval: u64) -> Self {
|
|
Self {
|
|
record: Mutex::new(LruCache::new(NonZeroUsize::new(2000).unwrap())),
|
|
amount,
|
|
interval,
|
|
}
|
|
}
|
|
|
|
pub fn check(&self, uid: i32) -> bool {
|
|
let t = SystemTime::now()
|
|
.duration_since(SystemTime::UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_secs();
|
|
let mut r = self.record.lock().unwrap();
|
|
if let Some(ts) = r.pop(&uid) {
|
|
let new_ts: Vec<u64> = ts
|
|
.into_iter()
|
|
.chain(iter::once(t))
|
|
.filter(|&tt| tt + self.interval > t)
|
|
.collect();
|
|
let len = new_ts.len() as u64;
|
|
r.put(uid, new_ts);
|
|
len < self.amount
|
|
} else {
|
|
r.put(uid, vec![t]);
|
|
true
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct MainLimiters {
|
|
post_min: Limiter,
|
|
post_hour: Limiter,
|
|
get_hour: Limiter,
|
|
}
|
|
|
|
impl MainLimiters {
|
|
pub fn init() -> Self {
|
|
Self {
|
|
post_min: Limiter::init(6, 60),
|
|
post_hour: Limiter::init(50, 3600),
|
|
get_hour: Limiter::init(1000, 3600),
|
|
}
|
|
}
|
|
|
|
pub fn check(&self, is_post: bool, uid: i32) -> bool {
|
|
if is_post {
|
|
self.post_hour.check(uid) && self.post_min.check(uid)
|
|
} else {
|
|
self.get_hour.check(uid)
|
|
}
|
|
}
|
|
}
|