2 files changed, 104 insertions, 0 deletions
diff --git a/src/writer/endable.rs b/src/writer/endable.rs
new file mode 100644
index 0000000..6d842f3
--- /dev/null
+++ b/src/writer/endable.rs
@@ -0,0 +1,37 @@
+use crate::Error;
+
+#[derive(Debug)]
+pub struct Endable<T> {
+    inner: T,
+    ended: bool,
+}
+
+impl<T> Endable<T> {
+    pub fn new(inner: T) -> Self {
+        Self {
+            inner,
+            ended: false,
+        }
+    }
+
+    pub fn end(&mut self) {
+        self.ended = true;
+    }
+
+    pub fn into_inner(self) -> T {
+        self.inner
+    }
+
+    #[inline(always)]
+    pub fn try_as_mut(&mut self) -> Result<&mut T, Error> {
+        if self.ended {
+            Err(Error::RootElementEnded)
+        } else {
+            Ok(&mut self.inner)
+        }
+    }
+
+    pub fn ignore_end(&mut self) -> &mut T {
+        &mut self.inner
+    }
+}
diff --git a/src/writer/loggable.rs b/src/writer/loggable.rs
new file mode 100644
index 0000000..dd69668
--- /dev/null
+++ b/src/writer/loggable.rs
@@ -0,0 +1,67 @@
+use std::{fmt::Display, mem, pin::pin, task::Poll};
+
+use futures::ready;
+use pin_project::pin_project;
+pub use tokio::io::AsyncWrite;
+
+#[pin_project]
+#[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 = str::from_utf8(&self.log_buffer).unwrap_or("buffer to string conversion failed");
+        f.write_str(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)))
+    }
+}