rapid-rs 0.5.0

Zero-config, batteries-included web framework for Rust - FastAPI meets Spring Boot
Documentation 
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

//! Rate limiting middleware implementation

use axum::{
    extract::{Request, State},
    http::StatusCode,
    middleware::Next,
    response::{IntoResponse, Response},
    Json,
};
use governor::{
    clock::DefaultClock,
    state::{InMemoryState, NotKeyed},
    Quota, RateLimiter as GovernorRateLimiter,
};
use serde::Serialize;
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;

/// Rate limit configuration
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Number of requests allowed per period
    pub requests_per_period: u32,
    
    /// Time period for rate limiting
    pub period: Duration,
    
    /// Burst size (max requests in a short burst)
    pub burst_size: u32,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            requests_per_period: 100,
            period: Duration::from_secs(60),
            burst_size: 10,
        }
    }
}

/// Rate limiter
#[derive(Clone)]
pub struct RateLimiter {
    limiter: Arc<GovernorRateLimiter<NotKeyed, InMemoryState, DefaultClock>>,
}

impl RateLimiter {
    pub fn new(config: RateLimitConfig) -> Self {
        let quota = Quota::with_period(config.period)
            .unwrap()
            .allow_burst(NonZeroU32::new(config.burst_size).unwrap());
        
        Self {
            limiter: Arc::new(GovernorRateLimiter::direct(quota)),
        }
    }
    
    /// Check if request is allowed
    pub fn check(&self) -> bool {
        self.limiter.check().is_ok()
    }
}

#[derive(Serialize)]
struct RateLimitError {
    code: String,
    message: String,
    retry_after_seconds: u64,
}

/// Rate limiting middleware
pub async fn rate_limit_middleware(
    State(limiter): State<RateLimiter>,
    request: Request,
    next: Next,
) -> Response {
    if limiter.check() {
        next.run(request).await
    } else {
        let error = RateLimitError {
            code: "RATE_LIMIT_EXCEEDED".to_string(),
            message: "Too many requests. Please try again later.".to_string(),
            retry_after_seconds: 60,
        };
        
        (StatusCode::TOO_MANY_REQUESTS, Json(error)).into_response()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_rate_limiter() {
        let config = RateLimitConfig {
            requests_per_period: 2,
            period: Duration::from_secs(1),
            burst_size: 2,
        };
        
        let limiter = RateLimiter::new(config);
        
        // First two requests should pass
        assert!(limiter.check());
        assert!(limiter.check());
        
        // Third should fail
        assert!(!limiter.check());
    }
}