From 72c9c1d59ec6a2dea0735e41dd0ad1aa12baf24d Mon Sep 17 00:00:00 2001
From: PZ Read <netfirewall@gmail.com>
Date: Fri, 26 May 2017 15:54:47 +0800
Subject: [PATCH 1/3] Fix plaintext write logic for limited rustls buffer.

---
 src/lib.rs | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 1a48ddc..8728219 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -209,17 +209,29 @@ impl<S, C> io::Write for TlsStream<S, C>
     where S: AsyncRead + AsyncWrite, C: Session
 {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        let output = self.session.write(buf)?;
+        loop {
+            let output = self.session.write(buf)?;
 
-        while self.session.wants_write() {
-            match self.session.write_tls(&mut self.io) {
-                Ok(_) => (),
-                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => break,
-                Err(e) => return Err(e)
+            while self.session.wants_write() {
+                match self.session.write_tls(&mut self.io) {
+                    Ok(_) => (),
+                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                        if output == 0 {
+                            // Both rustls buffer and IO buffer are blocking.
+                            return Err(io::Error::from(io::ErrorKind::WouldBlock));
+                        } else {
+                            break;
+                        }
+                    }
+                    Err(e) => return Err(e)
+                }
+            }
+
+            if output > 0 {
+                // Already wrote something out.
+                return Ok(output);
             }
         }
-
-        Ok(output)
     }
 
     fn flush(&mut self) -> io::Result<()> {

From 185f01093783b1356d3d135eaa1475b3584b723e Mon Sep 17 00:00:00 2001
From: PZ Read <netfirewall@gmail.com>
Date: Fri, 26 May 2017 16:38:02 +0800
Subject: [PATCH 2/3] Add async builders for custom session.

---
 src/lib.rs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 8728219..87f461f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -52,6 +52,15 @@ impl ClientConfigExt for Arc<ClientConfig> {
     }
 }
 
+pub fn connect_async_with_session<S>(stream: S, session: ClientSession)
+    -> ConnectAsync<S>
+    where S: AsyncRead + AsyncWrite
+{
+    ConnectAsync(MidHandshake {
+        inner: Some(TlsStream::new(stream, session))
+    })
+}
+
 impl ServerConfigExt for Arc<ServerConfig> {
     fn accept_async<S>(&self, stream: S)
         -> AcceptAsync<S>
@@ -63,6 +72,15 @@ impl ServerConfigExt for Arc<ServerConfig> {
     }
 }
 
+pub fn accept_async_with_session<S>(stream: S, session: ServerSession)
+    -> AcceptAsync<S>
+    where S: AsyncRead + AsyncWrite
+{
+    AcceptAsync(MidHandshake {
+        inner: Some(TlsStream::new(stream, session))
+    })
+}
+
 impl<S: AsyncRead + AsyncWrite> Future for ConnectAsync<S> {
     type Item = TlsStream<S, ClientSession>;
     type Error = io::Error;

From d6d06041d9735ba64d09e845b184873e36056e41 Mon Sep 17 00:00:00 2001
From: PZ Read <netfirewall@gmail.com>
Date: Fri, 26 May 2017 17:59:51 +0800
Subject: [PATCH 3/3] Fix empty buffer

---
 src/lib.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 87f461f..2c6a7e9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -227,6 +227,10 @@ impl<S, C> io::Write for TlsStream<S, C>
     where S: AsyncRead + AsyncWrite, C: Session
 {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        if buf.len() == 0 {
+            return Ok(0);
+        }
+
         loop {
             let output = self.session.write(buf)?;