2020-04-19 22:08:25 +00:00
|
|
|
pub mod rate_limiter;
|
|
|
|
|
|
|
|
use super::{IPAddr, Settings};
|
|
|
|
use crate::api::APIError;
|
2020-04-20 18:02:25 +00:00
|
|
|
use crate::get_ip;
|
2020-04-20 03:59:07 +00:00
|
|
|
use crate::settings::RateLimitConfig;
|
|
|
|
use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
|
2020-04-19 22:08:25 +00:00
|
|
|
use failure::Error;
|
2020-04-20 03:59:07 +00:00
|
|
|
use futures::future::{ok, Ready};
|
2020-04-20 14:45:14 +00:00
|
|
|
use log::debug;
|
2020-04-20 03:59:07 +00:00
|
|
|
use rate_limiter::{RateLimitType, RateLimiter};
|
2020-04-19 22:08:25 +00:00
|
|
|
use std::collections::HashMap;
|
2020-04-20 03:59:07 +00:00
|
|
|
use std::future::Future;
|
|
|
|
use std::pin::Pin;
|
2020-04-19 22:08:25 +00:00
|
|
|
use std::sync::Arc;
|
2020-04-20 03:59:07 +00:00
|
|
|
use std::task::{Context, Poll};
|
2020-04-19 22:08:25 +00:00
|
|
|
use std::time::SystemTime;
|
|
|
|
use strum::IntoEnumIterator;
|
2020-04-20 03:59:07 +00:00
|
|
|
use tokio::sync::Mutex;
|
2020-04-19 22:08:25 +00:00
|
|
|
|
|
|
|
#[derive(Debug, Clone)]
|
2020-04-20 17:51:42 +00:00
|
|
|
pub struct RateLimit {
|
2020-04-19 22:08:25 +00:00
|
|
|
pub rate_limiter: Arc<Mutex<RateLimiter>>,
|
2020-04-20 17:51:42 +00:00
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
|
|
|
|
#[derive(Debug, Clone)]
|
2020-04-20 17:51:42 +00:00
|
|
|
pub struct RateLimited {
|
|
|
|
rate_limiter: Arc<Mutex<RateLimiter>>,
|
|
|
|
type_: RateLimitType,
|
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
|
2020-04-20 17:51:42 +00:00
|
|
|
pub struct RateLimitedMiddleware<S> {
|
|
|
|
rate_limited: RateLimited,
|
|
|
|
service: S,
|
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
|
|
|
|
impl RateLimit {
|
|
|
|
pub fn message(&self) -> RateLimited {
|
|
|
|
self.kind(RateLimitType::Message)
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn post(&self) -> RateLimited {
|
|
|
|
self.kind(RateLimitType::Post)
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn register(&self) -> RateLimited {
|
|
|
|
self.kind(RateLimitType::Register)
|
|
|
|
}
|
|
|
|
|
|
|
|
fn kind(&self, type_: RateLimitType) -> RateLimited {
|
2020-04-20 17:51:42 +00:00
|
|
|
RateLimited {
|
|
|
|
rate_limiter: self.rate_limiter.clone(),
|
|
|
|
type_,
|
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl RateLimited {
|
|
|
|
pub async fn wrap<T, E>(
|
|
|
|
self,
|
|
|
|
ip_addr: String,
|
|
|
|
fut: impl Future<Output = Result<T, E>>,
|
|
|
|
) -> Result<T, E>
|
|
|
|
where
|
|
|
|
E: From<failure::Error>,
|
|
|
|
{
|
|
|
|
let rate_limit: RateLimitConfig = actix_web::web::block(move || {
|
|
|
|
// needs to be in a web::block because the RwLock in settings is from stdlib
|
2020-04-20 04:43:30 +00:00
|
|
|
Ok(Settings::get().rate_limit) as Result<_, failure::Error>
|
2020-04-20 03:59:07 +00:00
|
|
|
})
|
|
|
|
.await
|
|
|
|
.map_err(|e| match e {
|
|
|
|
actix_web::error::BlockingError::Error(e) => e,
|
|
|
|
_ => APIError::err("Operation canceled").into(),
|
|
|
|
})?;
|
|
|
|
|
|
|
|
// before
|
|
|
|
{
|
2020-04-20 17:51:42 +00:00
|
|
|
let mut limiter = self.rate_limiter.lock().await;
|
2020-04-20 03:59:07 +00:00
|
|
|
|
2020-04-20 17:51:42 +00:00
|
|
|
match self.type_ {
|
2020-04-20 03:59:07 +00:00
|
|
|
RateLimitType::Message => {
|
|
|
|
limiter.check_rate_limit_full(
|
2020-04-20 17:51:42 +00:00
|
|
|
self.type_,
|
2020-04-20 03:59:07 +00:00
|
|
|
&ip_addr,
|
|
|
|
rate_limit.message,
|
|
|
|
rate_limit.message_per_second,
|
|
|
|
false,
|
|
|
|
)?;
|
|
|
|
|
|
|
|
return fut.await;
|
|
|
|
}
|
|
|
|
RateLimitType::Post => {
|
|
|
|
limiter.check_rate_limit_full(
|
2020-04-20 17:51:42 +00:00
|
|
|
self.type_.clone(),
|
2020-04-20 03:59:07 +00:00
|
|
|
&ip_addr,
|
|
|
|
rate_limit.post,
|
|
|
|
rate_limit.post_per_second,
|
|
|
|
true,
|
|
|
|
)?;
|
|
|
|
}
|
|
|
|
RateLimitType::Register => {
|
|
|
|
limiter.check_rate_limit_full(
|
2020-04-20 17:51:42 +00:00
|
|
|
self.type_,
|
2020-04-20 03:59:07 +00:00
|
|
|
&ip_addr,
|
|
|
|
rate_limit.register,
|
|
|
|
rate_limit.register_per_second,
|
|
|
|
true,
|
|
|
|
)?;
|
|
|
|
}
|
|
|
|
};
|
|
|
|
}
|
|
|
|
|
|
|
|
let res = fut.await;
|
|
|
|
|
|
|
|
// after
|
|
|
|
{
|
2020-04-20 17:51:42 +00:00
|
|
|
let mut limiter = self.rate_limiter.lock().await;
|
2020-04-20 03:59:07 +00:00
|
|
|
if res.is_ok() {
|
2020-04-20 17:51:42 +00:00
|
|
|
match self.type_ {
|
2020-04-20 03:59:07 +00:00
|
|
|
RateLimitType::Post => {
|
|
|
|
limiter.check_rate_limit_full(
|
2020-04-20 17:51:42 +00:00
|
|
|
self.type_,
|
2020-04-20 03:59:07 +00:00
|
|
|
&ip_addr,
|
|
|
|
rate_limit.post,
|
|
|
|
rate_limit.post_per_second,
|
|
|
|
false,
|
|
|
|
)?;
|
|
|
|
}
|
|
|
|
RateLimitType::Register => {
|
|
|
|
limiter.check_rate_limit_full(
|
2020-04-20 17:51:42 +00:00
|
|
|
self.type_,
|
2020-04-20 03:59:07 +00:00
|
|
|
&ip_addr,
|
|
|
|
rate_limit.register,
|
|
|
|
rate_limit.register_per_second,
|
|
|
|
false,
|
|
|
|
)?;
|
|
|
|
}
|
|
|
|
_ => (),
|
|
|
|
};
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
res
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<S> Transform<S> for RateLimited
|
|
|
|
where
|
|
|
|
S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
|
|
|
|
S::Future: 'static,
|
|
|
|
{
|
|
|
|
type Request = S::Request;
|
|
|
|
type Response = S::Response;
|
|
|
|
type Error = actix_web::Error;
|
|
|
|
type InitError = ();
|
|
|
|
type Transform = RateLimitedMiddleware<S>;
|
|
|
|
type Future = Ready<Result<Self::Transform, Self::InitError>>;
|
|
|
|
|
|
|
|
fn new_transform(&self, service: S) -> Self::Future {
|
2020-04-20 17:51:42 +00:00
|
|
|
ok(RateLimitedMiddleware {
|
|
|
|
rate_limited: self.clone(),
|
|
|
|
service,
|
|
|
|
})
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-04-20 04:43:30 +00:00
|
|
|
type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
|
|
|
|
|
2020-04-20 03:59:07 +00:00
|
|
|
impl<S> Service for RateLimitedMiddleware<S>
|
|
|
|
where
|
|
|
|
S: Service<Request = ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
|
|
|
|
S::Future: 'static,
|
|
|
|
{
|
|
|
|
type Request = S::Request;
|
|
|
|
type Response = S::Response;
|
|
|
|
type Error = actix_web::Error;
|
2020-04-20 04:43:30 +00:00
|
|
|
type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
|
2020-04-20 03:59:07 +00:00
|
|
|
|
|
|
|
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
2020-04-20 17:51:42 +00:00
|
|
|
self.service.poll_ready(cx)
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
fn call(&mut self, req: S::Request) -> Self::Future {
|
2020-04-20 18:02:25 +00:00
|
|
|
let ip_addr = get_ip(&req.connection_info());
|
2020-04-20 03:59:07 +00:00
|
|
|
|
2020-04-20 17:51:42 +00:00
|
|
|
let fut = self
|
|
|
|
.rate_limited
|
|
|
|
.clone()
|
|
|
|
.wrap(ip_addr, self.service.call(req));
|
2020-04-20 03:59:07 +00:00
|
|
|
|
|
|
|
Box::pin(async move { fut.await.map_err(actix_web::Error::from) })
|
|
|
|
}
|
2020-04-19 22:08:25 +00:00
|
|
|
}
|