From 614a5068fa4cd73a162c457985666a7ba77e0335 Mon Sep 17 00:00:00 2001
From: Danny Browning <bdbrowning2@gmail.com>
Date: Tue, 10 Sep 2019 09:42:30 -0600
Subject: [PATCH] Tokio-Tungstenite Async Changes

Changes required for tokio-tungstenite async implementation.
---
 src/client.rs           |  4 ++--
 src/handshake/client.rs |  4 ++--
 src/handshake/mod.rs    | 12 +++++++++++-
 src/handshake/server.rs |  4 ++--
 src/server.rs           |  8 ++++----
 5 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index cea3086..f5f34b4 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -179,7 +179,7 @@ pub fn client_with_config<'t, Stream, Req>(
     config: Option<WebSocketConfig>,
 ) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
 where
-    Stream: Read + Write,
+    Stream: Read + Write + Unpin,
     Req: Into<Request<'t>>,
 {
     ClientHandshake::start(stream, request.into(), config).handshake()
@@ -195,7 +195,7 @@ pub fn client<'t, Stream, Req>(
     stream: Stream,
 ) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
 where
-    Stream: Read + Write,
+    Stream: Read + Write + Unpin,
     Req: Into<Request<'t>>,
 {
     client_with_config(request, stream, None)
diff --git a/src/handshake/client.rs b/src/handshake/client.rs
index 8203ee5..5491fbf 100644
--- a/src/handshake/client.rs
+++ b/src/handshake/client.rs
@@ -73,7 +73,7 @@ pub struct ClientHandshake<S> {
     _marker: PhantomData<S>,
 }
 
-impl<S: Read + Write> ClientHandshake<S> {
+impl<S: Read + Write + Unpin> ClientHandshake<S> {
     /// Initiate a client handshake.
     pub fn start(
         stream: S,
@@ -124,7 +124,7 @@ impl<S: Read + Write> ClientHandshake<S> {
     }
 }
 
-impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
+impl<S: Read + Write + Unpin> HandshakeRole for ClientHandshake<S> {
     type IncomingData = Response;
     type InternalStream = S;
     type FinalResult = (WebSocket<S>, Response);
diff --git a/src/handshake/mod.rs b/src/handshake/mod.rs
index ce862d9..90bfd4d 100644
--- a/src/handshake/mod.rs
+++ b/src/handshake/mod.rs
@@ -24,6 +24,16 @@ pub struct MidHandshake<Role: HandshakeRole> {
 }
 
 impl<Role: HandshakeRole> MidHandshake<Role> {
+    /// Allow access to machine
+    pub fn get_ref(&self) -> &HandshakeMachine<Role::InternalStream> {
+        &self.machine
+    }
+
+    /// Allow mutable access to machine
+    pub fn get_mut(&mut self) -> &mut HandshakeMachine<Role::InternalStream> {
+        &mut self.machine
+    }
+
     /// Restarts the handshake process.
     pub fn handshake(mut self) -> Result<Role::FinalResult, HandshakeError<Role>> {
         let mut mach = self.machine;
@@ -91,7 +101,7 @@ pub trait HandshakeRole {
     #[doc(hidden)]
     type IncomingData: TryParse;
     #[doc(hidden)]
-    type InternalStream: Read + Write;
+    type InternalStream: Read + Write + Unpin;
     #[doc(hidden)]
     type FinalResult;
     #[doc(hidden)]
diff --git a/src/handshake/server.rs b/src/handshake/server.rs
index 615a755..696f97e 100644
--- a/src/handshake/server.rs
+++ b/src/handshake/server.rs
@@ -153,7 +153,7 @@ pub struct ServerHandshake<S, C> {
     _marker: PhantomData<S>,
 }
 
-impl<S: Read + Write, C: Callback> ServerHandshake<S, C> {
+impl<S: Read + Write + Unpin, C: Callback> ServerHandshake<S, C> {
     /// Start server handshake. `callback` specifies a custom callback which the user can pass to
     /// the handshake, this callback will be called when the a websocket client connnects to the
     /// server, you can specify the callback if you want to add additional header to the client
@@ -172,7 +172,7 @@ impl<S: Read + Write, C: Callback> ServerHandshake<S, C> {
     }
 }
 
-impl<S: Read + Write, C: Callback> HandshakeRole for ServerHandshake<S, C> {
+impl<S: Read + Write + Unpin, C: Callback> HandshakeRole for ServerHandshake<S, C> {
     type IncomingData = Request;
     type InternalStream = S;
     type FinalResult = WebSocket<S>;
diff --git a/src/server.rs b/src/server.rs
index 725d892..d4c12fd 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -18,7 +18,7 @@ use std::io::{Read, Write};
 /// If you want TLS support, use `native_tls::TlsStream` or `openssl::ssl::SslStream`
 /// for the stream here. Any `Read + Write` streams are supported, including
 /// those from `Mio` and others.
-pub fn accept_with_config<S: Read + Write>(
+pub fn accept_with_config<S: Read + Write + Unpin>(
     stream: S,
     config: Option<WebSocketConfig>,
 ) -> Result<WebSocket<S>, HandshakeError<ServerHandshake<S, NoCallback>>> {
@@ -31,7 +31,7 @@ pub fn accept_with_config<S: Read + Write>(
 /// If you want TLS support, use `native_tls::TlsStream` or `openssl::ssl::SslStream`
 /// for the stream here. Any `Read + Write` streams are supported, including
 /// those from `Mio` and others.
-pub fn accept<S: Read + Write>(
+pub fn accept<S: Read + Write + Unpin>(
     stream: S,
 ) -> Result<WebSocket<S>, HandshakeError<ServerHandshake<S, NoCallback>>> {
     accept_with_config(stream, None)
@@ -45,7 +45,7 @@ pub fn accept<S: Read + Write>(
 /// This function does the same as `accept()` but accepts an extra callback
 /// for header processing. The callback receives headers of the incoming
 /// requests and is able to add extra headers to the reply.
-pub fn accept_hdr_with_config<S: Read + Write, C: Callback>(
+pub fn accept_hdr_with_config<S: Read + Write + Unpin, C: Callback>(
     stream: S,
     callback: C,
     config: Option<WebSocketConfig>,
@@ -58,7 +58,7 @@ pub fn accept_hdr_with_config<S: Read + Write, C: Callback>(
 /// This function does the same as `accept()` but accepts an extra callback
 /// for header processing. The callback receives headers of the incoming
 /// requests and is able to add extra headers to the reply.
-pub fn accept_hdr<S: Read + Write, C: Callback>(
+pub fn accept_hdr<S: Read + Write + Unpin, C: Callback>(
     stream: S,
     callback: C,
 ) -> Result<WebSocket<S>, HandshakeError<ServerHandshake<S, C>>> {