LemmyNet/lemmy
2
0
Fork 0
mirror of https://github.com/LemmyNet/lemmy.git synced 2024-11-26 22:31:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Give ratelimit fields explicit names

Browse source
This commit is contained in:
asonix 2020-04-20 12:51:42 -05:00
parent 904ebf34a5
commit 4df2031ae0
2 changed files with 36 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
server/src/main.rs
View file

@ -34,7 +34,9 @@ async fn main() -> io::Result<()> {
embedded_migrations::run(&conn).unwrap(); embedded_migrations::run(&conn).unwrap();
// Set up the rate limiter // Set up the rate limiter
let rate_limiter = RateLimit(Arc::new(Mutex::new(RateLimiter::default()))); let rate_limiter = RateLimit {
rate_limiter: Arc::new(Mutex::new(RateLimiter::default())),
};
// Set up websocket server // Set up websocket server
let server = ChatServer::startup(pool.clone(), rate_limiter.clone()).start(); let server = ChatServer::startup(pool.clone(), rate_limiter.clone()).start();

49
server/src/rate_limit/mod.rs
View file

@ -18,12 +18,20 @@ use strum::IntoEnumIterator;
use tokio::sync::Mutex; use tokio::sync::Mutex;
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct RateLimit(pub Arc<Mutex<RateLimiter>>); pub struct RateLimit {
pub rate_limiter: Arc<Mutex<RateLimiter>>,
}
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct RateLimited(Arc<Mutex<RateLimiter>>, RateLimitType); pub struct RateLimited {
rate_limiter: Arc<Mutex<RateLimiter>>,
type_: RateLimitType,
}
pub struct RateLimitedMiddleware<S>(RateLimited, S); pub struct RateLimitedMiddleware<S> {
rate_limited: RateLimited,
service: S,
}
impl RateLimit { impl RateLimit {
pub fn message(&self) -> RateLimited { pub fn message(&self) -> RateLimited {
@ -39,7 +47,10 @@ impl RateLimit {
} }
fn kind(&self, type_: RateLimitType) -> RateLimited { fn kind(&self, type_: RateLimitType) -> RateLimited {
RateLimited(self.0.clone(), type_) RateLimited {
rate_limiter: self.rate_limiter.clone(),
type_,
}
} }
} }
@ -64,12 +75,12 @@ impl RateLimited {
// before // before
{ {
let mut limiter = self.0.lock().await; let mut limiter = self.rate_limiter.lock().await;
match self.1 { match self.type_ {
RateLimitType::Message => { RateLimitType::Message => {
limiter.check_rate_limit_full( limiter.check_rate_limit_full(
self.1, self.type_,
&ip_addr, &ip_addr,
rate_limit.message, rate_limit.message,
rate_limit.message_per_second, rate_limit.message_per_second,
@ -80,7 +91,7 @@ impl RateLimited {
} }
RateLimitType::Post => { RateLimitType::Post => {
limiter.check_rate_limit_full( limiter.check_rate_limit_full(
self.1.clone(), self.type_.clone(),
&ip_addr, &ip_addr,
rate_limit.post, rate_limit.post,
rate_limit.post_per_second, rate_limit.post_per_second,
@ -89,7 +100,7 @@ impl RateLimited {
} }
RateLimitType::Register => { RateLimitType::Register => {
limiter.check_rate_limit_full( limiter.check_rate_limit_full(
self.1, self.type_,
&ip_addr, &ip_addr,
rate_limit.register, rate_limit.register,
rate_limit.register_per_second, rate_limit.register_per_second,
@ -103,12 +114,12 @@ impl RateLimited {
// after // after
{ {
let mut limiter = self.0.lock().await; let mut limiter = self.rate_limiter.lock().await;
if res.is_ok() { if res.is_ok() {
match self.1 { match self.type_ {
RateLimitType::Post => { RateLimitType::Post => {
limiter.check_rate_limit_full( limiter.check_rate_limit_full(
self.1, self.type_,
&ip_addr, &ip_addr,
rate_limit.post, rate_limit.post,
rate_limit.post_per_second, rate_limit.post_per_second,
@ -117,7 +128,7 @@ impl RateLimited {
} }
RateLimitType::Register => { RateLimitType::Register => {
limiter.check_rate_limit_full( limiter.check_rate_limit_full(
self.1, self.type_,
&ip_addr, &ip_addr,
rate_limit.register, rate_limit.register,
rate_limit.register_per_second, rate_limit.register_per_second,
@ -146,7 +157,10 @@ where
type Future = Ready<Result<Self::Transform, Self::InitError>>; type Future = Ready<Result<Self::Transform, Self::InitError>>;
fn new_transform(&self, service: S) -> Self::Future { fn new_transform(&self, service: S) -> Self::Future {
ok(RateLimitedMiddleware(self.clone(), service)) ok(RateLimitedMiddleware {
rate_limited: self.clone(),
service,
})
} }
} }
@ -163,7 +177,7 @@ where
type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>; type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.1.poll_ready(cx) self.service.poll_ready(cx)
} }
fn call(&mut self, req: S::Request) -> Self::Future { fn call(&mut self, req: S::Request) -> Self::Future {
@ -176,7 +190,10 @@ where
.unwrap_or("127.0.0.1") .unwrap_or("127.0.0.1")
.to_string(); .to_string();
let fut = self.0.clone().wrap(ip_addr, self.1.call(req)); let fut = self
.rate_limited
.clone()
.wrap(ip_addr, self.service.call(req));
Box::pin(async move { fut.await.map_err(actix_web::Error::from) }) Box::pin(async move { fut.await.map_err(actix_web::Error::from) })
} }