use async_native_tls::TlsStream;
use std::{
pin::Pin,
task::{Context, Poll},
};
use smol::{
io::{self, AsyncRead, AsyncWrite},
net::TcpStream,
};
pub enum SmolStream {
Plain(TcpStream),
Tls(TlsStream<TcpStream>),
}
impl AsyncRead for SmolStream {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
match &mut *self {
SmolStream::Plain(stream) => Pin::new(stream).poll_read(cx, buf),
SmolStream::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
}
}
}
impl AsyncWrite for SmolStream {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
match &mut *self {
SmolStream::Plain(stream) => Pin::new(stream).poll_write(cx, buf),
SmolStream::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
}
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
match &mut *self {
SmolStream::Plain(stream) => Pin::new(stream).poll_close(cx),
SmolStream::Tls(stream) => Pin::new(stream).poll_close(cx),
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
match &mut *self {
SmolStream::Plain(stream) => Pin::new(stream).poll_flush(cx),
SmolStream::Tls(stream) => Pin::new(stream).poll_flush(cx),
}
}
}