From 9739168d487ebc48974fb2ff7a5a8d4d963d837f Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Tue, 30 Jan 2018 12:44:14 -0800
Subject: [PATCH] fix limit usage for Request/Response Body future

---
 src/client/connect.rs  |  4 ++++
 src/client/mod.rs      |  4 ++--
 src/client/request.rs  |  2 +-
 src/client/response.rs |  7 ++++---
 src/httprequest.rs     | 43 ++++++++++++++++++++----------------------
 5 files changed, 31 insertions(+), 29 deletions(-)

diff --git a/src/client/connect.rs b/src/client/connect.rs
index 6d071813..6323649c 100644
--- a/src/client/connect.rs
+++ b/src/client/connect.rs
@@ -25,9 +25,12 @@ use server::IoStream;
 
 
 #[derive(Debug)]
+/// `Connect` type represents message that can be send to `ClientConnector`
+/// with connection request.
 pub struct Connect(pub Uri);
 
 impl Connect {
+    /// Create `Connect` message for specified `Uri`
     pub fn new<U>(uri: U) -> Result<Connect, HttpError> where Uri: HttpTryFrom<U> {
         Ok(Connect(Uri::try_from(uri).map_err(|e| e.into())?))
     }
@@ -38,6 +41,7 @@ impl ResponseType for Connect {
     type Error = ClientConnectorError;
 }
 
+/// A set of errors that can occur during connecting to a http host
 #[derive(Fail, Debug)]
 pub enum ClientConnectorError {
     /// Invalid url
diff --git a/src/client/mod.rs b/src/client/mod.rs
index ba5a88f9..645766f7 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -6,6 +6,6 @@ mod writer;
 
 pub(crate) use self::writer::HttpClientWriter;
 pub use self::request::{ClientRequest, ClientRequestBuilder};
-pub use self::response::ClientResponse;
-pub use self::parser::{HttpResponseParser, HttpResponseParserError};
+pub use self::response::{ClientResponse, JsonResponse};
+pub(crate) use self::parser::{HttpResponseParser, HttpResponseParserError};
 pub use self::connect::{Connect, Connection, ClientConnector, ClientConnectorError};
diff --git a/src/client/request.rs b/src/client/request.rs
index 99abad8e..980f1879 100644
--- a/src/client/request.rs
+++ b/src/client/request.rs
@@ -170,7 +170,7 @@ impl fmt::Debug for ClientRequest {
 }
 
 
-/// An HTTP client request builder
+/// An HTTP Client request builder
 ///
 /// This type can be used to construct an instance of `ClientRequest` through a
 /// builder-like pattern.
diff --git a/src/client/response.rs b/src/client/response.rs
index 820d0100..9ae45553 100644
--- a/src/client/response.rs
+++ b/src/client/response.rs
@@ -39,6 +39,7 @@ impl Default for ClientMessage {
     }
 }
 
+/// An HTTP Client response
 pub struct ClientResponse(Rc<UnsafeCell<ClientMessage>>);
 
 impl ClientResponse {
@@ -288,8 +289,8 @@ impl Future for ResponseBody {
         if let Some(resp) = self.resp.take() {
             if let Some(len) = resp.headers().get(header::CONTENT_LENGTH) {
                 if let Ok(s) = len.to_str() {
-                    if let Ok(len) = s.parse::<u64>() {
-                        if len > 262_144 {
+                    if let Ok(len) = s.parse::<usize>() {
+                        if len > self.limit {
                             return Err(PayloadError::Overflow);
                         }
                     } else {
@@ -321,7 +322,7 @@ impl Future for ResponseBody {
     }
 }
 
-/// Client response payload json parser that resolves to a deserialized `T` value.
+/// Client response json parser that resolves to a deserialized `T` value.
 ///
 /// Returns error:
 ///
diff --git a/src/httprequest.rs b/src/httprequest.rs
index e22dcb7e..75050233 100644
--- a/src/httprequest.rs
+++ b/src/httprequest.rs
@@ -700,36 +700,21 @@ pub struct RequestBody {
     pl: ReadAny,
     body: BytesMut,
     limit: usize,
-    error: Option<PayloadError>,
+    req: Option<HttpRequest<()>>,
 }
 
 impl RequestBody {
 
     /// Create `RequestBody` for request.
     pub fn from_request<S>(req: &HttpRequest<S>) -> RequestBody {
-        let mut body = RequestBody {
-            pl: req.payload().readany(),
+        let pl = req.payload().readany();
+        RequestBody {
+            pl: pl,
             body: BytesMut::new(),
             limit: 262_144,
-            error: None
-        };
-
-        if let Some(len) = req.headers().get(header::CONTENT_LENGTH) {
-            if let Ok(s) = len.to_str() {
-                if let Ok(len) = s.parse::<u64>() {
-                    if len > 262_144 {
-                        body.error = Some(PayloadError::Overflow);
-                    }
-                } else {
-                    body.error = Some(PayloadError::UnknownLength);
-                }
-            } else {
-                body.error = Some(PayloadError::UnknownLength);
-            }
+            req: Some(req.clone_without_state())
         }
-
-        body
-    }
+ }
 
     /// Change max size of payload. By default max size is 256Kb
     pub fn limit(mut self, limit: usize) -> Self {
@@ -743,8 +728,20 @@ impl Future for RequestBody {
     type Error = PayloadError;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        if let Some(err) = self.error.take() {
-            return Err(err)
+        if let Some(req) = self.req.take() {
+            if let Some(len) = req.headers().get(header::CONTENT_LENGTH) {
+                if let Ok(s) = len.to_str() {
+                    if let Ok(len) = s.parse::<usize>() {
+                        if len > self.limit {
+                            return Err(PayloadError::Overflow);
+                        }
+                    } else {
+                        return Err(PayloadError::UnknownLength);
+                    }
+                } else {
+                    return Err(PayloadError::UnknownLength);
+                }
+            }
         }
 
         loop {