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] 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<()> {