From cc3e14414d752bfa3b567ee5b0d2f1a79382a02b 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  | 30 ++++++++++++++++++++++++++----
 tests/network.rs |  9 +++++++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/response.rs b/src/response.rs
index c40c12fec..84834bb2e 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,22 @@ 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 set_chunked_threshold(&mut self, length: usize) {
+        self.chunked_threshold = Some(length);
+    }
+
+    /// 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 +278,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 +297,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,
+                                    self.chunked_threshold() /* TODO */));
 
         // add `Date` if not in the headers
         if self.headers.iter().find(|h| h.field.equiv(&"Date")).is_none() {
@@ -386,6 +406,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 +484,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..8afdfc730 100644
--- a/tests/network.rs
+++ b/tests/network.rs
@@ -173,3 +173,12 @@ fn connection_timeout() {
     tx_stop.send(());
 }
 */
+
+
+#[test]
+fn chunked_threshold() {
+    let mut resp = tiny_http::Response::from_string("test".to_string());
+    assert_eq!(resp.chunked_threshold(), 32768);
+    resp.set_chunked_threshold(42);
+    assert_eq!(resp.chunked_threshold(), 42);
+}