From fff2f4a73b19559deb4faa4dabf237cde353b2a6 Mon Sep 17 00:00:00 2001
From: quininer <quininer@live.com>
Date: Sat, 31 Mar 2018 15:16:31 +0800
Subject: [PATCH] fix(tokio_impl): shutdown WouldBlock

---
 Cargo.toml        |  4 ++--
 src/tokio_impl.rs |  8 +++++++-
 tests/test.rs     | 14 ++++++++------
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index bc6d23c..a7a1bb5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "tokio-rustls"
-version = "0.6.0-alpha"
+version = "0.6.0-alpha.1"
 authors = ["quininer kel <quininer@live.com>"]
 license = "MIT/Apache-2.0"
 repository = "https://github.com/quininer/tokio-rustls"
@@ -24,5 +24,5 @@ webpki = "0.18.0-alpha"
 tokio = "0.1"
 
 [features]
-default = [ "unstable-futures", "tokio" ]
+default = [ "tokio" ]
 unstable-futures = [ "futures", "tokio/unstable-futures" ]
diff --git a/src/tokio_impl.rs b/src/tokio_impl.rs
index f5d3c6c..936c14b 100644
--- a/src/tokio_impl.rs
+++ b/src/tokio_impl.rs
@@ -64,7 +64,13 @@ impl<S, C> AsyncWrite for TlsStream<S, C>
             self.session.send_close_notify();
             self.is_shutdown = true;
         }
-        self.session.complete_io(&mut self.io)?;
+
+        match self.session.complete_io(&mut self.io) {
+            Ok(_) => (),
+            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(Async::NotReady),
+            Err(e) => return Err(e)
+        }
+
         self.io.shutdown()
     }
 }
diff --git a/tests/test.rs b/tests/test.rs
index 92c904a..246e85a 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -73,10 +73,11 @@ fn start_client(addr: &SocketAddr, domain: &str, chain: Option<BufReader<Cursor<
         .and_then(|stream| config.connect_async(domain, stream))
         .and_then(|stream| aio::write_all(stream, HELLO_WORLD))
         .and_then(|(stream, _)| aio::read_exact(stream, vec![0; HELLO_WORLD.len()]))
-        .and_then(|(_, buf)| {
+        .and_then(|(stream, buf)| {
             assert_eq!(buf, HELLO_WORLD);
-            Ok(())
-        });
+            aio::shutdown(stream)
+        })
+        .map(drop);
 
     done.wait()
 }
@@ -98,10 +99,11 @@ fn start_client2(addr: &SocketAddr, domain: &str, chain: Option<BufReader<Cursor
         .and_then(|stream| config.connect_async(domain, stream))
         .and_then(|stream| stream.write_all(HELLO_WORLD))
         .and_then(|(stream, _)| stream.read_exact(vec![0; HELLO_WORLD.len()]))
-        .and_then(|(_, buf)| {
+        .and_then(|(stream, buf)| {
             assert_eq!(buf, HELLO_WORLD);
-            Ok(())
-        });
+            stream.close()
+        })
+        .map(drop);
 
     block_on(done)
 }