73 lines
2.2 KiB
Rust
73 lines
2.2 KiB
Rust
use std::io::{self, IoSlice};
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
|
|
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
|
use tokio::net::TcpStream;
|
|
|
|
// TLS state is large; plain TCP keeps its direct, allocation-free representation.
|
|
pub(crate) enum Transport {
|
|
Tcp(TcpStream),
|
|
Tls(Box<tokio_rustls::TlsStream<TcpStream>>),
|
|
}
|
|
|
|
macro_rules! dispatch {
|
|
($value:expr, $method:ident $(, $argument:expr)*) => {
|
|
match $value {
|
|
Transport::Tcp(stream) => Pin::new(stream).$method($($argument),*),
|
|
Transport::Tls(stream) => Pin::new(stream.as_mut()).$method($($argument),*),
|
|
}
|
|
};
|
|
}
|
|
|
|
impl AsyncRead for Transport {
|
|
fn poll_read(
|
|
self: Pin<&mut Self>,
|
|
context: &mut Context<'_>,
|
|
buffer: &mut ReadBuf<'_>,
|
|
) -> Poll<io::Result<()>> {
|
|
dispatch!(self.get_mut(), poll_read, context, buffer)
|
|
}
|
|
}
|
|
|
|
impl AsyncWrite for Transport {
|
|
fn poll_write(
|
|
self: Pin<&mut Self>,
|
|
context: &mut Context<'_>,
|
|
bytes: &[u8],
|
|
) -> Poll<io::Result<usize>> {
|
|
dispatch!(self.get_mut(), poll_write, context, bytes)
|
|
}
|
|
fn poll_write_vectored(
|
|
self: Pin<&mut Self>,
|
|
context: &mut Context<'_>,
|
|
buffers: &[IoSlice<'_>],
|
|
) -> Poll<io::Result<usize>> {
|
|
dispatch!(self.get_mut(), poll_write_vectored, context, buffers)
|
|
}
|
|
fn is_write_vectored(&self) -> bool {
|
|
match self {
|
|
Self::Tcp(stream) => stream.is_write_vectored(),
|
|
Self::Tls(stream) => stream.is_write_vectored(),
|
|
}
|
|
}
|
|
fn poll_flush(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
dispatch!(self.get_mut(), poll_flush, context)
|
|
}
|
|
fn poll_shutdown(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
dispatch!(self.get_mut(), poll_shutdown, context)
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn connect(
|
|
host: &str,
|
|
port: i64,
|
|
tls: Option<&crate::tls::ClientTls>,
|
|
) -> Result<Transport, String> {
|
|
let stream = crate::net::connect_stream(host, port).await?;
|
|
match tls {
|
|
Some(config) => crate::tls::connect(stream, host, config).await,
|
|
None => Ok(Transport::Tcp(stream)),
|
|
}
|
|
}
|