Better websocket

src/ws.rs

@@ -0,0 +1,65 @@
+use std::borrow::Cow;
+
+use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
+use futures_util::{
+    SinkExt, StreamExt,
+    stream::{SplitSink, SplitStream},
+};
+use tokio::sync::Mutex;
+
+use crate::protocol::{ClientMethod, ServerMethod};
+
+pub struct EnclaveWebSocket {
+    tx: Mutex<SplitSink<WebSocket, Message>>,
+    rx: Mutex<SplitStream<WebSocket>>,
+}
+
+impl EnclaveWebSocket {
+    pub fn new(ws: WebSocket) -> Self {
+        let (tx, rx) = ws.split();
+
+        Self {
+            tx: Mutex::new(tx),
+            rx: Mutex::new(rx),
+        }
+    }
+
+    pub async fn read(&self) -> anyhow::Result<Option<ServerMethod>> {
+        match self.rx.lock().await.next().await.transpose()? {
+            Some(Message::Text(text)) => match serde_json::from_str(&text.to_string()) {
+                Ok(msg) => Ok(Some(msg)),
+
+                Err(e) => {
+                    self.send(&ClientMethod::Error {
+                        error: Cow::Owned(format!("Unable to parse message: {e}")),
+                    })
+                    .await?;
+
+                    Ok(None)
+                }
+            },
+
+            Some(Message::Ping(v)) => {
+                self.tx.lock().await.send(Message::Pong(v)).await?;
+
+                Ok(None)
+            }
+
+            Some(_) => Ok(None),
+
+            None => Ok(None),
+        }
+    }
+
+    pub async fn send(&self, message: &ClientMethod) -> anyhow::Result<()> {
+        self.tx
+            .lock()
+            .await
+            .send(Message::Text(Utf8Bytes::from(serde_json::to_string(
+                message,
+            )?)))
+            .await?;
+
+        Ok(())
+    }
+}