2 settings::structs::{RateLimitConfig, Settings},
7 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
8 use futures::future::{ok, Ready};
9 use rate_limiter::{RateLimitType, RateLimiter};
14 task::{Context, Poll},
16 use tokio::sync::Mutex;
20 #[derive(Debug, Clone)]
21 pub struct RateLimit {
22 // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
23 // across await points
24 pub rate_limiter: Arc<Mutex<RateLimiter>>,
27 #[derive(Debug, Clone)]
28 pub struct RateLimited {
29 rate_limiter: Arc<Mutex<RateLimiter>>,
33 pub struct RateLimitedMiddleware<S> {
34 rate_limited: RateLimited,
39 pub fn message(&self) -> RateLimited {
40 self.kind(RateLimitType::Message)
43 pub fn post(&self) -> RateLimited {
44 self.kind(RateLimitType::Post)
47 pub fn register(&self) -> RateLimited {
48 self.kind(RateLimitType::Register)
51 pub fn image(&self) -> RateLimited {
52 self.kind(RateLimitType::Image)
55 fn kind(&self, type_: RateLimitType) -> RateLimited {
57 rate_limiter: self.rate_limiter.clone(),
64 pub async fn wrap<T, E>(
67 fut: impl Future<Output = Result<T, E>>,
72 // Does not need to be blocking because the RwLock in settings never held across await points,
73 // and the operation here locks only long enough to clone
74 let rate_limit: RateLimitConfig = Settings::get().rate_limit();
78 let mut limiter = self.rate_limiter.lock().await;
81 RateLimitType::Message => {
82 limiter.check_rate_limit_full(
86 rate_limit.message_per_second,
93 RateLimitType::Post => {
94 limiter.check_rate_limit_full(
98 rate_limit.post_per_second,
102 RateLimitType::Register => {
103 limiter.check_rate_limit_full(
107 rate_limit.register_per_second,
111 RateLimitType::Image => {
112 limiter.check_rate_limit_full(
116 rate_limit.image_per_second,
127 let mut limiter = self.rate_limiter.lock().await;
130 RateLimitType::Post => {
131 limiter.check_rate_limit_full(
135 rate_limit.post_per_second,
139 RateLimitType::Register => {
140 limiter.check_rate_limit_full(
144 rate_limit.register_per_second,
157 impl<S> Transform<S> for RateLimited
159 S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
162 type Request = S::Request;
163 type Response = S::Response;
164 type Error = actix_web::Error;
166 type Transform = RateLimitedMiddleware<S>;
167 type Future = Ready<Result<Self::Transform, Self::InitError>>;
169 fn new_transform(&self, service: S) -> Self::Future {
170 ok(RateLimitedMiddleware {
171 rate_limited: self.clone(),
177 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
179 impl<S> Service for RateLimitedMiddleware<S>
181 S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
184 type Request = S::Request;
185 type Response = S::Response;
186 type Error = actix_web::Error;
187 type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
189 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
190 self.service.poll_ready(cx)
193 fn call(&mut self, req: S::Request) -> Self::Future {
194 let ip_addr = get_ip(&req.connection_info());
199 .wrap(ip_addr, self.service.call(req));
201 Box::pin(async move { fut.await.map_err(actix_web::Error::from) })