|
|
|
@ -1,7 +1,7 @@
|
|
|
|
|
use crate::{error::LemmyError, utils::get_ip, IpAddr};
|
|
|
|
|
use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
|
|
|
|
|
use futures::future::{ok, Ready};
|
|
|
|
|
use rate_limiter::{RateLimitType, RateLimiter};
|
|
|
|
|
use rate_limiter::{RateLimitStorage, RateLimitType};
|
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
|
use std::{
|
|
|
|
|
future::Future,
|
|
|
|
@ -10,6 +10,7 @@ use std::{
|
|
|
|
|
sync::{Arc, Mutex},
|
|
|
|
|
task::{Context, Poll},
|
|
|
|
|
};
|
|
|
|
|
use tokio::sync::{mpsc, mpsc::Sender, OnceCell};
|
|
|
|
|
use typed_builder::TypedBuilder;
|
|
|
|
|
|
|
|
|
|
pub mod rate_limiter;
|
|
|
|
@ -55,65 +56,102 @@ pub struct RateLimitConfig {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[derive(Debug, Clone)]
|
|
|
|
|
pub struct RateLimit {
|
|
|
|
|
// it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
|
|
|
|
|
// across await points
|
|
|
|
|
pub rate_limiter: Arc<Mutex<RateLimiter>>,
|
|
|
|
|
struct RateLimit {
|
|
|
|
|
pub rate_limiter: RateLimitStorage,
|
|
|
|
|
pub rate_limit_config: RateLimitConfig,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[derive(Debug, Clone)]
|
|
|
|
|
pub struct RateLimited {
|
|
|
|
|
rate_limiter: Arc<Mutex<RateLimiter>>,
|
|
|
|
|
rate_limit_config: RateLimitConfig,
|
|
|
|
|
pub struct RateLimitedGuard {
|
|
|
|
|
rate_limit: Arc<Mutex<RateLimit>>,
|
|
|
|
|
type_: RateLimitType,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
pub struct RateLimitedMiddleware<S> {
|
|
|
|
|
rate_limited: RateLimited,
|
|
|
|
|
service: Rc<S>,
|
|
|
|
|
/// Single instance of rate limit config and buckets, which is shared across all threads.
|
|
|
|
|
#[derive(Clone)]
|
|
|
|
|
pub struct RateLimitCell {
|
|
|
|
|
tx: Sender<RateLimitConfig>,
|
|
|
|
|
rate_limit: Arc<Mutex<RateLimit>>,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl RateLimit {
|
|
|
|
|
pub fn message(&self) -> RateLimited {
|
|
|
|
|
impl RateLimitCell {
|
|
|
|
|
/// Initialize cell if it wasnt initialized yet. Otherwise returns the existing cell.
|
|
|
|
|
pub async fn new(rate_limit_config: RateLimitConfig) -> &'static Self {
|
|
|
|
|
static LOCAL_INSTANCE: OnceCell<RateLimitCell> = OnceCell::const_new();
|
|
|
|
|
LOCAL_INSTANCE
|
|
|
|
|
.get_or_init(|| async {
|
|
|
|
|
let (tx, mut rx) = mpsc::channel::<RateLimitConfig>(4);
|
|
|
|
|
let rate_limit = Arc::new(Mutex::new(RateLimit {
|
|
|
|
|
rate_limiter: Default::default(),
|
|
|
|
|
rate_limit_config,
|
|
|
|
|
}));
|
|
|
|
|
let rate_limit2 = rate_limit.clone();
|
|
|
|
|
tokio::spawn(async move {
|
|
|
|
|
while let Some(r) = rx.recv().await {
|
|
|
|
|
rate_limit2
|
|
|
|
|
.lock()
|
|
|
|
|
.expect("Failed to lock rate limit mutex for updating")
|
|
|
|
|
.rate_limit_config = r;
|
|
|
|
|
}
|
|
|
|
|
});
|
|
|
|
|
RateLimitCell { tx, rate_limit }
|
|
|
|
|
})
|
|
|
|
|
.await
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Call this when the config was updated, to update all in-memory cells.
|
|
|
|
|
pub async fn send(&self, config: RateLimitConfig) -> Result<(), LemmyError> {
|
|
|
|
|
self.tx.send(config).await?;
|
|
|
|
|
Ok(())
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
pub fn message(&self) -> RateLimitedGuard {
|
|
|
|
|
self.kind(RateLimitType::Message)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
pub fn post(&self) -> RateLimited {
|
|
|
|
|
pub fn post(&self) -> RateLimitedGuard {
|
|
|
|
|
self.kind(RateLimitType::Post)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
pub fn register(&self) -> RateLimited {
|
|
|
|
|
pub fn register(&self) -> RateLimitedGuard {
|
|
|
|
|
self.kind(RateLimitType::Register)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
pub fn image(&self) -> RateLimited {
|
|
|
|
|
pub fn image(&self) -> RateLimitedGuard {
|
|
|
|
|
self.kind(RateLimitType::Image)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
pub fn comment(&self) -> RateLimited {
|
|
|
|
|
pub fn comment(&self) -> RateLimitedGuard {
|
|
|
|
|
self.kind(RateLimitType::Comment)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
pub fn search(&self) -> RateLimited {
|
|
|
|
|
pub fn search(&self) -> RateLimitedGuard {
|
|
|
|
|
self.kind(RateLimitType::Search)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn kind(&self, type_: RateLimitType) -> RateLimited {
|
|
|
|
|
RateLimited {
|
|
|
|
|
rate_limiter: self.rate_limiter.clone(),
|
|
|
|
|
rate_limit_config: self.rate_limit_config.clone(),
|
|
|
|
|
fn kind(&self, type_: RateLimitType) -> RateLimitedGuard {
|
|
|
|
|
RateLimitedGuard {
|
|
|
|
|
rate_limit: self.rate_limit.clone(),
|
|
|
|
|
type_,
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl RateLimited {
|
|
|
|
|
pub struct RateLimitedMiddleware<S> {
|
|
|
|
|
rate_limited: RateLimitedGuard,
|
|
|
|
|
service: Rc<S>,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl RateLimitedGuard {
|
|
|
|
|
/// Returns true if the request passed the rate limit, false if it failed and should be rejected.
|
|
|
|
|
pub fn check(self, ip_addr: IpAddr) -> bool {
|
|
|
|
|
// Does not need to be blocking because the RwLock in settings never held across await points,
|
|
|
|
|
// and the operation here locks only long enough to clone
|
|
|
|
|
let rate_limit = self.rate_limit_config;
|
|
|
|
|
let mut guard = self
|
|
|
|
|
.rate_limit
|
|
|
|
|
.lock()
|
|
|
|
|
.expect("Failed to lock rate limit mutex for reading");
|
|
|
|
|
let rate_limit = &guard.rate_limit_config;
|
|
|
|
|
|
|
|
|
|
let (kind, interval) = match self.type_ {
|
|
|
|
|
RateLimitType::Message => (rate_limit.message, rate_limit.message_per_second),
|
|
|
|
@ -123,13 +161,13 @@ impl RateLimited {
|
|
|
|
|
RateLimitType::Comment => (rate_limit.comment, rate_limit.comment_per_second),
|
|
|
|
|
RateLimitType::Search => (rate_limit.search, rate_limit.search_per_second),
|
|
|
|
|
};
|
|
|
|
|
let mut limiter = self.rate_limiter.lock().expect("mutex poison error");
|
|
|
|
|
let limiter = &mut guard.rate_limiter;
|
|
|
|
|
|
|
|
|
|
limiter.check_rate_limit_full(self.type_, &ip_addr, kind, interval)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl<S> Transform<S, ServiceRequest> for RateLimited
|
|
|
|
|
impl<S> Transform<S, ServiceRequest> for RateLimitedGuard
|
|
|
|
|
where
|
|
|
|
|
S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
|
|
|
|
|
S::Future: 'static,
|
|
|
|
|