66 lines
1.6 KiB
Rust
66 lines
1.6 KiB
Rust
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(())
|
|
}
|
|
}
|