use crate::decision::Decision;
use crate::quota::Quota;
#[derive(Debug, Clone)]
pub struct RateLimitExt {
pub key: String,
pub quota: Quota,
pub decision: Decision,
pub allowed: bool,
pub remaining: u64,
pub limit: u64,
pub reset_seconds: u64,
}
impl RateLimitExt {
pub fn new(key: impl Into<String>, quota: Quota, decision: Decision) -> Self {
let info = decision.info();
Self {
key: key.into(),
allowed: decision.is_allowed(),
remaining: info.remaining,
limit: info.limit,
reset_seconds: info.reset_seconds(),
quota,
decision,
}
}
pub fn is_allowed(&self) -> bool {
self.allowed
}
pub fn is_denied(&self) -> bool {
!self.allowed
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct RateLimitResponse {
pub allowed: bool,
pub limit: u64,
pub remaining: u64,
pub reset_in_seconds: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub retry_after_seconds: Option<u64>,
}
impl From<&RateLimitExt> for RateLimitResponse {
fn from(ext: &RateLimitExt) -> Self {
Self {
allowed: ext.allowed,
limit: ext.limit,
remaining: ext.remaining,
reset_in_seconds: ext.reset_seconds,
retry_after_seconds: ext
.decision
.info()
.retry_after
.map(|d| d.as_secs()),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::decision::RateLimitInfo;
use std::time::{Duration, Instant};
#[test]
fn test_rate_limit_ext() {
let info = RateLimitInfo::new(100, 50, Instant::now() + Duration::from_secs(60), Instant::now());
let decision = Decision::allowed(info);
let quota = Quota::per_minute(100);
let ext = RateLimitExt::new("user:123", quota, decision);
assert!(ext.is_allowed());
assert!(!ext.is_denied());
assert_eq!(ext.remaining, 50);
assert_eq!(ext.limit, 100);
}
#[test]
fn test_rate_limit_response_serialization() {
let info = RateLimitInfo::new(100, 0, Instant::now() + Duration::from_secs(30), Instant::now())
.with_retry_after(Duration::from_secs(30));
let decision = Decision::denied(info);
let quota = Quota::per_minute(100);
let ext = RateLimitExt::new("user:123", quota, decision);
let response: RateLimitResponse = (&ext).into();
assert!(!response.allowed);
assert_eq!(response.limit, 100);
assert_eq!(response.remaining, 0);
assert!(response.retry_after_seconds.is_some());
}
}