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()
76 .unwrap_or_else(RateLimitConfig::default);
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,
129 let mut limiter = self.rate_limiter.lock().await;
132 RateLimitType::Post => {
133 limiter.check_rate_limit_full(
137 rate_limit.post_per_second,
141 RateLimitType::Register => {
142 limiter.check_rate_limit_full(
146 rate_limit.register_per_second,
159 impl<S> Transform<S, ServiceRequest> for RateLimited
161 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
164 type Response = S::Response;
165 type Error = actix_web::Error;
167 type Transform = RateLimitedMiddleware<S>;
168 type Future = Ready<Result<Self::Transform, Self::InitError>>;
170 fn new_transform(&self, service: S) -> Self::Future {
171 ok(RateLimitedMiddleware {
172 rate_limited: self.clone(),
178 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
180 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
182 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
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(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
190 self.service.poll_ready(cx)
193 fn call(&self, req: ServiceRequest) -> 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) })