1 use crate::{utils::get_ip, IpAddr};
3 dev::{Service, ServiceRequest, ServiceResponse, Transform},
6 use futures::future::{ok, Ready};
7 use rate_limiter::{RateLimitType, RateLimiter};
8 use serde::{Deserialize, Serialize};
14 task::{Context, Poll},
16 use typed_builder::TypedBuilder;
20 #[derive(Debug, Deserialize, Serialize, Clone, TypedBuilder)]
21 pub struct RateLimitConfig {
22 #[builder(default = 180)]
23 /// Maximum number of messages created in interval
25 #[builder(default = 60)]
26 /// Interval length for message limit, in seconds
27 pub message_per_second: i32,
28 #[builder(default = 6)]
29 /// Maximum number of posts created in interval
31 #[builder(default = 300)]
32 /// Interval length for post limit, in seconds
33 pub post_per_second: i32,
34 #[builder(default = 3)]
35 /// Maximum number of registrations in interval
37 #[builder(default = 3600)]
38 /// Interval length for registration limit, in seconds
39 pub register_per_second: i32,
40 #[builder(default = 6)]
41 /// Maximum number of image uploads in interval
43 #[builder(default = 3600)]
44 /// Interval length for image uploads, in seconds
45 pub image_per_second: i32,
46 #[builder(default = 6)]
47 /// Maximum number of comments created in interval
49 #[builder(default = 600)]
50 /// Interval length for comment limit, in seconds
51 pub comment_per_second: i32,
52 #[builder(default = 60)]
53 /// Maximum number of searches created in interval
55 #[builder(default = 600)]
56 /// Interval length for search limit, in seconds
57 pub search_per_second: i32,
60 #[derive(Debug, Clone)]
61 pub struct RateLimit {
62 // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
63 // across await points
64 pub rate_limiter: Arc<Mutex<RateLimiter>>,
65 pub rate_limit_config: RateLimitConfig,
68 #[derive(Debug, Clone)]
69 pub struct RateLimited {
70 rate_limiter: Arc<Mutex<RateLimiter>>,
71 rate_limit_config: RateLimitConfig,
75 pub struct RateLimitedMiddleware<S> {
76 rate_limited: RateLimited,
81 pub fn message(&self) -> RateLimited {
82 self.kind(RateLimitType::Message)
85 pub fn post(&self) -> RateLimited {
86 self.kind(RateLimitType::Post)
89 pub fn register(&self) -> RateLimited {
90 self.kind(RateLimitType::Register)
93 pub fn image(&self) -> RateLimited {
94 self.kind(RateLimitType::Image)
97 pub fn comment(&self) -> RateLimited {
98 self.kind(RateLimitType::Comment)
101 pub fn search(&self) -> RateLimited {
102 self.kind(RateLimitType::Search)
105 fn kind(&self, type_: RateLimitType) -> RateLimited {
107 rate_limiter: self.rate_limiter.clone(),
108 rate_limit_config: self.rate_limit_config.clone(),
115 /// Returns true if the request passed the rate limit, false if it failed and should be rejected.
116 pub fn check(self, ip_addr: IpAddr) -> bool {
117 // Does not need to be blocking because the RwLock in settings never held across await points,
118 // and the operation here locks only long enough to clone
119 let rate_limit = self.rate_limit_config;
121 let (kind, interval) = match self.type_ {
122 RateLimitType::Message => (rate_limit.message, rate_limit.message_per_second),
123 RateLimitType::Post => (rate_limit.post, rate_limit.post_per_second),
124 RateLimitType::Register => (rate_limit.register, rate_limit.register_per_second),
125 RateLimitType::Image => (rate_limit.image, rate_limit.image_per_second),
126 RateLimitType::Comment => (rate_limit.comment, rate_limit.comment_per_second),
127 RateLimitType::Search => (rate_limit.search, rate_limit.search_per_second),
129 let mut limiter = self.rate_limiter.lock().expect("mutex poison error");
131 limiter.check_rate_limit_full(self.type_, &ip_addr, kind, interval)
135 impl<S> Transform<S, ServiceRequest> for RateLimited
137 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
140 type Response = S::Response;
141 type Error = actix_web::Error;
143 type Transform = RateLimitedMiddleware<S>;
144 type Future = Ready<Result<Self::Transform, Self::InitError>>;
146 fn new_transform(&self, service: S) -> Self::Future {
147 ok(RateLimitedMiddleware {
148 rate_limited: self.clone(),
149 service: Rc::new(service),
154 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
156 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
158 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
161 type Response = S::Response;
162 type Error = actix_web::Error;
163 type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
165 fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
166 self.service.poll_ready(cx)
169 fn call(&self, req: ServiceRequest) -> Self::Future {
170 let ip_addr = get_ip(&req.connection_info());
172 let rate_limited = self.rate_limited.clone();
173 let service = self.service.clone();
175 Box::pin(async move {
176 if rate_limited.check(ip_addr) {
177 service.call(req).await
179 let (http_req, _) = req.into_parts();
180 // if rate limit was hit, respond with http 400
181 Ok(ServiceResponse::new(
183 HttpResponse::BadRequest().finish(),