Start to wrap incoming stream

src/controllers/server.rs

@@ -1,24 +1,42 @@
-use actix_web::{App, HttpResponse, HttpServer, web, http, FromRequest, HttpRequest, HttpMessage};
-
-use crate::controllers::routes::{get_routes, Route, RequestResult};
-use crate::data::config::Config;
-use crate::controllers::routes::Method::{GET, POST};
-use crate::data::http_request_handler::{HttpRequestHandler, RequestValue};
-use actix_web::dev::{PayloadStream, Payload, Decompress};
-use actix_web::web::{BytesMut};
-use futures::future::{LocalBoxFuture};
-use futures::{FutureExt, StreamExt};
-use actix_web::error::{ErrorBadRequest};
-use encoding_rs::UTF_8;
 use std::collections::HashMap;
-use percent_encoding::{percent_decode_str};
-use crate::api_data::http_error::HttpError;
+use std::io::Take;
+use std::pin::Pin;
 
+use actix_web::{App, FromRequest, http, HttpMessage, HttpRequest, HttpResponse, HttpServer, web};
+use actix_web::dev::{Decompress, Payload, PayloadStream};
+use actix_web::error::{ErrorBadRequest, PayloadError};
+use actix_web::web::{Bytes, BytesMut};
+use encoding_rs::UTF_8;
+use futures::{FutureExt, Stream, StreamExt};
+use futures::future::LocalBoxFuture;
+use futures::task::{Context, Poll};
+use percent_encoding::percent_decode_str;
+
+use crate::api_data::http_error::HttpError;
+use crate::constants::MAX_REQUEST_SIZE;
+use crate::controllers::routes::{get_routes, RequestResult, Route};
+use crate::controllers::routes::Method::{GET, POST};
+use crate::data::config::Config;
+use crate::data::http_request_handler::{HttpRequestHandler, RequestValue};
 
 /// Main server functions
 ///
 /// @author Pierre Hubert
 
+/// Custom stream to give it a limit
+struct LimitedStream {
+    stream: Box<dyn Stream<Item=Result<Bytes, PayloadError>> + Unpin + 'static>
+}
+
+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)
+    }
+}
+
 /// Custom request value
 struct CustomRequest {
     req: web::HttpRequest,
@@ -33,9 +51,25 @@ impl FromRequest for CustomRequest {
 
     fn from_request(req: &HttpRequest, payload: &mut Payload<PayloadStream>) -> Self::Future {
         let req = req.clone();
-        let payload = payload.take();
+        let payload = Box::new(payload.take());
+
+        let payload = LimitedStream {
+            stream: payload
+        };
+
 
         async move {
+
+            // 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 {
+                        return Err(actix_web::error::ErrorBadRequest("Request too big!"));
+                    }
+                }
+            }
+
+
             let mut body_args = HashMap::new();
 
             // Process "application/x-www-form-urlencoded" requests
@@ -88,12 +122,21 @@ impl FromRequest for CustomRequest {
                         // Add the value to the body
                         body_args.insert(
                             percent_decode_str(args[0]).decode_utf8_lossy().to_string(),
-                            RequestValue::string(percent_decode_str(args[1]).decode_utf8_lossy().to_string())
+                            RequestValue::string(percent_decode_str(args[1]).decode_utf8_lossy().to_string()),
                         );
                     }
                 }
             }
 
+            // Process "multipart/form-data" request
+            else if req.content_type().starts_with("multipart/form-data") {
+                let mut req = actix_multipart::Multipart::new(req.headers(), payload);
+
+                while let Some(el) = req.next().await {
+                    let mut field = el?;
+                }
+            }
+
             Ok(CustomRequest {
                 req: req.clone(),
                 body: body_args,
@@ -153,8 +196,7 @@ async fn process_request(custom_req: CustomRequest) -> HttpResponse {
 
         // Set default error response if required
         Err(e) => {
-
-            let err_msg= e.to_string();
+            let err_msg = e.to_string();
 
             if !request.has_response() {
                 request.internal_error(e).unwrap_err();