use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio_native_tls::TlsStream;
pub enum Stream {
Tcp(TcpStream),
Tls(TlsStream<TcpStream>),
}
impl Stream {
pub async fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
match self {
Stream::Tcp(stream) => stream.read(buf).await,
Stream::Tls(stream) => stream.read(buf).await,
}
}
pub async fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
match self {
Stream::Tcp(stream) => stream.write_all(buf).await,
Stream::Tls(stream) => stream.write_all(buf).await,
}
}
pub async fn flush(&mut self) -> std::io::Result<()> {
match self {
Stream::Tcp(stream) => stream.flush().await,
Stream::Tls(stream) => stream.flush().await,
}
}
pub async fn shutdown(&mut self) -> std::io::Result<()> {
match self {
Stream::Tcp(stream) => stream.shutdown().await,
Stream::Tls(stream) => stream.shutdown().await,
}
}
}