arampoire/nut-rs
1
0
Fork 0
mirror of https://github.com/arampoire/nut-rs.git synced 2025-12-01 00:30:23 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add tokio buffered stream

Browse source
This commit is contained in:
Aram 🍐 2021-08-04 22:06:15 -04:00
parent 07034d2cec
commit 97e3731df2
3 changed files with 57 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

56
rups/src/tokio/stream.rs
View file

@ -1,13 +1,19 @@
use crate::Error;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, BufReader, ReadBuf};
use tokio::net::TcpStream;
/// A wrapper for various Tokio stream types.
pub enum ConnectionStream {
/// A plain TCP stream.
Plain(TcpStream),
Tcp(TcpStream),
/// A stream wrapped with `BufReader`.
///
/// Use `.buffered()` to wrap any stream with `BufReader`.
/// It can then be un-wrapped with `.unbuffered()`.
Buffered(Box<BufReader<ConnectionStream>>),
/// A client stream wrapped with SSL using `rustls`.
#[cfg(feature = "async-ssl")]
@ -50,10 +56,14 @@ impl AsyncRead for ConnectionStream {
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
match self.get_mut() {
Self::Plain(stream) => {
Self::Tcp(stream) => {
let pinned = Pin::new(stream);
pinned.poll_read(cx, buf)
}
Self::Buffered(reader) => {
let pinned = Pin::new(reader.get_mut());
pinned.poll_read(cx, buf)
}
#[cfg(feature = "async-ssl")]
Self::SslClient(stream) => {
let pinned = Pin::new(stream);
@ -68,6 +78,28 @@ impl AsyncRead for ConnectionStream {
}
}
impl AsyncBufRead for ConnectionStream {
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<&[u8]>> {
match self.get_mut() {
Self::Buffered(reader) => {
let pinned = Pin::new(reader.get_mut());
pinned.poll_fill_buf(cx)
}
_ => core::task::Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"Stream is not buffered",
))),
}
}
fn consume(self: Pin<&mut Self>, amt: usize) {
if let Self::Buffered(reader) = self.get_mut() {
let pinned = Pin::new(reader.get_mut());
pinned.consume(amt)
}
}
}
impl AsyncWrite for ConnectionStream {
fn poll_write(
self: Pin<&mut Self>,
@ -75,10 +107,14 @@ impl AsyncWrite for ConnectionStream {
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
match self.get_mut() {
Self::Plain(stream) => {
Self::Tcp(stream) => {
let pinned = Pin::new(stream);
pinned.poll_write(cx, buf)
}
Self::Buffered(reader) => {
let pinned = Pin::new(reader.get_mut());
pinned.poll_write(cx, buf)
}
#[cfg(feature = "async-ssl")]
Self::SslClient(stream) => {
let pinned = Pin::new(stream);
@ -94,10 +130,14 @@ impl AsyncWrite for ConnectionStream {
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
Self::Plain(stream) => {
Self::Tcp(stream) => {
let pinned = Pin::new(stream);
pinned.poll_flush(cx)
}
Self::Buffered(reader) => {
let pinned = Pin::new(reader.get_mut());
pinned.poll_flush(cx)
}
#[cfg(feature = "async-ssl")]
Self::SslClient(stream) => {
let pinned = Pin::new(stream);
@ -113,10 +153,14 @@ impl AsyncWrite for ConnectionStream {
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
Self::Plain(stream) => {
Self::Tcp(stream) => {
let pinned = Pin::new(stream);
pinned.poll_shutdown(cx)
}
Self::Buffered(reader) => {
let pinned = Pin::new(reader.get_mut());
pinned.poll_shutdown(cx)
}
#[cfg(feature = "async-ssl")]
Self::SslClient(stream) => {
let pinned = Pin::new(stream);