2 settings::structs::{RateLimitConfig, Settings},
6 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
7 use futures::future::{ok, Ready};
8 use rate_limiter::{RateLimitType, RateLimiter};
13 task::{Context, Poll},
15 use tokio::sync::Mutex;
19 #[derive(Debug, Clone)]
20 pub struct RateLimit {
21 // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
22 // across await points
23 pub rate_limiter: Arc<Mutex<RateLimiter>>,
26 #[derive(Debug, Clone)]
27 pub struct RateLimited {
28 rate_limiter: Arc<Mutex<RateLimiter>>,
32 pub struct RateLimitedMiddleware<S> {
33 rate_limited: RateLimited,
38 pub fn message(&self) -> RateLimited {
39 self.kind(RateLimitType::Message)
42 pub fn post(&self) -> RateLimited {
43 self.kind(RateLimitType::Post)
46 pub fn register(&self) -> RateLimited {
47 self.kind(RateLimitType::Register)
50 pub fn image(&self) -> RateLimited {
51 self.kind(RateLimitType::Image)
54 fn kind(&self, type_: RateLimitType) -> RateLimited {
56 rate_limiter: self.rate_limiter.clone(),
63 pub async fn wrap<T, E>(
66 fut: impl Future<Output = Result<T, E>>,
71 // Does not need to be blocking because the RwLock in settings never held across await points,
72 // and the operation here locks only long enough to clone
73 let rate_limit: RateLimitConfig = Settings::get().rate_limit();
77 let mut limiter = self.rate_limiter.lock().await;
80 RateLimitType::Message => {
81 limiter.check_rate_limit_full(
85 rate_limit.message_per_second,
92 RateLimitType::Post => {
93 limiter.check_rate_limit_full(
97 rate_limit.post_per_second,
101 RateLimitType::Register => {
102 limiter.check_rate_limit_full(
106 rate_limit.register_per_second,
110 RateLimitType::Image => {
111 limiter.check_rate_limit_full(
115 rate_limit.image_per_second,
126 let mut limiter = self.rate_limiter.lock().await;
129 RateLimitType::Post => {
130 limiter.check_rate_limit_full(
134 rate_limit.post_per_second,
138 RateLimitType::Register => {
139 limiter.check_rate_limit_full(
143 rate_limit.register_per_second,
156 impl<S> Transform<S> for RateLimited
158 S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
161 type Request = S::Request;
162 type Response = S::Response;
163 type Error = actix_web::Error;
165 type Transform = RateLimitedMiddleware<S>;
166 type Future = Ready<Result<Self::Transform, Self::InitError>>;
168 fn new_transform(&self, service: S) -> Self::Future {
169 ok(RateLimitedMiddleware {
170 rate_limited: self.clone(),
176 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
178 impl<S> Service for RateLimitedMiddleware<S>
180 S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
183 type Request = S::Request;
184 type Response = S::Response;
185 type Error = actix_web::Error;
186 type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
188 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
189 self.service.poll_ready(cx)
192 fn call(&mut self, req: S::Request) -> Self::Future {
193 let ip_addr = get_ip(&req.connection_info());
198 .wrap(ip_addr, self.service.call(req));
200 Box::pin(async move { fut.await.map_err(actix_web::Error::from) })