From 1f98d87a620cf23b09fb851e8154fda2f3e4ccde Mon Sep 17 00:00:00 2001
From: quininer <quininer@live.com>
Date: Mon, 24 Sep 2018 23:02:30 +0800
Subject: [PATCH] more complete handshake

---
 src/common/mod.rs |  5 +++--
 src/lib.rs        |  4 +++-
 src/tokio_impl.rs | 39 +++++++++++++++++++++++++--------------
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/src/common/mod.rs b/src/common/mod.rs
index 7db198e..e18aa6e 100644
--- a/src/common/mod.rs
+++ b/src/common/mod.rs
@@ -12,9 +12,10 @@ use rustls::WriteV;
 #[cfg(feature = "tokio-support")]
 use tokio::io::AsyncWrite;
 
+
 pub struct Stream<'a, S: 'a, IO: 'a> {
-    session: &'a mut S,
-    io: &'a mut IO
+    pub session: &'a mut S,
+    pub io: &'a mut IO
 }
 
 pub trait WriteTls<'a, S: Session, IO: Read + Write>: Read + Write {
diff --git a/src/lib.rs b/src/lib.rs
index 736cad6..ee1524d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -129,9 +129,11 @@ impl<IO, S> TlsStream<IO, S> {
     }
 }
 
-impl<IO, S> From<(IO, S)> for TlsStream<IO, S> {
+impl<IO, S: Session> From<(IO, S)> for TlsStream<IO, S> {
     #[inline]
     fn from((io, session): (IO, S)) -> TlsStream<IO, S> {
+        assert!(!session.is_handshaking());
+
         TlsStream {
             is_shutdown: false,
             eof: false,
diff --git a/src/tokio_impl.rs b/src/tokio_impl.rs
index 00b4722..644e4f0 100644
--- a/src/tokio_impl.rs
+++ b/src/tokio_impl.rs
@@ -5,6 +5,17 @@ use tokio::prelude::Poll;
 use common::Stream;
 
 
+macro_rules! try_async {
+    ( $e:expr ) => {
+        match $e {
+            Ok(n) => n,
+            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock =>
+                return Ok(Async::NotReady),
+            Err(e) => return Err(e)
+        }
+    }
+}
+
 impl<IO: AsyncRead + AsyncWrite> Future for Connect<IO> {
     type Item = TlsStream<IO, ClientSession>;
     type Error = io::Error;
@@ -24,7 +35,9 @@ impl<IO: AsyncRead + AsyncWrite> Future for Accept<IO> {
 }
 
 impl<IO, S> Future for MidHandshake<IO, S>
-    where IO: io::Read + io::Write, S: Session
+where
+    IO: io::Read + io::Write,
+    S: Session
 {
     type Item = TlsStream<IO, S>;
     type Error = io::Error;
@@ -32,15 +45,15 @@ impl<IO, S> Future for MidHandshake<IO, S>
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
         {
             let stream = self.inner.as_mut().unwrap();
-            if stream.session.is_handshaking() {
-                let (io, session) = stream.get_mut();
-                let mut stream = Stream::new(session, io);
+            let (io, session) = stream.get_mut();
+            let mut stream = Stream::new(session, io);
 
-                match stream.complete_io() {
-                    Ok(_) => (),
-                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(Async::NotReady),
-                    Err(e) => return Err(e)
-                }
+            if stream.session.is_handshaking() {
+                try_async!(stream.complete_io());
+            }
+
+            if stream.session.wants_write() {
+                try_async!(stream.complete_io());
             }
         }
 
@@ -69,12 +82,10 @@ impl<IO, S> AsyncWrite for TlsStream<IO, S>
             self.is_shutdown = true;
         }
 
-        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)
+        {
+            let mut stream = Stream::new(&mut self.session, &mut self.io);
+            try_async!(stream.complete_io());
         }
-
         self.io.shutdown()
     }
 }