refactor: use asyncronous handshakes

pull/1/head
Alexey Galakhov 8 years ago
parent 0a21d17042
commit 4b646308be
  1. 2
      Cargo.toml
  2. 7
      examples/client.rs
  3. 10
      examples/server.rs
  4. 219
      src/lib.rs

@ -8,5 +8,5 @@ futures = "*"
tokio-core = "*" tokio-core = "*"
url = "*" url = "*"
tungstenite = { git = "ssh://git@bitbucket.org/mikogo/tungstenite-rs.git" } tungstenite = { git = "ssh://git@bitbucket.org/mikogo/tungstenite-rs.git", branch = "handshake-refactor" }

@ -25,10 +25,10 @@ use futures::sync::mpsc;
use futures::{Future, Sink, Stream}; use futures::{Future, Sink, Stream};
use tokio_core::net::TcpStream; use tokio_core::net::TcpStream;
use tokio_core::reactor::Core; use tokio_core::reactor::Core;
use tokio_tungstenite::ClientHandshakeExt;
use tungstenite::handshake::client::{ClientHandshake, Request};
use tungstenite::protocol::Message; use tungstenite::protocol::Message;
use tokio_tungstenite::client_async;
fn main() { fn main() {
// Specify the server address to which the client will be connecting. // Specify the server address to which the client will be connecting.
let connect_addr = env::args().nth(1).unwrap_or_else(|| { let connect_addr = env::args().nth(1).unwrap_or_else(|| {
@ -69,8 +69,7 @@ fn main() {
// more work from the remote then we can exit. // more work from the remote then we can exit.
let mut stdout = io::stdout(); let mut stdout = io::stdout();
let client = tcp.and_then(|stream| { let client = tcp.and_then(|stream| {
let req = Request { url: url }; client_async(url, stream).and_then(|ws_stream| {
ClientHandshake::<TcpStream>::new_async(stream, req).and_then(|ws_stream| {
println!("WebSocket handshake has been successfully completed"); println!("WebSocket handshake has been successfully completed");
// `sink` is the stream of messages going out. // `sink` is the stream of messages going out.

@ -29,13 +29,13 @@ use std::io::{Error, ErrorKind};
use std::rc::Rc; use std::rc::Rc;
use futures::stream::Stream; use futures::stream::Stream;
use futures::{Future}; use futures::Future;
use tokio_core::net::{TcpListener, TcpStream}; use tokio_core::net::TcpListener;
use tokio_core::reactor::Core; use tokio_core::reactor::Core;
use tokio_tungstenite::ServerHandshakeExt;
use tungstenite::handshake::server::ServerHandshake;
use tungstenite::protocol::Message; use tungstenite::protocol::Message;
use tokio_tungstenite::accept_async;
fn main() { fn main() {
let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string()); let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string());
let addr = addr.parse().unwrap(); let addr = addr.parse().unwrap();
@ -60,7 +60,7 @@ fn main() {
let connections_inner = connections.clone(); let connections_inner = connections.clone();
let handle_inner = handle.clone(); let handle_inner = handle.clone();
ServerHandshake::<TcpStream>::new_async(stream).and_then(move |ws_stream| { accept_async(stream).and_then(move |ws_stream| {
println!("New WebSocket connection: {}", addr); println!("New WebSocket connection: {}", addr);
// Create a channel for our stream, which other sockets will use to // Create a channel for our stream, which other sockets will use to

@ -13,7 +13,12 @@
//! functionality provided by the `tungestenite` crate, on which this crate is //! functionality provided by the `tungestenite` crate, on which this crate is
//! built. Configuration is done through `tungestenite` crate as well. //! built. Configuration is done through `tungestenite` crate as well.
#![deny(missing_docs)] #![deny(
missing_docs,
unused_must_use,
unused_mut,
unused_imports,
unused_import_braces)]
#[macro_use] #[macro_use]
extern crate futures; extern crate futures;
@ -23,14 +28,55 @@ extern crate url;
use std::io::ErrorKind; use std::io::ErrorKind;
use futures::{Poll, Future, Async, AsyncSink, Stream, Sink, StartSend, task}; use futures::{Poll, Future, Async, AsyncSink, Stream, Sink, StartSend};
use tokio_core::io::Io; use tokio_core::io::Io;
use tungstenite::handshake::client::{ClientHandshake, Request}; use url::Url;
use tungstenite::handshake::client::ClientHandshake;
use tungstenite::handshake::server::ServerHandshake; use tungstenite::handshake::server::ServerHandshake;
use tungstenite::handshake::{Handshake, HandshakeResult}; use tungstenite::handshake::{HandshakeRole, HandshakeError};
use tungstenite::protocol::{WebSocket, Message}; use tungstenite::protocol::{WebSocket, Message};
use tungstenite::error::Error as WsError; use tungstenite::error::Error as WsError;
use tungstenite::client;
use tungstenite::server;
/// Create a handshake provided stream and assuming the provided request.
///
/// This function will internally call `client::client` to create a
/// handshake representation and returns a future representing the
/// resolution of the WebSocket handshake. The returned future will resolve
/// to either `WebSocketStream<S>` or `Error` depending if it's successful
/// or not.
///
/// This is typically used for clients who have already established, for
/// example, a TCP connection to the remove server.
pub fn client_async<S: Io>(url: Url, stream: S) -> ConnectAsync<S> {
ConnectAsync {
inner: MidHandshake {
inner: Some(client::client(url, stream))
}
}
}
/// Accepts a new WebSocket connection with the provided stream.
///
/// This function will internally call `server::accept` to create a
/// handshake representation and returns a future representing the
/// resolution of the WebSocket handshake. The returned future will resolve
/// to either `WebSocketStream<S>` or `Error` depending if it's successful
/// or not.
///
/// This is typically used after a socket has been accepted from a
/// `TcpListener`. That socket is then passed to this function to perform
/// the server half of the accepting a client's websocket connection.
pub fn accept_async<S: Io>(stream: S) -> AcceptAsync<S> {
AcceptAsync {
inner: MidHandshake {
inner: Some(server::accept(stream))
}
}
}
/// A wrapper around an underlying raw stream which implements the WebSocket /// A wrapper around an underlying raw stream which implements the WebSocket
/// protocol. /// protocol.
@ -45,113 +91,81 @@ pub struct WebSocketStream<S> {
inner: WebSocket<S>, inner: WebSocket<S>,
} }
/// Future returned from `ClientHandshakeExt::new_async` which will resolve impl<T> Stream for WebSocketStream<T> where T: Io {
/// once the connection handshake has finished. type Item = Message;
pub struct ClientHandshakeAsync<S> { type Error = WsError;
inner: Option<ClientHandshake<S>>,
}
/// Future returned from `ServerHandshakeExt::new_async` which will resolve fn poll(&mut self) -> Poll<Option<Message>, WsError> {
/// once the connection handshake has finished. self.inner.read_message().map(|m| Some(m)).to_async()
pub struct ServerHandshakeAsync<S: Io> { }
inner: Option<ServerHandshake<S>>,
} }
/// Extension trait for the `ClientHandshake` type in the `tungstenite` crate. impl<T> Sink for WebSocketStream<T> where T: Io {
pub trait ClientHandshakeExt { type SinkItem = Message;
/// Create a handshake provided stream and assuming the provided request. type SinkError = WsError;
///
/// This function will internally call `ClientHandshake::new` to create a fn start_send(&mut self, item: Message) -> StartSend<Message, WsError> {
/// handshake representation and returns a future representing the try!(self.inner.write_message(item).to_async());
/// resolution of the WebSocket handshake. The returned future will resolve Ok(AsyncSink::Ready)
/// to either `WebSocketStream<S>` or `Error` depending if it's successful }
/// or not.
/// fn poll_complete(&mut self) -> Poll<(), WsError> {
/// This is typically used for clients who have already established, for self.inner.write_pending().to_async()
/// example, a TCP connection to the remove server. }
fn new_async<S: Io>(stream: S, request: Request) -> ClientHandshakeAsync<S>;
} }
/// Extension trait for the `ServerHandshake` type in the `tungstenite` crate. /// Future returned from connect_async() which will resolve
pub trait ServerHandshakeExt { /// once the connection handshake has finished.
/// Accepts a new WebSocket connection with the provided stream. pub struct ConnectAsync<S> {
/// inner: MidHandshake<S, ClientHandshake>,
/// This function will internally call `ServerHandshake::new` to create a
/// handshake representation and returns a future representing the
/// resolution of the WebSocket handshake. The returned future will resolve
/// to either `WebSocketStream<S>` or `Error` depending if it's successful
/// or not.
///
/// This is typically used after a socket has been accepted from a
/// `TcpListener`. That socket is then passed to this function to perform
/// the server half of the accepting a client's websocket connection.
fn new_async<S: Io>(stream: S) -> ServerHandshakeAsync<S>;
} }
impl<S: Io> ClientHandshakeExt for ClientHandshake<S> { impl<S: Io> Future for ConnectAsync<S> {
fn new_async<Stream: Io>(stream: Stream, request: Request) -> ClientHandshakeAsync<Stream> { type Item = WebSocketStream<S>;
ClientHandshakeAsync { type Error = WsError;
inner: Some(ClientHandshake::new(stream, request)),
} fn poll(&mut self) -> Poll<WebSocketStream<S>, WsError> {
self.inner.poll()
} }
} }
impl<S: Io> ServerHandshakeExt for ServerHandshake<S> { /// Future returned from accept_async() which will resolve
fn new_async<Stream: Io>(stream: Stream) -> ServerHandshakeAsync<Stream> { /// once the connection handshake has finished.
ServerHandshakeAsync { pub struct AcceptAsync<S> {
inner: Some(ServerHandshake::new(stream)), inner: MidHandshake<S, ServerHandshake>,
}
}
} }
// FIXME: `ClientHandshakeAsync<S>` and `ServerHandshakeAsync<S>` have the same implementation, we impl<S: Io> Future for AcceptAsync<S> {
// have to get rid of this copy-pasting one day. But currently I don't see an elegant way to write
// it.
impl<S: Io> Future for ClientHandshakeAsync<S> {
type Item = WebSocketStream<S>; type Item = WebSocketStream<S>;
type Error = WsError; type Error = WsError;
fn poll(&mut self) -> Poll<WebSocketStream<S>, WsError> { fn poll(&mut self) -> Poll<WebSocketStream<S>, WsError> {
let hs = self.inner.take().expect("Cannot poll a handshake twice"); self.inner.poll()
match hs.handshake()? {
HandshakeResult::Done(stream) => {
Ok(WebSocketStream { inner: stream }.into())
},
HandshakeResult::Incomplete(handshake) => {
// FIXME: Remove this line after we have a guarantee that the underlying handshake
// calls to both `read()`/`write()`. Or replace it by `poll_read()` and
// `poll_write()` (this requires making the handshake's stream public).
task::park().unpark();
self.inner = Some(handshake);
Ok(Async::NotReady)
},
}
} }
} }
// FIXME: `ClientHandshakeAsync<S>` and `ServerHandshakeAsync<S>` have the same implementation, we struct MidHandshake<S, R> {
// have to get rid of this copy-pasting one day. But currently I don't see an elegant way to write inner: Option<Result<WebSocket<S>, HandshakeError<S, R>>>,
// it. }
impl<S: Io> Future for ServerHandshakeAsync<S> {
impl<S: Io, R: HandshakeRole> Future for MidHandshake<S, R> {
type Item = WebSocketStream<S>; type Item = WebSocketStream<S>;
type Error = WsError; type Error = WsError;
fn poll(&mut self) -> Poll<WebSocketStream<S>, WsError> { fn poll(&mut self) -> Poll<WebSocketStream<S>, WsError> {
let hs = self.inner.take().expect("Cannot poll a handshake twice"); match self.inner.take().expect("cannot poll MidHandshake twice") {
match hs.handshake()? { Ok(stream) => Ok(WebSocketStream { inner: stream }.into()),
HandshakeResult::Done(stream) => { Err(HandshakeError::Failure(e)) => Err(e),
Ok(WebSocketStream { inner: stream }.into()) Err(HandshakeError::Interrupted(s)) => {
}, match s.handshake() {
HandshakeResult::Incomplete(handshake) => { Ok(stream) => Ok(WebSocketStream { inner: stream }.into()),
// FIXME: Remove this line after we have a guarantee that the underlying handshake Err(HandshakeError::Failure(e)) => Err(e),
// calls to both `read()`/`write()`. Or replace it by `poll_read()` and Err(HandshakeError::Interrupted(s)) => {
// `poll_write()` (this requires making the handshake's stream public). self.inner = Some(Err(HandshakeError::Interrupted(s)));
task::park().unpark(); Ok(Async::NotReady)
}
self.inner = Some(handshake); }
Ok(Async::NotReady) }
},
} }
} }
} }
@ -176,29 +190,6 @@ impl<T> ToAsync for Result<T, WsError> {
} }
} }
impl<T> Stream for WebSocketStream<T> where T: Io {
type Item = Message;
type Error = WsError;
fn poll(&mut self) -> Poll<Option<Message>, WsError> {
self.inner.read_message().map(|m| Some(m)).to_async()
}
}
impl<T> Sink for WebSocketStream<T> where T: Io {
type SinkItem = Message;
type SinkError = WsError;
fn start_send(&mut self, item: Message) -> StartSend<Message, WsError> {
try!(self.inner.write_message(item).to_async());
Ok(AsyncSink::Ready)
}
fn poll_complete(&mut self) -> Poll<(), WsError> {
self.inner.write_pending().to_async()
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@ -209,8 +200,6 @@ mod tests {
use futures::{Future, Stream}; use futures::{Future, Stream};
use tokio_core::net::{TcpStream, TcpListener}; use tokio_core::net::{TcpStream, TcpListener};
use tokio_core::reactor::Core; use tokio_core::reactor::Core;
use tungstenite::handshake::server::ServerHandshake;
use tungstenite::handshake::client::{ClientHandshake, Request};
#[test] #[test]
fn handshakes() { fn handshakes() {
@ -227,7 +216,7 @@ mod tests {
let connections = listener.incoming(); let connections = listener.incoming();
tx.send(()).unwrap(); tx.send(()).unwrap();
let handshakes = connections.and_then(|(connection, _)| { let handshakes = connections.and_then(|(connection, _)| {
ServerHandshake::<TcpStream>::new_async(connection) accept_async(connection)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e)) .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}); });
let server = handshakes.for_each(|_| { let server = handshakes.for_each(|_| {
@ -244,7 +233,7 @@ mod tests {
let tcp = TcpStream::connect(&address, &handle); let tcp = TcpStream::connect(&address, &handle);
let handshake = tcp.and_then(|stream| { let handshake = tcp.and_then(|stream| {
let url = url::Url::parse("ws://localhost:12345/").unwrap(); let url = url::Url::parse("ws://localhost:12345/").unwrap();
ClientHandshake::<TcpStream>::new_async(stream, Request { url: url }) client_async(url, stream)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e)) .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}); });
let client = handshake.and_then(|_| { let client = handshake.and_then(|_| {

Loading…
Cancel
Save