native-tls: fix buffered io hang

tokio-native-tls/src/lib.rs

@@ -360,20 +360,45 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Future for MidHandshake<S> {
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let mut_self = self.get_mut();
-        let mut s = mut_self.0.take().expect("future polled after completion");
-
-        s.get_mut().context = cx as *mut _ as *mut ();
-        match s.handshake() {
-            Ok(mut s) => {
-                s.get_mut().context = null_mut();
-                Poll::Ready(Ok(TlsStream(s)))
+        let mut already_flushed = false;
+
+        loop {
+            let mut s = mut_self.0.take().expect("future polled after completion");
+
+            s.get_mut().context = cx as *mut _ as *mut ();
+
+            match s.handshake() {
+                Ok(mut s) => {
+                    s.get_mut().context = null_mut();
+                    return Poll::Ready(Ok(TlsStream(s)));
+                }
+                Err(HandshakeError::WouldBlock(mut s)) => {
+                    s.get_mut().context = null_mut();
+
+                    match Pin::new(&mut s.get_mut().get_mut()).poll_flush(cx) {
+                        Poll::Ready(Ok(())) => {
+                            mut_self.0 = Some(s);
+
+                            if !already_flushed {
+                                already_flushed = true;
+                                continue;
+                            } else {
+                                return Poll::Pending;
+                            }
+                        }
+                        Poll::Ready(Err(_e)) => {
+                            mut_self.0 = Some(s);
+                            todo!("figure out how to bubble up the io error")
+                        }
+                        Poll::Pending => {}
+                    }
+
+                    mut_self.0 = Some(s);
+
+                    return Poll::Pending;
+                }
+                Err(HandshakeError::Failure(e)) => return Poll::Ready(Err(e)),
             }
-            Err(HandshakeError::WouldBlock(mut s)) => {
-                s.get_mut().context = null_mut();
-                mut_self.0 = Some(s);
-                Poll::Pending
-            }
-            Err(HandshakeError::Failure(e)) => Poll::Ready(Err(e)),
         }
     }
 }

tokio-native-tls/tests/buffered.rs

@@ -0,0 +1,22 @@
+use native_tls::TlsConnector;
+use tokio::{io::BufWriter, net::TcpStream};
+
+#[tokio::test]
+async fn connect_using_bufwriter() {
+    drop(env_logger::try_init());
+
+    let socket = BufWriter::new(
+        TcpStream::connect(("example.com", 443))
+            .await
+            .expect("connect socket"),
+    );
+
+    let connector = TlsConnector::builder()
+        .build()
+        .expect("build TLS connector");
+
+    tokio_native_tls::TlsConnector::from(connector)
+        .connect("example.com", BufWriter::new(socket))
+        .await
+        .expect("connect TLS");
+}