From 2c9e62ad8cb08913c6da27f2b74720d032649993 Mon Sep 17 00:00:00 2001
From: Martin Algesten <martin@algesten.se>
Date: Tue, 23 Oct 2018 20:18:24 +0100
Subject: [PATCH] fix dealloc issues

---
 src/pool.rs     | 27 ++++++++++++++++-----------
 src/response.rs | 21 +++++++++++++++++++--
 src/stream.rs   | 10 ----------
 3 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/src/pool.rs b/src/pool.rs
index 715dbfa..5b2999f 100644
--- a/src/pool.rs
+++ b/src/pool.rs
@@ -89,19 +89,18 @@ impl<R: Read + Sized> PoolReturnRead<R> {
                 return;
             }
             let state = &mut unit.agent.lock().unwrap();
+            // bring back stream here to either go into pool or dealloc
+            let stream = unsafe { *Box::from_raw(self.stream) };
+            self.stream = ::std::ptr::null_mut();
             if let Some(agent) = state.as_mut() {
-                unsafe {
-                    let stream = *Box::from_raw(self.stream);
-                    self.stream = ::std::ptr::null_mut();
-                    if !stream.is_poolable() {
-                        // just let it deallocate
-                        return;
-                    }
-                    // insert back into pool
-                    let key = PoolKey::new(&unit.url);
-                    agent.pool().recycle.insert(key, stream);
+                if !stream.is_poolable() {
+                    // just let it deallocate
+                    return;
                 }
-            };
+                // insert back into pool
+                let key = PoolKey::new(&unit.url);
+                agent.pool().recycle.insert(key, stream);
+            }
         }
     }
 
@@ -124,3 +123,9 @@ impl<R: Read + Sized> Read for PoolReturnRead<R> {
         Ok(amount)
     }
 }
+
+impl<R: Read + Sized> Drop for PoolReturnRead<R> {
+    fn drop(&mut self) {
+        self.return_connection();
+    }
+}
\ No newline at end of file
diff --git a/src/response.rs b/src/response.rs
index 03eff69..52173d7 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -276,7 +276,7 @@ impl Response {
 
         let stream = Box::new(self.stream.expect("No reader in response?!"));
         let stream_ptr = Box::into_raw(stream);
-        let yolo = YoloRead { stream: stream_ptr };
+        let mut yolo = YoloRead { stream: stream_ptr, dealloc: false };
         let unit = self.unit;
 
         match (use_chunked, limit_bytes) {
@@ -290,7 +290,10 @@ impl Response {
                 stream_ptr,
                 LimitedRead::new(yolo, len),
             )),
-            (false, None) => Box::new(yolo),
+            (false, None) => {
+                yolo.dealloc = true; // dealloc when read drops.
+                Box::new(yolo)
+            },
         }
     }
 
@@ -521,6 +524,7 @@ fn read_next_line<R: Read>(reader: &mut R) -> IoResult<AsciiString> {
 /// *Internal API*
 struct YoloRead {
     stream: *mut Stream,
+    dealloc: bool, // whether we are to dealloc stream on drop
 }
 
 impl Read for YoloRead {
@@ -531,6 +535,9 @@ impl Read for YoloRead {
             }
             let amount = (*self.stream).read(buf)?;
             if amount == 0 {
+                if self.dealloc {
+                    let _stream = Box::from_raw(self.stream);
+                }
                 self.stream = ::std::ptr::null_mut();
             }
             Ok(amount)
@@ -538,6 +545,16 @@ impl Read for YoloRead {
     }
 }
 
+impl Drop for YoloRead {
+    fn drop(&mut self) {
+        if self.dealloc && !self.stream.is_null() {
+            unsafe {
+                let _stream = Box::from_raw(self.stream);
+            }
+        }
+    }
+}
+
 /// Limits a YoloRead to a content size (as set by a "Content-Length" header).
 ///
 /// *Internal API*
diff --git a/src/stream.rs b/src/stream.rs
index ae09308..b316ed5 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -18,16 +18,6 @@ pub enum Stream {
     Test(Box<dyn Read + Send>, Vec<u8>),
 }
 
-impl Drop for Stream {
-    fn drop(&mut self) {
-        match self {
-            #[cfg(feature = "tls")]
-            Stream::Https(s) => { s.shutdown().ok(); },
-            _ => (),
-        }
-    }
-}
-
 impl ::std::fmt::Debug for Stream {
     fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
         write!(