2 use crate::{get_ip, LemmyError};
3 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
4 use futures::future::{ok, Ready};
5 use lemmy_utils::settings::{RateLimitConfig, Settings};
6 use rate_limiter::{RateLimitType, RateLimiter};
11 task::{Context, Poll},
13 use tokio::sync::Mutex;
17 #[derive(Debug, Clone)]
18 pub struct RateLimit {
19 // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
20 // across await points
21 pub rate_limiter: Arc<Mutex<RateLimiter>>,
24 #[derive(Debug, Clone)]
25 pub struct RateLimited {
26 rate_limiter: Arc<Mutex<RateLimiter>>,
30 pub struct RateLimitedMiddleware<S> {
31 rate_limited: RateLimited,
36 pub fn message(&self) -> RateLimited {
37 self.kind(RateLimitType::Message)
40 pub fn post(&self) -> RateLimited {
41 self.kind(RateLimitType::Post)
44 pub fn register(&self) -> RateLimited {
45 self.kind(RateLimitType::Register)
48 fn kind(&self, type_: RateLimitType) -> RateLimited {
50 rate_limiter: self.rate_limiter.clone(),
57 pub async fn wrap<T, E>(
60 fut: impl Future<Output = Result<T, E>>,
65 // Does not need to be blocking because the RwLock in settings never held across await points,
66 // and the operation here locks only long enough to clone
67 let rate_limit: RateLimitConfig = Settings::get().rate_limit;
71 let mut limiter = self.rate_limiter.lock().await;
74 RateLimitType::Message => {
75 limiter.check_rate_limit_full(
79 rate_limit.message_per_second,
86 RateLimitType::Post => {
87 limiter.check_rate_limit_full(
91 rate_limit.post_per_second,
95 RateLimitType::Register => {
96 limiter.check_rate_limit_full(
100 rate_limit.register_per_second,
111 let mut limiter = self.rate_limiter.lock().await;
114 RateLimitType::Post => {
115 limiter.check_rate_limit_full(
119 rate_limit.post_per_second,
123 RateLimitType::Register => {
124 limiter.check_rate_limit_full(
128 rate_limit.register_per_second,
141 impl<S> Transform<S> for RateLimited
143 S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
146 type Request = S::Request;
147 type Response = S::Response;
148 type Error = actix_web::Error;
150 type Transform = RateLimitedMiddleware<S>;
151 type Future = Ready<Result<Self::Transform, Self::InitError>>;
153 fn new_transform(&self, service: S) -> Self::Future {
154 ok(RateLimitedMiddleware {
155 rate_limited: self.clone(),
161 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
163 impl<S> Service for RateLimitedMiddleware<S>
165 S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
168 type Request = S::Request;
169 type Response = S::Response;
170 type Error = actix_web::Error;
171 type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
173 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
174 self.service.poll_ready(cx)
177 fn call(&mut self, req: S::Request) -> Self::Future {
178 let ip_addr = get_ip(&req.connection_info());
183 .wrap(ip_addr, self.service.call(req));
185 Box::pin(async move { fut.await.map_err(actix_web::Error::from) })