common/
io.rs

1use std::io;
2
3use async_trait::async_trait;
4use tokio::io::{AsyncReadExt, AsyncWriteExt};
5
6#[async_trait]
7pub trait AsyncReader: Send {
8    async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> io::Result<usize>;
9    async fn read_exact<'a>(&'a mut self, buf: &'a mut [u8]) -> io::Result<usize>;
10}
11
12pub(crate) struct AsyncReadWrapper<R: AsyncReadExt + Unpin> {
13    reader: R,
14}
15
16impl<R: AsyncReadExt + Unpin> AsyncReadWrapper<R> {
17    pub(crate) fn new(reader: R) -> Self {
18        Self { reader }
19    }
20}
21
22#[async_trait]
23impl<R: AsyncReadExt + Unpin + Send> AsyncReader for AsyncReadWrapper<R> {
24    async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> io::Result<usize>
25    where
26        Self: Unpin,
27    {
28        self.reader.read(buf).await
29    }
30
31    async fn read_exact<'a>(&'a mut self, buf: &'a mut [u8]) -> io::Result<usize>
32    where
33        Self: Unpin,
34    {
35        self.reader.read_exact(buf).await
36    }
37}
38
39#[async_trait]
40pub trait AsyncWriter: Send {
41    async fn write<'a>(&'a mut self, src: &'a [u8]) -> io::Result<usize>;
42    async fn shutdown(&mut self) -> io::Result<()>;
43}
44
45pub(crate) struct AsyncWriteWrapper<W: AsyncWriteExt + Unpin> {
46    writer: W,
47}
48
49impl<W: AsyncWriteExt + Unpin> AsyncWriteWrapper<W> {
50    pub(crate) fn new(writer: W) -> Self {
51        Self { writer }
52    }
53}
54
55#[async_trait]
56impl<W: AsyncWriteExt + Unpin + Send> AsyncWriter for AsyncWriteWrapper<W> {
57    async fn write<'a>(&'a mut self, src: &'a [u8]) -> io::Result<usize>
58    where
59        Self: Unpin,
60    {
61        self.writer.write(src).await
62    }
63
64    async fn shutdown(&mut self) -> io::Result<()> {
65        self.writer.shutdown().await
66    }
67}
68
69pub async fn copy(
70    from: &mut Box<dyn AsyncReader>,
71    to: &mut Box<dyn AsyncWriter>,
72    debug_log_message: &str,
73) -> tokio::io::Result<()> {
74    let mut data = vec![0; 4096];
75    loop {
76        let size = from.read(&mut data).await?;
77        if size == 0 {
78            break;
79        }
80        log::debug!("{} {:?}", debug_log_message, &data[..size]);
81        to.write(&data[..size]).await?;
82    }
83    to.shutdown().await
84}
85
86pub struct Stream {
87    pub reader: Box<dyn AsyncReader>,
88    pub writer: Box<dyn AsyncWriter>,
89}
90
91impl Stream {
92    pub fn new<
93        R: AsyncReadExt + Unpin + Send + 'static,
94        W: AsyncWriteExt + Unpin + Send + 'static,
95    >(
96        reader: R,
97        writer: W,
98    ) -> Self {
99        Self {
100            reader: Box::new(AsyncReadWrapper::new(reader)),
101            writer: Box::new(AsyncWriteWrapper::new(writer)),
102        }
103    }
104}
common/io.rs

common/
io.rs
