Dominik Nakamura
3 years ago
parent
bc1b88b820
commit
32450ae5af
@ -0,0 +1,219 @@ |
|||||||
|
//! Connection helper.
|
||||||
|
use std::io::{Read, Write}; |
||||||
|
|
||||||
|
use crate::{ |
||||||
|
client::{client_with_config, uri_mode, IntoClientRequest}, |
||||||
|
error::UrlError, |
||||||
|
handshake::client::Response, |
||||||
|
protocol::WebSocketConfig, |
||||||
|
stream::MaybeTlsStream, |
||||||
|
ClientHandshake, Error, HandshakeError, Result, WebSocket, |
||||||
|
}; |
||||||
|
|
||||||
|
/// A connector that can be used when establishing connections, allowing to control whether
|
||||||
|
/// `native-tls` or `rustls` is used to create a TLS connection. Or TLS can be disabled with the
|
||||||
|
/// `Plain` variant.
|
||||||
|
#[non_exhaustive] |
||||||
|
#[allow(missing_debug_implementations)] |
||||||
|
pub enum Connector { |
||||||
|
/// Plain (non-TLS) connector.
|
||||||
|
Plain, |
||||||
|
/// `native-tls` TLS connector.
|
||||||
|
#[cfg(feature = "native-tls")] |
||||||
|
NativeTls(native_tls_crate::TlsConnector), |
||||||
|
/// `rustls` TLS connector.
|
||||||
|
#[cfg(feature = "__rustls-tls")] |
||||||
|
Rustls(std::sync::Arc<rustls::ClientConfig>), |
||||||
|
} |
||||||
|
|
||||||
|
mod encryption { |
||||||
|
#[cfg(feature = "native-tls")] |
||||||
|
pub mod native_tls { |
||||||
|
use native_tls_crate::{HandshakeError as TlsHandshakeError, TlsConnector}; |
||||||
|
|
||||||
|
use std::io::{Read, Write}; |
||||||
|
|
||||||
|
use crate::{ |
||||||
|
error::TlsError, |
||||||
|
stream::{MaybeTlsStream, Mode}, |
||||||
|
Error, Result, |
||||||
|
}; |
||||||
|
|
||||||
|
pub fn wrap_stream<S>( |
||||||
|
socket: S, |
||||||
|
domain: &str, |
||||||
|
mode: Mode, |
||||||
|
tls_connector: Option<TlsConnector>, |
||||||
|
) -> Result<MaybeTlsStream<S>> |
||||||
|
where |
||||||
|
S: Read + Write, |
||||||
|
{ |
||||||
|
match mode { |
||||||
|
Mode::Plain => Ok(MaybeTlsStream::Plain(socket)), |
||||||
|
Mode::Tls => { |
||||||
|
let try_connector = tls_connector.map_or_else(TlsConnector::new, Ok); |
||||||
|
let connector = try_connector.map_err(TlsError::Native)?; |
||||||
|
let connected = connector.connect(domain, socket); |
||||||
|
match connected { |
||||||
|
Err(e) => match e { |
||||||
|
TlsHandshakeError::Failure(f) => Err(Error::Tls(f.into())), |
||||||
|
TlsHandshakeError::WouldBlock(_) => { |
||||||
|
panic!("Bug: TLS handshake not blocked") |
||||||
|
} |
||||||
|
}, |
||||||
|
Ok(s) => Ok(MaybeTlsStream::NativeTls(s)), |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
#[cfg(feature = "__rustls-tls")] |
||||||
|
pub mod rustls { |
||||||
|
use rustls::{ClientConfig, ClientSession, StreamOwned}; |
||||||
|
use webpki::DNSNameRef; |
||||||
|
|
||||||
|
use std::{ |
||||||
|
io::{Read, Write}, |
||||||
|
sync::Arc, |
||||||
|
}; |
||||||
|
|
||||||
|
use crate::{ |
||||||
|
error::TlsError, |
||||||
|
stream::{MaybeTlsStream, Mode}, |
||||||
|
Result, |
||||||
|
}; |
||||||
|
|
||||||
|
pub fn wrap_stream<S>( |
||||||
|
socket: S, |
||||||
|
domain: &str, |
||||||
|
mode: Mode, |
||||||
|
tls_connector: Option<Arc<ClientConfig>>, |
||||||
|
) -> Result<MaybeTlsStream<S>> |
||||||
|
where |
||||||
|
S: Read + Write, |
||||||
|
{ |
||||||
|
match mode { |
||||||
|
Mode::Plain => Ok(MaybeTlsStream::Plain(socket)), |
||||||
|
Mode::Tls => { |
||||||
|
let config = match tls_connector { |
||||||
|
Some(config) => config, |
||||||
|
None => { |
||||||
|
#[allow(unused_mut)] |
||||||
|
let mut config = ClientConfig::new(); |
||||||
|
#[cfg(feature = "rustls-tls-native-roots")] |
||||||
|
{ |
||||||
|
config.root_store = rustls_native_certs::load_native_certs() |
||||||
|
.map_err(|(_, err)| err)?; |
||||||
|
} |
||||||
|
#[cfg(feature = "rustls-tls-webpki-roots")] |
||||||
|
{ |
||||||
|
config |
||||||
|
.root_store |
||||||
|
.add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS); |
||||||
|
} |
||||||
|
|
||||||
|
Arc::new(config) |
||||||
|
} |
||||||
|
}; |
||||||
|
let domain = DNSNameRef::try_from_ascii_str(domain).map_err(TlsError::Dns)?; |
||||||
|
let client = ClientSession::new(&config, domain); |
||||||
|
let stream = StreamOwned::new(client, socket); |
||||||
|
|
||||||
|
Ok(MaybeTlsStream::Rustls(stream)) |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
pub mod plain { |
||||||
|
use std::io::{Read, Write}; |
||||||
|
|
||||||
|
use crate::{ |
||||||
|
error::UrlError, |
||||||
|
stream::{MaybeTlsStream, Mode}, |
||||||
|
Error, Result, |
||||||
|
}; |
||||||
|
|
||||||
|
pub fn wrap_stream<S>(socket: S, mode: Mode) -> Result<MaybeTlsStream<S>> |
||||||
|
where |
||||||
|
S: Read + Write, |
||||||
|
{ |
||||||
|
match mode { |
||||||
|
Mode::Plain => Ok(MaybeTlsStream::Plain(socket)), |
||||||
|
Mode::Tls => Err(Error::Url(UrlError::TlsFeatureNotEnabled)), |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
type TlsHandshakeError<S> = HandshakeError<ClientHandshake<MaybeTlsStream<S>>>; |
||||||
|
|
||||||
|
/// Creates a WebSocket handshake from a request and a stream,
|
||||||
|
/// upgrading the stream to TLS if required.
|
||||||
|
pub fn client_tls<R, S>( |
||||||
|
request: R, |
||||||
|
stream: S, |
||||||
|
) -> Result<(WebSocket<MaybeTlsStream<S>>, Response), TlsHandshakeError<S>> |
||||||
|
where |
||||||
|
R: IntoClientRequest, |
||||||
|
S: Read + Write, |
||||||
|
{ |
||||||
|
client_tls_with_config(request, stream, None, None) |
||||||
|
} |
||||||
|
|
||||||
|
/// The same as [`client_tls()`] but one can specify a websocket configuration,
|
||||||
|
/// and an optional connector. If no connector is specified, a default one will
|
||||||
|
/// be created.
|
||||||
|
///
|
||||||
|
/// Please refer to [`client_tls()`] for more details.
|
||||||
|
pub fn client_tls_with_config<R, S>( |
||||||
|
request: R, |
||||||
|
stream: S, |
||||||
|
config: Option<WebSocketConfig>, |
||||||
|
connector: Option<Connector>, |
||||||
|
) -> Result<(WebSocket<MaybeTlsStream<S>>, Response), TlsHandshakeError<S>> |
||||||
|
where |
||||||
|
R: IntoClientRequest, |
||||||
|
S: Read + Write, |
||||||
|
{ |
||||||
|
let request = request.into_client_request()?; |
||||||
|
|
||||||
|
#[cfg(any(feature = "native-tls", feature = "__rustls-tls"))] |
||||||
|
let domain = match request.uri().host() { |
||||||
|
Some(d) => Ok(d.to_string()), |
||||||
|
None => Err(Error::Url(UrlError::NoHostName)), |
||||||
|
}?; |
||||||
|
|
||||||
|
let mode = uri_mode(&request.uri())?; |
||||||
|
|
||||||
|
let stream = match connector { |
||||||
|
Some(conn) => match conn { |
||||||
|
#[cfg(feature = "native-tls")] |
||||||
|
Connector::NativeTls(conn) => { |
||||||
|
self::encryption::native_tls::wrap_stream(stream, &domain, mode, Some(conn)) |
||||||
|
} |
||||||
|
#[cfg(feature = "__rustls-tls")] |
||||||
|
Connector::Rustls(conn) => { |
||||||
|
self::encryption::rustls::wrap_stream(stream, &domain, mode, Some(conn)) |
||||||
|
} |
||||||
|
Connector::Plain => self::encryption::plain::wrap_stream(stream, mode), |
||||||
|
}, |
||||||
|
None => { |
||||||
|
#[cfg(feature = "native-tls")] |
||||||
|
{ |
||||||
|
self::encryption::native_tls::wrap_stream(stream, &domain, mode, None) |
||||||
|
} |
||||||
|
#[cfg(all(feature = "__rustls-tls", not(feature = "native-tls")))] |
||||||
|
{ |
||||||
|
self::encryption::rustls::wrap_stream(stream, &domain, mode, None) |
||||||
|
} |
||||||
|
#[cfg(not(any(feature = "native-tls", feature = "__rustls-tls")))] |
||||||
|
{ |
||||||
|
self::encryption::plain::wrap_stream(stream, mode) |
||||||
|
} |
||||||
|
} |
||||||
|
}?; |
||||||
|
|
||||||
|
client_with_config(request, stream, config) |
||||||
|
} |
||||||
Loading…
Reference in new issue