From 9e1e3e167ba5d65c6fb62a4eb6ec16e9d6f5b1b7 Mon Sep 17 00:00:00 2001 From: Aiden McClelland <3732071+dr-bonez@users.noreply.github.com> Date: Fri, 12 Nov 2021 15:00:57 -0700 Subject: [PATCH] Feature/rate limiting (#786) * rate limiting * 10s rate limit --- appmgr/src/middleware/auth.rs | 33 ++++++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/appmgr/src/middleware/auth.rs b/appmgr/src/middleware/auth.rs index e04ca2c1d..c7a9b118a 100644 --- a/appmgr/src/middleware/auth.rs +++ b/appmgr/src/middleware/auth.rs @@ -1,4 +1,6 @@ use std::borrow::Borrow; +use std::sync::Arc; +use std::time::{Duration, Instant}; use basic_cookies::Cookie; use color_eyre::eyre::eyre; @@ -15,6 +17,7 @@ use rpc_toolkit::yajrc::RpcMethod; use rpc_toolkit::Metadata; use serde::{Deserialize, Serialize}; use sha2::Sha256; +use tokio::sync::Mutex; use crate::context::RpcContext; use crate::{Error, ResultExt}; @@ -178,21 +181,23 @@ impl Borrow for HashSessionToken { } pub fn auth(ctx: RpcContext) -> DynMiddleware { + let rate_limiter = Arc::new(Mutex::new(Instant::now())); Box::new( move |req: &mut Request, metadata: M| -> BoxFuture>, HttpError>> { let ctx = ctx.clone(); + let rate_limiter = rate_limiter.clone(); async move { let mut header_stub = Request::new(Body::empty()); *header_stub.headers_mut() = req.headers().clone(); let m2: DynMiddlewareStage2 = Box::new(move |req, rpc_req| { async move { - if metadata - .get(rpc_req.method.as_str(), "authenticated") - .unwrap_or(true) - { - if let Err(e) = HasValidSession::from_request_parts(req, &ctx).await { + if let Err(e) = HasValidSession::from_request_parts(req, &ctx).await { + if metadata + .get(rpc_req.method.as_str(), "authenticated") + .unwrap_or(true) + { let (res_parts, _) = Response::new(()).into_parts(); return Ok(Err(to_response( &req.headers, @@ -200,6 +205,24 @@ pub fn auth(ctx: RpcContext) -> DynMiddleware { Err(e.into()), |_| StatusCode::OK, )?)); + } else { + let mut guard = rate_limiter.lock().await; + if guard.elapsed() < Duration::from_secs(10) { + let (res_parts, _) = Response::new(()).into_parts(); + return Ok(Err(to_response( + &req.headers, + res_parts, + Err(Error::new( + eyre!( + "Please limit login attempts to 1 per 10 seconds." + ), + crate::ErrorKind::RateLimited, + ) + .into()), + |_| StatusCode::OK, + )?)); + } + *guard = Instant::now(); } } Ok(Ok(noop3()))