better rate limiting (#1296)

* better rate limiting

* actually reset counter
This commit is contained in:
Aiden McClelland
2022-02-28 13:06:27 -07:00
committed by GitHub
parent b57f12f7ae
commit 86823e063a

View File

@@ -181,7 +181,7 @@ impl Borrow<String> for HashSessionToken {
} }
pub fn auth<M: Metadata>(ctx: RpcContext) -> DynMiddleware<M> { pub fn auth<M: Metadata>(ctx: RpcContext) -> DynMiddleware<M> {
let rate_limiter = Arc::new(Mutex::new(Instant::now())); let rate_limiter = Arc::new(Mutex::new((0_usize, Instant::now())));
Box::new( Box::new(
move |req: &mut Request<Body>, move |req: &mut Request<Body>,
metadata: M| metadata: M|
@@ -205,24 +205,29 @@ pub fn auth<M: Metadata>(ctx: RpcContext) -> DynMiddleware<M> {
Err(e.into()), Err(e.into()),
|_| StatusCode::OK, |_| StatusCode::OK,
)?)); )?));
} else { } else if rpc_req.method.as_str() == "auth.login" {
let mut guard = rate_limiter.lock().await; let mut guard = rate_limiter.lock().await;
if guard.elapsed() < Duration::from_secs(10) { guard.0 += 1;
let (res_parts, _) = Response::new(()).into_parts(); if guard.1.elapsed() < Duration::from_secs(20) {
return Ok(Err(to_response( if guard.0 >= 3 {
&req.headers, let (res_parts, _) = Response::new(()).into_parts();
res_parts, return Ok(Err(to_response(
Err(Error::new( &req.headers,
eyre!( res_parts,
"Please limit login attempts to 1 per 10 seconds." Err(Error::new(
eyre!(
"Please limit login attempts to 3 per 20 seconds."
), ),
crate::ErrorKind::RateLimited, crate::ErrorKind::RateLimited,
) )
.into()), .into()),
|_| StatusCode::OK, |_| StatusCode::OK,
)?)); )?));
}
} else {
guard.0 = 0;
} }
*guard = Instant::now(); guard.1 = Instant::now();
} }
} }
Ok(Ok(noop3())) Ok(Ok(noop3()))