diff --git a/src/controllers/server.rs b/src/controllers/server.rs index 4e3a35b..665d4b4 100644 --- a/src/controllers/server.rs +++ b/src/controllers/server.rs @@ -28,6 +28,7 @@ use crate::helpers::requests_limit_helper; struct LimitedStream { stream: Box> + Unpin + 'static>, already_read: usize, + max_size: usize, } impl<'a> Stream for LimitedStream @@ -35,7 +36,7 @@ impl<'a> Stream for LimitedStream type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - if self.already_read >= MAX_REQUEST_SIZE { // TODO : check if works + if self.already_read >= self.max_size { return Poll::Ready(None); } @@ -65,23 +66,29 @@ impl FromRequest for CustomRequest { let req = req.clone(); let payload = Box::new(payload.take()); - let payload = LimitedStream { - stream: payload, - already_read: 0, - }; - - async move { + let content_length_size; // Check the size, if provided if req.headers().contains_key("Content-Length") { if let Some(v) = req.headers().get("Content-Length") { - if String::from_utf8_lossy(v.as_bytes()).parse::().unwrap_or(0) > MAX_REQUEST_SIZE { + content_length_size = String::from_utf8_lossy(v.as_bytes()).parse::().unwrap_or(0); + if content_length_size > MAX_REQUEST_SIZE { return Err(actix_web::error::ErrorBadRequest("Request too big!")); } + } else { + unreachable!(); } + } else { + return Err(actix_web::error::ErrorBadRequest("Content-Length header is required!")); } + let payload = LimitedStream { + stream: payload, + already_read: 0, + max_size: content_length_size, + }; + let mut body_args = HashMap::new();