initial

rust/src/transport.rs

@@ -0,0 +1,55 @@
+use anyhow::Result;
+use futures_util::{SinkExt, StreamExt};
+use tokio::net::TcpStream;
+use tokio_tungstenite::{
+    connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream,
+};
+use tracing::info;
+
+/// A WebSocket connection (either client or server side).
+pub type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
+
+/// Connect to a WebSocket server as a client.
+pub async fn connect_to_server(url: &str) -> Result<WsStream> {
+    info!("Connecting to WebSocket server: {}", url);
+    let (ws_stream, response) = connect_async(url).await?;
+    info!("WebSocket connected, status: {}", response.status());
+    Ok(ws_stream)
+}
+
+/// Send a binary message over the WebSocket.
+pub async fn send_binary(ws: &mut WsStream, data: Vec<u8>) -> Result<()> {
+    ws.send(Message::Binary(data.into())).await?;
+    Ok(())
+}
+
+/// Receive the next binary message from the WebSocket.
+/// Returns None if the connection is closed.
+pub async fn recv_binary(ws: &mut WsStream) -> Result<Option<Vec<u8>>> {
+    loop {
+        match ws.next().await {
+            Some(Ok(Message::Binary(data))) => return Ok(Some(data.to_vec())),
+            Some(Ok(Message::Close(_))) => return Ok(None),
+            Some(Ok(Message::Ping(data))) => {
+                ws.send(Message::Pong(data)).await?;
+            }
+            Some(Ok(_)) => continue,
+            Some(Err(e)) => return Err(anyhow::anyhow!("WebSocket error: {}", e)),
+            None => return Ok(None),
+        }
+    }
+}
+
+/// Send a close frame.
+pub async fn close(ws: &mut WsStream) -> Result<()> {
+    ws.close(None).await?;
+    Ok(())
+}
+
+/// WebSocket server acceptor — accepts a TcpStream and performs the WebSocket upgrade.
+pub async fn accept_connection(
+    stream: TcpStream,
+) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>> {
+    let ws = tokio_tungstenite::accept_async(MaybeTlsStream::Plain(stream)).await?;
+    Ok(ws)
+}