Work directly on the tokio AsyncRead/AsyncWrite traits and convert internally

Fixes https://github.com/sdroege/async-tungstenite/issues/28
pull/40/head 0.7.0
Sebastian Dröge 4 years ago committed by Sebastian Dröge
parent b1e44a5e77
commit 3b770316d6
  1. 176
      src/tokio.rs

@ -1,6 +1,7 @@
//! `tokio` integration. //! `tokio` integration.
use tungstenite::client::IntoClientRequest; use tungstenite::client::IntoClientRequest;
use tungstenite::handshake::client::{Request, Response}; use tungstenite::handshake::client::{Request, Response};
use tungstenite::handshake::server::{Callback, NoCallback};
use tungstenite::protocol::WebSocketConfig; use tungstenite::protocol::WebSocketConfig;
use tungstenite::Error; use tungstenite::Error;
@ -20,15 +21,13 @@ pub(crate) mod tokio_tls {
use tungstenite::stream::Mode; use tungstenite::stream::Mode;
use tungstenite::Error; use tungstenite::Error;
use futures_io::{AsyncRead, AsyncWrite};
use crate::stream::Stream as StreamSwitcher; use crate::stream::Stream as StreamSwitcher;
use crate::{client_async_with_config, domain, Response, WebSocketConfig, WebSocketStream}; use crate::{client_async_with_config, domain, Response, WebSocketConfig, WebSocketStream};
use super::TokioAdapter; use super::TokioAdapter;
/// A stream that might be protected with TLS. /// A stream that might be protected with TLS.
pub type MaybeTlsStream<S> = StreamSwitcher<S, TokioAdapter<TlsStream<TokioAdapter<S>>>>; pub type MaybeTlsStream<S> = StreamSwitcher<TokioAdapter<S>, TokioAdapter<TlsStream<S>>>;
pub type AutoStream<S> = MaybeTlsStream<S>; pub type AutoStream<S> = MaybeTlsStream<S>;
@ -41,10 +40,10 @@ pub(crate) mod tokio_tls {
mode: Mode, mode: Mode,
) -> Result<AutoStream<S>, Error> ) -> Result<AutoStream<S>, Error>
where where
S: 'static + AsyncRead + AsyncWrite + Unpin, S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{ {
match mode { match mode {
Mode::Plain => Ok(StreamSwitcher::Plain(socket)), Mode::Plain => Ok(StreamSwitcher::Plain(TokioAdapter(socket))),
Mode::Tls => { Mode::Tls => {
let stream = { let stream = {
let connector = if let Some(connector) = connector { let connector = if let Some(connector) = connector {
@ -54,7 +53,7 @@ pub(crate) mod tokio_tls {
AsyncTlsConnector::from(connector) AsyncTlsConnector::from(connector)
}; };
connector connector
.connect(&domain, TokioAdapter(socket)) .connect(&domain, socket)
.await .await
.map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err))? .map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err))?
}; };
@ -74,7 +73,7 @@ pub(crate) mod tokio_tls {
) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error> ) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error>
where where
R: IntoClientRequest + Unpin, R: IntoClientRequest + Unpin,
S: 'static + AsyncRead + AsyncWrite + Unpin, S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
AutoStream<S>: Unpin, AutoStream<S>: Unpin,
{ {
let request: Request = request.into_client_request()?; let request: Request = request.into_client_request()?;
@ -91,16 +90,16 @@ pub(crate) mod tokio_tls {
#[cfg(not(any(feature = "async-tls", feature = "tokio-native-tls")))] #[cfg(not(any(feature = "async-tls", feature = "tokio-native-tls")))]
pub(crate) mod dummy_tls { pub(crate) mod dummy_tls {
use futures_io::{AsyncRead, AsyncWrite};
use tungstenite::client::{uri_mode, IntoClientRequest}; use tungstenite::client::{uri_mode, IntoClientRequest};
use tungstenite::handshake::client::Request; use tungstenite::handshake::client::Request;
use tungstenite::stream::Mode; use tungstenite::stream::Mode;
use tungstenite::Error; use tungstenite::Error;
use super::TokioAdapter;
use crate::{client_async_with_config, domain, Response, WebSocketConfig, WebSocketStream}; use crate::{client_async_with_config, domain, Response, WebSocketConfig, WebSocketStream};
pub type AutoStream<S> = S; pub type AutoStream<S> = TokioAdapter<S>;
type Connector = (); type Connector = ();
async fn wrap_stream<S>( async fn wrap_stream<S>(
@ -110,10 +109,10 @@ pub(crate) mod dummy_tls {
mode: Mode, mode: Mode,
) -> Result<AutoStream<S>, Error> ) -> Result<AutoStream<S>, Error>
where where
S: 'static + AsyncRead + AsyncWrite + Unpin, S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{ {
match mode { match mode {
Mode::Plain => Ok(socket), Mode::Plain => Ok(TokioAdapter(socket)),
Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())), Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
} }
} }
@ -126,7 +125,7 @@ pub(crate) mod dummy_tls {
) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error> ) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error>
where where
R: IntoClientRequest + Unpin, R: IntoClientRequest + Unpin,
S: 'static + AsyncRead + AsyncWrite + Unpin, S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
AutoStream<S>: Unpin, AutoStream<S>: Unpin,
{ {
let request: Request = request.into_client_request()?; let request: Request = request.into_client_request()?;
@ -145,13 +144,147 @@ pub(crate) mod dummy_tls {
use self::dummy_tls::{client_async_tls_with_connector_and_config, AutoStream}; use self::dummy_tls::{client_async_tls_with_connector_and_config, AutoStream};
#[cfg(all(feature = "async-tls", not(feature = "tokio-native-tls")))] #[cfg(all(feature = "async-tls", not(feature = "tokio-native-tls")))]
use crate::async_tls::{client_async_tls_with_connector_and_config, AutoStream}; pub(crate) mod async_tls_adapter {
use super::{
Error, IntoClientRequest, Response, TokioAdapter, WebSocketConfig, WebSocketStream,
};
use crate::stream::Stream as StreamSwitcher;
use std::marker::Unpin;
/// Creates a WebSocket handshake from a request and a stream,
/// upgrading the stream to TLS if required and using the given
/// connector and WebSocket configuration.
pub async fn client_async_tls_with_connector_and_config<R, S>(
request: R,
stream: S,
connector: Option<Connector>,
config: Option<WebSocketConfig>,
) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error>
where
R: IntoClientRequest + Unpin,
S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
AutoStream<S>: Unpin,
{
crate::async_tls::client_async_tls_with_connector_and_config(
request,
TokioAdapter(stream),
connector,
config,
)
.await
}
pub type Connector = real_async_tls::TlsConnector;
pub type MaybeTlsStream<S> = StreamSwitcher<S, real_async_tls::client::TlsStream<S>>;
pub type AutoStream<S> = MaybeTlsStream<TokioAdapter<S>>;
}
#[cfg(all(feature = "async-tls", not(feature = "tokio-native-tls")))]
pub use self::async_tls_adapter::client_async_tls_with_connector_and_config;
#[cfg(all(feature = "async-tls", not(feature = "tokio-native-tls")))] #[cfg(all(feature = "async-tls", not(feature = "tokio-native-tls")))]
type Connector = real_async_tls::TlsConnector; use self::async_tls_adapter::{AutoStream, Connector};
#[cfg(feature = "tokio-native-tls")] #[cfg(feature = "tokio-native-tls")]
use self::tokio_tls::{client_async_tls_with_connector_and_config, AutoStream, Connector}; use self::tokio_tls::{client_async_tls_with_connector_and_config, AutoStream, Connector};
/// Creates a WebSocket handshake from a request and a stream.
/// For convenience, the user may call this with a url string, a URL,
/// or a `Request`. Calling with `Request` allows the user to add
/// a WebSocket protocol or other custom headers.
///
/// Internally, this custom creates 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 on whether the handshake is successful.
///
/// This is typically used for clients who have already established, for
/// example, a TCP connection to the remote server.
pub async fn client_async<'a, R, S>(
request: R,
stream: S,
) -> Result<(WebSocketStream<TokioAdapter<S>>, Response), Error>
where
R: IntoClientRequest + Unpin,
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
client_async_with_config(request, stream, None).await
}
/// The same as `client_async()` but the one can specify a websocket configuration.
/// Please refer to `client_async()` for more details.
pub async fn client_async_with_config<'a, R, S>(
request: R,
stream: S,
config: Option<WebSocketConfig>,
) -> Result<(WebSocketStream<TokioAdapter<S>>, Response), Error>
where
R: IntoClientRequest + Unpin,
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
crate::client_async_with_config(request, TokioAdapter(stream), config).await
}
/// 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 async fn accept_async<S>(stream: S) -> Result<WebSocketStream<TokioAdapter<S>>, Error>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
accept_hdr_async(stream, NoCallback).await
}
/// The same as `accept_async()` but the one can specify a websocket configuration.
/// Please refer to `accept_async()` for more details.
pub async fn accept_async_with_config<S>(
stream: S,
config: Option<WebSocketConfig>,
) -> Result<WebSocketStream<TokioAdapter<S>>, Error>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
accept_hdr_async_with_config(stream, NoCallback, config).await
}
/// Accepts a new WebSocket connection with the provided stream.
///
/// This function does the same as `accept_async()` 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 async fn accept_hdr_async<S, C>(
stream: S,
callback: C,
) -> Result<WebSocketStream<TokioAdapter<S>>, Error>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
C: Callback + Unpin,
{
accept_hdr_async_with_config(stream, callback, None).await
}
/// The same as `accept_hdr_async()` but the one can specify a websocket configuration.
/// Please refer to `accept_hdr_async()` for more details.
pub async fn accept_hdr_async_with_config<S, C>(
stream: S,
callback: C,
config: Option<WebSocketConfig>,
) -> Result<WebSocketStream<TokioAdapter<S>>, Error>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
C: Callback + Unpin,
{
crate::accept_hdr_async_with_config(TokioAdapter(stream), callback, config).await
}
/// Type alias for the stream type of the `client_async()` functions. /// Type alias for the stream type of the `client_async()` functions.
pub type ClientStream<S> = AutoStream<S>; pub type ClientStream<S> = AutoStream<S>;
@ -164,7 +297,7 @@ pub async fn client_async_tls<R, S>(
) -> Result<(WebSocketStream<ClientStream<S>>, Response), Error> ) -> Result<(WebSocketStream<ClientStream<S>>, Response), Error>
where where
R: IntoClientRequest + Unpin, R: IntoClientRequest + Unpin,
S: 'static + AsyncRead + AsyncWrite + Unpin, S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
AutoStream<S>: Unpin, AutoStream<S>: Unpin,
{ {
client_async_tls_with_connector_and_config(request, stream, None, None).await client_async_tls_with_connector_and_config(request, stream, None, None).await
@ -181,7 +314,7 @@ pub async fn client_async_tls_with_config<R, S>(
) -> Result<(WebSocketStream<ClientStream<S>>, Response), Error> ) -> Result<(WebSocketStream<ClientStream<S>>, Response), Error>
where where
R: IntoClientRequest + Unpin, R: IntoClientRequest + Unpin,
S: 'static + AsyncRead + AsyncWrite + Unpin, S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
AutoStream<S>: Unpin, AutoStream<S>: Unpin,
{ {
client_async_tls_with_connector_and_config(request, stream, None, config).await client_async_tls_with_connector_and_config(request, stream, None, config).await
@ -198,14 +331,14 @@ pub async fn client_async_tls_with_connector<R, S>(
) -> Result<(WebSocketStream<ClientStream<S>>, Response), Error> ) -> Result<(WebSocketStream<ClientStream<S>>, Response), Error>
where where
R: IntoClientRequest + Unpin, R: IntoClientRequest + Unpin,
S: 'static + AsyncRead + AsyncWrite + Unpin, S: 'static + tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
AutoStream<S>: Unpin, AutoStream<S>: Unpin,
{ {
client_async_tls_with_connector_and_config(request, stream, connector, None).await client_async_tls_with_connector_and_config(request, stream, connector, None).await
} }
/// Type alias for the stream type of the `connect_async()` functions. /// Type alias for the stream type of the `connect_async()` functions.
pub type ConnectStream = ClientStream<TokioAdapter<TcpStream>>; pub type ConnectStream = ClientStream<TcpStream>;
/// Connect to a given URL. /// Connect to a given URL.
pub async fn connect_async<R>( pub async fn connect_async<R>(
@ -232,7 +365,7 @@ where
let try_socket = TcpStream::connect((domain.as_str(), port)).await; let try_socket = TcpStream::connect((domain.as_str(), port)).await;
let socket = try_socket.map_err(Error::Io)?; let socket = try_socket.map_err(Error::Io)?;
client_async_tls_with_connector_and_config(request, TokioAdapter(socket), None, config).await client_async_tls_with_connector_and_config(request, socket, None, config).await
} }
#[cfg(any(feature = "async-tls", feature = "tokio-native-tls"))] #[cfg(any(feature = "async-tls", feature = "tokio-native-tls"))]
@ -264,8 +397,7 @@ where
let try_socket = TcpStream::connect((domain.as_str(), port)).await; let try_socket = TcpStream::connect((domain.as_str(), port)).await;
let socket = try_socket.map_err(Error::Io)?; let socket = try_socket.map_err(Error::Io)?;
client_async_tls_with_connector_and_config(request, TokioAdapter(socket), connector, config) client_async_tls_with_connector_and_config(request, socket, connector, config).await
.await
} }
use pin_project::pin_project; use pin_project::pin_project;

Loading…
Cancel
Save