Start9/start-os
2
0
Fork 0
mirror of https://github.com/Start9Labs/start-os.git synced 2026-03-31 04:23:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feature/rate limiting (#786)

Browse Source 
* rate limiting

* 10s rate limit
This commit is contained in:
Aiden McClelland
2021-11-12 15:00:57 -07:00
parent c723ee6a15
commit 9e1e3e167b
1 changed files with 28 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
appmgr/src/middleware/auth.rs
View File

@@ -1,4 +1,6 @@
use std::borrow::Borrow; use std::borrow::Borrow;
use std::sync::Arc;
use std::time::{Duration, Instant};
use basic_cookies::Cookie; use basic_cookies::Cookie;
use color_eyre::eyre::eyre; use color_eyre::eyre::eyre;
@@ -15,6 +17,7 @@ use rpc_toolkit::yajrc::RpcMethod;
use rpc_toolkit::Metadata; use rpc_toolkit::Metadata;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sha2::Sha256; use sha2::Sha256;
use tokio::sync::Mutex;
use crate::context::RpcContext; use crate::context::RpcContext;
use crate::{Error, ResultExt}; use crate::{Error, ResultExt};
@@ -178,21 +181,23 @@ 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()));
Box::new( Box::new(
move |req: &mut Request<Body>, move |req: &mut Request<Body>,
metadata: M| metadata: M|
-> BoxFuture<Result<Result<DynMiddlewareStage2, Response<Body>>, HttpError>> { -> BoxFuture<Result<Result<DynMiddlewareStage2, Response<Body>>, HttpError>> {
let ctx = ctx.clone(); let ctx = ctx.clone();
let rate_limiter = rate_limiter.clone();
async move { async move {
let mut header_stub = Request::new(Body::empty()); let mut header_stub = Request::new(Body::empty());
*header_stub.headers_mut() = req.headers().clone(); *header_stub.headers_mut() = req.headers().clone();
let m2: DynMiddlewareStage2 = Box::new(move |req, rpc_req| { let m2: DynMiddlewareStage2 = Box::new(move |req, rpc_req| {
async move { async move {
if metadata if let Err(e) = HasValidSession::from_request_parts(req, &ctx).await {
.get(rpc_req.method.as_str(), "authenticated") if metadata
.unwrap_or(true) .get(rpc_req.method.as_str(), "authenticated")
{ .unwrap_or(true)
if let Err(e) = HasValidSession::from_request_parts(req, &ctx).await { {
let (res_parts, _) = Response::new(()).into_parts(); let (res_parts, _) = Response::new(()).into_parts();
return Ok(Err(to_response( return Ok(Err(to_response(
&req.headers, &req.headers,
@@ -200,6 +205,24 @@ pub fn auth<M: Metadata>(ctx: RpcContext) -> DynMiddleware<M> {
Err(e.into()), Err(e.into()),
|_| StatusCode::OK, |_| 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())) Ok(Ok(noop3()))