diff --git a/Cargo.toml b/Cargo.toml index 5ba5ac0..57dcaac 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,11 +10,36 @@ documentation = "https://docs.rs/tokio-tungstenite" repository = "https://github.com/snapview/tokio-tungstenite" version = "0.2.1" +[features] +default = ["connect", "tls"] +connect = ["tokio-dns-unofficial", "tokio-core"] +tls = ["tokio-tls", "native-tls", "bytes"] + [dependencies] futures = "*" tokio-io = "*" tungstenite = "*" url = "*" +[dependencies.bytes] +optional = true +version = "*" + +[dependencies.native-tls] +optional = true +version = "*" + +[dependencies.tokio-dns-unofficial] +optional = true +version = "*" + +[dependencies.tokio-core] +optional = true +version = "*" + +[dependencies.tokio-tls] +optional = true +version = "*" + [dev-dependencies] tokio-core = "*" diff --git a/src/connect.rs b/src/connect.rs new file mode 100644 index 0000000..6743019 --- /dev/null +++ b/src/connect.rs @@ -0,0 +1,93 @@ +//! Connection helper. + +extern crate tokio_dns; +extern crate tokio_core; + +use self::tokio_dns::tcp_connect; +use self::tokio_core::reactor::Remote; + +use futures::{Future, BoxFuture}; +use futures::future; + +use super::{WebSocketStream, Request, client_async}; +use tungstenite::Error; +use tungstenite::client::url_mode; + +#[cfg(feature="tls")] +mod encryption { + extern crate native_tls; + extern crate tokio_tls; + + use super::tokio_core::net::TcpStream; + + use self::native_tls::TlsConnector; + use self::tokio_tls::{TlsConnectorExt, TlsStream}; + + use futures::{Future, BoxFuture}; + use futures::future; + + use tungstenite::Error; + use tungstenite::stream::Mode; + + pub use stream::Stream as StreamSwitcher; + pub type AutoStream = StreamSwitcher>; + + pub fn wrap_stream(socket: TcpStream, domain: String, mode: Mode) -> BoxFuture { + match mode { + Mode::Plain => future::ok(StreamSwitcher::Plain(socket)).boxed(), + Mode::Tls => { + future::result(TlsConnector::builder()) + .and_then(move |builder| future::result(builder.build())) + .and_then(move |connector| connector.connect_async(&domain, socket)) + .map(|s| StreamSwitcher::Tls(s)) + .map_err(|e| Error::Tls(e)) + .boxed() + } + } + } +} + +#[cfg(not(feature="tls"))] +mod encryption { + use super::tokio_core::net::TcpStream; + + use futures::{Future, BoxFuture}; + use futures::future; + + use tungstenite::Error; + use tungstenite::stream::Mode; + + pub type AutoStream = TcpStream; + + pub fn wrap_stream(socket: TcpStream, _domain: String, mode: Mode) -> BoxFuture { + match mode { + Mode::Plain => future::ok(socket).boxed(), + Mode::Tls => future::err(Error::Url("TLS support not compiled in.".into())).boxed(), + } + } +} + +use self::encryption::{AutoStream, wrap_stream}; + +/// Connect to a given URL. +pub fn connect_async(request: R, handle: Remote) -> BoxFuture, Error> +where R: Into> +{ + let request: Request = request.into(); + + // Make sure we check domain and mode first. URL must be valid. + let mode = match url_mode(&request.url) { + Ok(m) => m, + Err(e) => return future::err(e.into()).boxed(), + }; + let domain = match request.url.host_str() { + Some(d) => d.to_string(), + None => return future::err(Error::Url("No host name in the URL".into())).boxed(), + }; + let port = request.url.port_or_known_default().expect("Bug: port unknown"); + + tcp_connect((domain.as_str(), port), handle).map_err(|e| e.into()) + .and_then(move |socket| wrap_stream(socket, domain, mode)) + .and_then(move |stream| client_async(request, stream)) + .boxed() +} diff --git a/src/lib.rs b/src/lib.rs index 0d40efd..6076272 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,6 +25,11 @@ extern crate tokio_io; extern crate tungstenite; extern crate url; +#[cfg(feature="connect")] +mod connect; +#[cfg(all(feature="connect", feature="tls"))] +mod stream; + use std::io::ErrorKind; use futures::{Poll, Future, Async, AsyncSink, Stream, Sink, StartSend}; @@ -39,6 +44,9 @@ use tungstenite::protocol::{WebSocket, Message}; use tungstenite::error::Error as WsError; use tungstenite::server; +#[cfg(feature="connect")] +pub use connect::connect_async; + /// A WebSocket request pub struct Request<'a> { /// URL of the request. diff --git a/src/stream.rs b/src/stream.rs new file mode 100644 index 0000000..51992e3 --- /dev/null +++ b/src/stream.rs @@ -0,0 +1,73 @@ +//! Convenience wrapper for streams to switch between plain TCP and TLS at runtime. +//! +//! There is no dependency on actual TLS implementations. Everything like +//! `native_tls` or `openssl` will work as long as there is a TLS stream supporting standard +//! `Read + Write` traits. + +extern crate bytes; + +use std::io::{Read, Write, Result as IoResult, Error as IoError}; + +use self::bytes::{Buf, BufMut}; +use futures::Poll; +use tokio_io::{AsyncRead, AsyncWrite}; + +/// Stream, either plain TCP or TLS. +pub enum Stream { + Plain(S), + Tls(T), +} + +impl Read for Stream { + fn read(&mut self, buf: &mut [u8]) -> IoResult { + match *self { + Stream::Plain(ref mut s) => s.read(buf), + Stream::Tls(ref mut s) => s.read(buf), + } + } +} + +impl Write for Stream { + fn write(&mut self, buf: &[u8]) -> IoResult { + match *self { + Stream::Plain(ref mut s) => s.write(buf), + Stream::Tls(ref mut s) => s.write(buf), + } + } + fn flush(&mut self) -> IoResult<()> { + match *self { + Stream::Plain(ref mut s) => s.flush(), + Stream::Tls(ref mut s) => s.flush(), + } + } +} + +impl AsyncRead for Stream { + unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool { + match *self { + Stream::Plain(ref s) => s.prepare_uninitialized_buffer(buf), + Stream::Tls(ref s) => s.prepare_uninitialized_buffer(buf), + } + } + fn read_buf(&mut self, buf: &mut B) -> Poll { + match *self { + Stream::Plain(ref mut s) => s.read_buf(buf), + Stream::Tls(ref mut s) => s.read_buf(buf), + } + } +} + +impl AsyncWrite for Stream { + fn shutdown(&mut self) -> Poll<(), IoError> { + match *self { + Stream::Plain(ref mut s) => s.shutdown(), + Stream::Tls(ref mut s) => s.shutdown(), + } + } + fn write_buf(&mut self, buf: &mut B) -> Poll { + match *self { + Stream::Plain(ref mut s) => s.write_buf(buf), + Stream::Tls(ref mut s) => s.write_buf(buf), + } + } +}