crypto handshake and encryption

src/ws.rs

@@ -1,26 +1,31 @@
 use std::borrow::Cow;
 
-use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
+use axum::extract::ws::{Message, WebSocket};
 use futures_util::{
     SinkExt, StreamExt,
     stream::{SplitSink, SplitStream},
 };
-use tokio::sync::{Mutex, MutexGuard};
+use tokio::sync::Mutex;
 
-use crate::protocol::{ClientMethod, ServerMethod};
+use crate::{
+    crypto::SessionCipher,
+    protocol::{ClientMethod, ServerMethod},
+};
 
 pub struct EnclaveWebSocket {
     tx: Mutex<SplitSink<WebSocket, Message>>,
     rx: Mutex<SplitStream<WebSocket>>,
+    cihper: Mutex<SessionCipher>,
 }
 
 impl EnclaveWebSocket {
-    pub fn new(ws: WebSocket) -> Self {
+    pub fn new(ws: WebSocket, cipher: SessionCipher) -> Self {
         let (tx, rx) = ws.split();
 
         Self {
             tx: Mutex::new(tx),
             rx: Mutex::new(rx),
+            cihper: Mutex::new(cipher),
         }
     }
 
@@ -39,6 +44,23 @@ impl EnclaveWebSocket {
                 }
             },
 
+            Some(Message::Binary(encrypted)) => {
+                let text = String::from_utf8(self.cihper.lock().await.decrypt(&encrypted)?)?;
+
+                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?;
 
@@ -52,18 +74,16 @@ impl EnclaveWebSocket {
     }
 
     pub async fn send(&self, message: &ClientMethod) -> anyhow::Result<()> {
+        let text = serde_json::to_string(message)?;
+
+        let encrypted = self.cihper.lock().await.encrypt(text.as_bytes())?;
+
         self.tx
             .lock()
             .await
-            .send(Message::Text(Utf8Bytes::from(serde_json::to_string(
-                message,
-            )?)))
+            .send(Message::Binary(encrypted.into()))
             .await?;
 
         Ok(())
     }
-
-    pub async fn tx(&self) -> MutexGuard<'_, SplitSink<WebSocket, Message>> {
-        self.tx.lock().await
-    }
 }