atm0s_reverse_proxy_protocol/
stream.rs

1use std::{
2    io,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use tokio::io::{AsyncRead, AsyncWrite};
8
9pub struct TunnelStream<R: AsyncRead, W: AsyncWrite> {
10    read: R,
11    write: W,
12}
13
14impl<R: AsyncRead, W: AsyncWrite> TunnelStream<R, W> {
15    pub fn new(read: R, write: W) -> Self {
16        Self { read, write }
17    }
18}
19
20impl<R: AsyncRead + Unpin, W: AsyncWrite + Unpin> AsyncRead for TunnelStream<R, W> {
21    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut tokio::io::ReadBuf<'_>) -> Poll<io::Result<()>> {
22        Pin::new(&mut self.get_mut().read).poll_read(cx, buf)
23    }
24}
25
26impl<R: AsyncRead + Unpin, W: AsyncWrite + Unpin> AsyncWrite for TunnelStream<R, W> {
27    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
28        Pin::new(&mut self.get_mut().write).poll_write(cx, buf)
29    }
30
31    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
32        Pin::new(&mut self.get_mut().write).poll_flush(cx)
33    }
34
35    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
36        Pin::new(&mut self.get_mut().write).poll_shutdown(cx)
37    }
38}