skp_ratelimit/
extensions.rs

1//! Request extensions for accessing rate limit info in handlers.
2//!
3//! This module provides extension types that can be injected into
4//! request handlers to access rate limit information.
5//!
6//! # Example
7//!
8//! ```ignore
9//! use axum::Extension;
10//! use oc_ratelimit_advanced::extensions::RateLimitExt;
11//!
12//! async fn handler(Extension(rate_limit): Extension<RateLimitExt>) {
13//!     println!("Remaining: {}", rate_limit.remaining);
14//! }
15//! ```
16
17use crate::decision::Decision;
18use crate::quota::Quota;
19
20/// Rate limit information available via request extensions.
21///
22/// This is automatically added to requests when using the rate limit middleware.
23#[derive(Debug, Clone)]
24pub struct RateLimitExt {
25    /// The key used for rate limiting this request.
26    pub key: String,
27    /// The quota applied to this request.
28    pub quota: Quota,
29    /// The rate limit decision.
30    pub decision: Decision,
31    /// Whether the request was allowed.
32    pub allowed: bool,
33    /// Remaining requests in the current window.
34    pub remaining: u64,
35    /// Maximum requests allowed.
36    pub limit: u64,
37    /// Seconds until reset.
38    pub reset_seconds: u64,
39}
40
41impl RateLimitExt {
42    /// Create a new rate limit extension from a decision.
43    pub fn new(key: impl Into<String>, quota: Quota, decision: Decision) -> Self {
44        let info = decision.info();
45        Self {
46            key: key.into(),
47            allowed: decision.is_allowed(),
48            remaining: info.remaining,
49            limit: info.limit,
50            reset_seconds: info.reset_seconds(),
51            quota,
52            decision,
53        }
54    }
55
56    /// Check if the request was allowed.
57    pub fn is_allowed(&self) -> bool {
58        self.allowed
59    }
60
61    /// Check if the request was denied.
62    pub fn is_denied(&self) -> bool {
63        !self.allowed
64    }
65}
66
67/// Rate limit info that can be serialized to JSON.
68///
69/// Useful for returning rate limit information in API responses.
70#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
71pub struct RateLimitResponse {
72    /// Whether the request was allowed.
73    pub allowed: bool,
74    /// Maximum requests allowed per window.
75    pub limit: u64,
76    /// Remaining requests in current window.
77    pub remaining: u64,
78    /// Seconds until the rate limit resets.
79    pub reset_in_seconds: u64,
80    /// If denied, the reason.
81    #[serde(skip_serializing_if = "Option::is_none")]
82    pub retry_after_seconds: Option<u64>,
83}
84
85impl From<&RateLimitExt> for RateLimitResponse {
86    fn from(ext: &RateLimitExt) -> Self {
87        Self {
88            allowed: ext.allowed,
89            limit: ext.limit,
90            remaining: ext.remaining,
91            reset_in_seconds: ext.reset_seconds,
92            retry_after_seconds: ext
93                .decision
94                .info()
95                .retry_after
96                .map(|d| d.as_secs()),
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::decision::RateLimitInfo;
105    use std::time::{Duration, Instant};
106
107    #[test]
108    fn test_rate_limit_ext() {
109        let info = RateLimitInfo::new(100, 50, Instant::now() + Duration::from_secs(60), Instant::now());
110        let decision = Decision::allowed(info);
111        let quota = Quota::per_minute(100);
112
113        let ext = RateLimitExt::new("user:123", quota, decision);
114
115        assert!(ext.is_allowed());
116        assert!(!ext.is_denied());
117        assert_eq!(ext.remaining, 50);
118        assert_eq!(ext.limit, 100);
119    }
120
121    #[test]
122    fn test_rate_limit_response_serialization() {
123        let info = RateLimitInfo::new(100, 0, Instant::now() + Duration::from_secs(30), Instant::now())
124            .with_retry_after(Duration::from_secs(30));
125        let decision = Decision::denied(info);
126        let quota = Quota::per_minute(100);
127
128        let ext = RateLimitExt::new("user:123", quota, decision);
129        let response: RateLimitResponse = (&ext).into();
130
131        assert!(!response.allowed);
132        assert_eq!(response.limit, 100);
133        assert_eq!(response.remaining, 0);
134        assert!(response.retry_after_seconds.is_some());
135    }
136}
skp_ratelimit/extensions.rs

skp_ratelimit/
extensions.rs
