1 use crate::{settings::structs::RateLimitConfig, utils::get_ip, IpAddr, LemmyError};
2 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
3 use futures::future::{ok, Ready};
4 use rate_limiter::{RateLimitType, RateLimiter};
11 use tokio::sync::Mutex;
15 #[derive(Debug, Clone)]
16 pub struct RateLimit {
17 // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
18 // across await points
19 pub rate_limiter: Arc<Mutex<RateLimiter>>,
20 pub rate_limit_config: RateLimitConfig,
23 #[derive(Debug, Clone)]
24 pub struct RateLimited {
25 rate_limiter: Arc<Mutex<RateLimiter>>,
26 rate_limit_config: RateLimitConfig,
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 pub fn image(&self) -> RateLimited {
49 self.kind(RateLimitType::Image)
52 pub fn comment(&self) -> RateLimited {
53 self.kind(RateLimitType::Comment)
56 fn kind(&self, type_: RateLimitType) -> RateLimited {
58 rate_limiter: self.rate_limiter.clone(),
59 rate_limit_config: self.rate_limit_config.clone(),
66 pub async fn wrap<T, E>(
69 fut: impl Future<Output = Result<T, E>>,
74 // Does not need to be blocking because the RwLock in settings never held across await points,
75 // and the operation here locks only long enough to clone
76 let rate_limit = self.rate_limit_config;
80 let mut limiter = self.rate_limiter.lock().await;
83 RateLimitType::Message => {
84 limiter.check_rate_limit_full(
88 rate_limit.message_per_second,
95 RateLimitType::Post => {
96 limiter.check_rate_limit_full(
100 rate_limit.post_per_second,
104 RateLimitType::Register => {
105 limiter.check_rate_limit_full(
109 rate_limit.register_per_second,
113 RateLimitType::Image => {
114 limiter.check_rate_limit_full(
118 rate_limit.image_per_second,
122 RateLimitType::Comment => {
123 limiter.check_rate_limit_full(
127 rate_limit.comment_per_second,
138 let mut limiter = self.rate_limiter.lock().await;
141 RateLimitType::Post => {
142 limiter.check_rate_limit_full(
146 rate_limit.post_per_second,
150 RateLimitType::Register => {
151 limiter.check_rate_limit_full(
155 rate_limit.register_per_second,
168 impl<S> Transform<S, ServiceRequest> for RateLimited
170 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
173 type Response = S::Response;
174 type Error = actix_web::Error;
176 type Transform = RateLimitedMiddleware<S>;
177 type Future = Ready<Result<Self::Transform, Self::InitError>>;
179 fn new_transform(&self, service: S) -> Self::Future {
180 ok(RateLimitedMiddleware {
181 rate_limited: self.clone(),
187 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
189 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
191 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
194 type Response = S::Response;
195 type Error = actix_web::Error;
196 type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
198 fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
199 self.service.poll_ready(cx)
202 fn call(&self, req: ServiceRequest) -> Self::Future {
203 let ip_addr = get_ip(&req.connection_info());
208 .wrap(ip_addr, self.service.call(req));
210 Box::pin(async move { fut.await.map_err(actix_web::Error::from) })