From 64ba67b1ca8b15606ab404df27f0715ca9a0d042 Mon Sep 17 00:00:00 2001
From: Martin Algesten <martin@algesten.se>
Date: Wed, 31 Oct 2018 14:24:25 +0100
Subject: [PATCH] Expose chunked_threshold on Response. Fix #149

---
 src/response.rs  | 31 +++++++++++++++++++++++++++----
 tests/network.rs |  8 ++++++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/response.rs b/src/response.rs
index c40c12fec..33f54d824 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -53,6 +53,7 @@ pub struct Response<R> where R: Read {
     status_code: StatusCode,
     headers: Vec<Header>,
     data_length: Option<usize>,
+    chunked_threshold: Option<usize>
 }
 
 /// A `Response` without a template parameter.
@@ -113,7 +114,8 @@ fn write_message_header<W>(mut writer: W, http_version: &HTTPVersion,
 }
 
 fn choose_transfer_encoding(request_headers: &[Header], http_version: &HTTPVersion,
-                            entity_length: &Option<usize>, has_additional_headers: bool)
+                            entity_length: &Option<usize>, has_additional_headers: bool,
+                            chunked_threshold: usize)
     -> TransferEncoding
 {
     use util;
@@ -165,8 +167,7 @@ fn choose_transfer_encoding(request_headers: &[Header], http_version: &HTTPVersi
     }
 
     // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
-    let chunks_threshold = 32768;
-    if entity_length.as_ref().map_or(true, |val| *val >= chunks_threshold) {
+    if entity_length.as_ref().map_or(true, |val| *val >= chunked_threshold) {
         return TransferEncoding::Chunked;
     }
 
@@ -191,6 +192,7 @@ impl<R> Response<R> where R: Read {
             status_code: status_code,
             headers: Vec::with_capacity(16),
             data_length: data_length,
+            chunked_threshold: None,
         };
 
         for h in headers {
@@ -207,6 +209,23 @@ impl<R> Response<R> where R: Read {
         response
     }
 
+    /// Set a threshold for `Content-Length` where we chose chunked
+    /// transfer. Notice that chunked transfer might happen regardless of
+    /// this threshold, for instance when the request headers indicate
+    /// it is wanted or when there is no `Content-Length`.
+    pub fn with_chunked_threshold(mut self, length: usize) -> Response<R>{
+        self.chunked_threshold = Some(length);
+        self
+    }
+
+    /// The current `Content-Length` threshold for switching over to
+    /// chunked transfer. The default is 32768 bytes. Notice that
+    /// chunked transfer is mutually exclusive with sending a
+    /// `Content-Length` header as per the HTTP spec.
+    pub fn chunked_threshold(&self) -> usize {
+        self.chunked_threshold.unwrap_or(32768)
+    }
+
     /// Adds a header to the list.
     /// Does all the checks.
     pub fn add_header<H>(&mut self, header: H) where H: Into<Header> {
@@ -260,6 +279,7 @@ impl<R> Response<R> where R: Read {
             headers: self.headers,
             status_code: self.status_code,
             data_length: data_length,
+            chunked_threshold: None,
         }
     }
 
@@ -278,7 +298,8 @@ impl<R> Response<R> where R: Read {
                                -> IoResult<()>
     {
         let mut transfer_encoding = Some(choose_transfer_encoding(request_headers,
-                                    &http_version, &self.data_length, false /* TODO */));
+                                    &http_version, &self.data_length, false /* TODO */,
+                                    self.chunked_threshold()));
 
         // add `Date` if not in the headers
         if self.headers.iter().find(|h| h.field.equiv(&"Date")).is_none() {
@@ -386,6 +407,7 @@ impl<R> Response<R> where R: Read + Send + 'static {
             status_code: self.status_code,
             headers: self.headers,
             data_length: self.data_length,
+            chunked_threshold: None,
         }
     }
 }
@@ -463,6 +485,7 @@ impl Clone for Response<io::Empty> {
             status_code: self.status_code.clone(),
             headers: self.headers.clone(),
             data_length: self.data_length.clone(),
+            chunked_threshold: self.chunked_threshold.clone(),
         }
     }
 }
diff --git a/tests/network.rs b/tests/network.rs
index 0d6260967..cdae914b7 100644
--- a/tests/network.rs
+++ b/tests/network.rs
@@ -173,3 +173,11 @@ fn connection_timeout() {
     tx_stop.send(());
 }
 */
+
+
+#[test]
+fn chunked_threshold() {
+    let resp = tiny_http::Response::from_string("test".to_string());
+    assert_eq!(resp.chunked_threshold(), 32768);
+    assert_eq!(resp.with_chunked_threshold(42).chunked_threshold(), 42);
+}