Improve content-length header check

src/controllers/server.rs

@@ -28,6 +28,7 @@ use crate::helpers::requests_limit_helper;
 struct LimitedStream {
     stream: Box<dyn Stream<Item=Result<Bytes, PayloadError>> + 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<Bytes, PayloadError>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        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::<usize>().unwrap_or(0) > MAX_REQUEST_SIZE {
+                    content_length_size = String::from_utf8_lossy(v.as_bytes()).parse::<usize>().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();