1 files changed, 68 insertions, 0 deletions
diff --git a/src/writer/loggable.rs b/src/writer/loggable.rs
new file mode 100644
index 0000000..1b9f64a
--- /dev/null
+++ b/src/writer/loggable.rs
@@ -0,0 +1,68 @@
+use std::{fmt::Display, mem, pin::pin, task::Poll};
+
+use futures::ready;
+use pin_project::pin_project;
+pub use tokio::io::AsyncWrite;
+
+#[pin_project]
+/// Wrapper struct for logging writes to `AsyncWrite` implementors.
+#[derive(Debug)]
+pub struct Loggable<W> {
+    log_buffer: Vec<u8>,
+    #[pin]
+    writer: W,
+}
+
+impl<W> Loggable<W> {
+    pub fn new(writer: W) -> Self {
+        Self {
+            log_buffer: Vec::new(),
+            writer,
+        }
+    }
+
+    pub fn into_inner(self) -> W {
+        self.writer
+    }
+
+    pub fn take_log(&mut self) -> Vec<u8> {
+        let log: Vec<u8> = mem::replace(&mut self.log_buffer, Vec::new());
+        log
+    }
+}
+
+impl<W: AsyncWrite + Unpin + Send> Display for Loggable<W> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let str = String::from_utf8_lossy(&self.log_buffer);
+        write!(f, "{str}")
+    }
+}
+
+impl<W: AsyncWrite + Unpin + Send> AsyncWrite for Loggable<W> {
+    fn poll_write(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> std::task::Poll<Result<usize, std::io::Error>> {
+        let this = self.as_mut().project();
+        let ready = ready!(this.writer.poll_write(cx, buf));
+        self.log_buffer.extend_from_slice(buf);
+        Poll::Ready(ready)
+    }
+
+    fn poll_flush(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), std::io::Error>> {
+        let this = self.project();
+        Poll::Ready(ready!(this.writer.poll_flush(cx)))
+    }
+
+    fn poll_shutdown(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), std::io::Error>> {
+        let this = self.project();
+        Poll::Ready(ready!(this.writer.poll_shutdown(cx)))
+    }
+}