From bab1fe8272e0997a7103a8c7e750541b0dbf3cf2 Mon Sep 17 00:00:00 2001
From: Pierre HUBERT <pierre.git@communiquons.org>
Date: Sat, 20 Jun 2020 10:47:06 +0200
Subject: [PATCH] Apply stream read limit

---
 src/constants.rs          |  2 +-
 src/controllers/server.rs | 23 +++++++++++++++++------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/constants.rs b/src/constants.rs
index 3c2f767..9a2a351 100644
--- a/src/constants.rs
+++ b/src/constants.rs
@@ -40,4 +40,4 @@ pub const DEFAULT_ACCOUNT_IMAGE: &str = "avatars/0Reverse.png";
 pub const ERROR_ACCOUNT_IMAGE: &str = "avatars/0Red.png";
 
 /// Maximum requests size (50 Mo)
-pub const MAX_REQUEST_SIZE: u64 = 50000000;
\ No newline at end of file
+pub const MAX_REQUEST_SIZE: usize = 50000000;
\ No newline at end of file
diff --git a/src/controllers/server.rs b/src/controllers/server.rs
index 6404df9..5a07c93 100644
--- a/src/controllers/server.rs
+++ b/src/controllers/server.rs
@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::io::Take;
 use std::pin::Pin;
 
 use actix_web::{App, FromRequest, http, HttpMessage, HttpRequest, HttpResponse, HttpServer, web};
@@ -23,9 +22,10 @@ use crate::data::http_request_handler::{HttpRequestHandler, RequestValue};
 ///
 /// @author Pierre Hubert
 
-/// Custom stream to give it a limit
+/// Custom stream to give a limit to requests size
 struct LimitedStream {
-    stream: Box<dyn Stream<Item=Result<Bytes, PayloadError>> + Unpin + 'static>
+    stream: Box<dyn Stream<Item=Result<Bytes, PayloadError>> + Unpin + 'static>,
+    already_read: usize,
 }
 
 impl<'a> Stream for LimitedStream
@@ -33,7 +33,17 @@ 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>> {
-        Pin::new(self.stream.as_mut()).poll_next(cx)
+        if self.already_read >= MAX_REQUEST_SIZE { // TODO : check if works
+            return Poll::Ready(None);
+        }
+
+        let res = Pin::new(self.stream.as_mut()).poll_next(cx);
+
+        if let Poll::Ready(Some(Ok(d))) = &res {
+            self.already_read = self.already_read + d.len();
+        }
+
+        res
     }
 }
 
@@ -54,7 +64,8 @@ impl FromRequest for CustomRequest {
         let payload = Box::new(payload.take());
 
         let payload = LimitedStream {
-            stream: payload
+            stream: payload,
+            already_read: 0,
         };
 
 
@@ -63,7 +74,7 @@ impl FromRequest for CustomRequest {
             // 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::<u64>().unwrap_or(0) > MAX_REQUEST_SIZE {
+                    if String::from_utf8_lossy(v.as_bytes()).parse::<usize>().unwrap_or(0) > MAX_REQUEST_SIZE {
                         return Err(actix_web::error::ErrorBadRequest("Request too big!"));
                     }
                 }