From 0a3683e2fa750a83f8b56d450a8e9f12adb72dd1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jind=C5=99ich=20Moravec?= <xmorave2@fi.muni.cz>
Date: Mon, 11 Dec 2023 16:49:30 +0100
Subject: [PATCH] feat(proto): add generic proto writer

---
 proto/src/lib.rs                |  1 +
 proto/src/writer/mod.rs         |  3 +++
 proto/src/writer/oneway.rs      | 31 ++++++++++++++++++++++++++++++
 proto/src/writer/protowriter.rs | 34 +++++++++++++++++++++++++++++++++
 4 files changed, 69 insertions(+)
 create mode 100644 proto/src/writer/mod.rs
 create mode 100644 proto/src/writer/oneway.rs
 create mode 100644 proto/src/writer/protowriter.rs

diff --git a/proto/src/lib.rs b/proto/src/lib.rs
index e216a50..c65ece0 100644
--- a/proto/src/lib.rs
+++ b/proto/src/lib.rs
@@ -1 +1,2 @@
 pub mod message;
+pub mod writer;
diff --git a/proto/src/writer/mod.rs b/proto/src/writer/mod.rs
new file mode 100644
index 0000000..fac68e8
--- /dev/null
+++ b/proto/src/writer/mod.rs
@@ -0,0 +1,3 @@
+
+pub mod oneway;
+pub mod protowriter;
diff --git a/proto/src/writer/oneway.rs b/proto/src/writer/oneway.rs
new file mode 100644
index 0000000..1649fd8
--- /dev/null
+++ b/proto/src/writer/oneway.rs
@@ -0,0 +1,31 @@
+use crate::message::proto_message::ProtoMessage;
+use crate::writer::protowriter::ProtoWriter;
+use async_trait::async_trait;
+use tokio::io::{AsyncWrite, AsyncWriteExt};
+
+#[async_trait]
+pub trait OneWayProtoWriter<T>
+where
+    T: ProtoMessage,
+{
+    async fn write_proto(&mut self, message: T) -> anyhow::Result<()>;
+}
+
+#[async_trait]
+impl<W, T> OneWayProtoWriter<T> for ProtoWriter<W>
+where
+    W: AsyncWrite + Unpin + Send,
+    T: ProtoMessage + Send + 'static,
+{
+    async fn write_proto(&mut self, message: T) -> anyhow::Result<()> {
+        let variant = message.variant();
+        let mut data = message.serialize()?;
+        let length = data.len() as i32 + 4;
+
+        self.inner.write_u8(variant).await?;
+        self.inner.write_i32(length).await?;
+        self.inner.write_all(&mut data).await?;
+
+        Ok(())
+    }
+}
diff --git a/proto/src/writer/protowriter.rs b/proto/src/writer/protowriter.rs
new file mode 100644
index 0000000..848a727
--- /dev/null
+++ b/proto/src/writer/protowriter.rs
@@ -0,0 +1,34 @@
+use async_trait::async_trait;
+use tokio::io::{AsyncWrite, AsyncWriteExt};
+
+pub struct ProtoWriter<W>
+where
+    W: AsyncWrite + Unpin + Send,
+{
+    pub(super) inner: W,
+}
+
+impl<W> ProtoWriter<W>
+where
+    W: AsyncWrite + Unpin + Send,
+{
+    pub fn new(writer: W) -> ProtoWriter<W> {
+        ProtoWriter { inner: writer }
+    }
+}
+
+#[async_trait]
+pub trait ProtoFlush {
+    async fn flush(&mut self) -> anyhow::Result<()>;
+}
+
+#[async_trait]
+impl<W> ProtoFlush for ProtoWriter<W>
+where
+    W: AsyncWrite + Unpin + Send,
+{
+    async fn flush(&mut self) -> anyhow::Result<()> {
+        self.inner.flush().await?;
+        Ok(())
+    }
+}