Merge pull request #10 from furakus/master

Allowing to use custom sessions to build async connector/acceptor

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;
@@ -209,17 +227,33 @@ 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)?;
-
-        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)
-            }
+        if buf.len() == 0 {
+            return Ok(0);
         }
 
-        Ok(output)
+        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 => {
+                        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);
+            }
+        }
     }
 
     fn flush(&mut self) -> io::Result<()> {