capnweb_transport/
websocket.rs

1use crate::{RpcTransport, TransportError};
2use async_trait::async_trait;
3use capnweb_core::{decode_message, encode_message, Message};
4use futures::{SinkExt, StreamExt};
5use std::sync::Arc;
6use tokio::sync::{mpsc, Mutex};
7use tokio_tungstenite::{tungstenite, WebSocketStream};
8use tungstenite::protocol::Message as WsMessage;
9
10/// WebSocket transport implementation
11pub struct WebSocketTransport<S> {
12    /// The WebSocket stream
13    stream: Arc<Mutex<WebSocketStream<S>>>,
14    /// Channel for sending messages
15    tx: mpsc::UnboundedSender<Message>,
16    /// Channel for receiving messages
17    rx: Arc<Mutex<mpsc::UnboundedReceiver<Message>>>,
18}
19
20impl<S> WebSocketTransport<S>
21where
22    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Sync + Unpin + 'static,
23{
24    /// Create a new WebSocket transport from a stream
25    pub fn new(stream: WebSocketStream<S>) -> Self {
26        let (tx, mut rx_internal) = mpsc::unbounded_channel();
27        let (tx_internal, rx) = mpsc::unbounded_channel();
28
29        let stream = Arc::new(Mutex::new(stream));
30        let stream_clone = stream.clone();
31
32        // Spawn task to handle incoming messages
33        tokio::spawn(async move {
34            loop {
35                let mut stream = stream_clone.lock().await;
36                match stream.next().await {
37                    Some(Ok(msg)) => {
38                        if let WsMessage::Binary(data) = msg {
39                            // Decode the message
40                            match decode_message(&data) {
41                                Ok(message) => {
42                                    if tx_internal.send(message).is_err() {
43                                        break;
44                                    }
45                                }
46                                Err(e) => {
47                                    eprintln!("Failed to decode message: {}", e);
48                                }
49                            }
50                        } else if let WsMessage::Close(_) = msg {
51                            break;
52                        }
53                    }
54                    Some(Err(e)) => {
55                        eprintln!("WebSocket error: {}", e);
56                        break;
57                    }
58                    None => break,
59                }
60            }
61        });
62
63        // Spawn task to handle outgoing messages
64        let stream_clone2 = stream.clone();
65        tokio::spawn(async move {
66            while let Some(message) = rx_internal.recv().await {
67                let encoded = match encode_message(&message) {
68                    Ok(data) => data,
69                    Err(e) => {
70                        eprintln!("Failed to encode message: {}", e);
71                        continue;
72                    }
73                };
74
75                let mut stream = stream_clone2.lock().await;
76                if let Err(e) = stream.send(WsMessage::Binary(encoded)).await {
77                    eprintln!("Failed to send WebSocket message: {}", e);
78                    break;
79                }
80            }
81        });
82
83        Self {
84            stream,
85            tx,
86            rx: Arc::new(Mutex::new(rx)),
87        }
88    }
89}
90
91#[async_trait]
92impl<S> RpcTransport for WebSocketTransport<S>
93where
94    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Sync + Unpin + 'static,
95{
96    async fn send(&mut self, message: Message) -> Result<(), TransportError> {
97        self.tx
98            .send(message)
99            .map_err(|_| TransportError::ConnectionClosed)
100    }
101
102    async fn recv(&mut self) -> Result<Option<Message>, TransportError> {
103        let mut rx = self.rx.lock().await;
104        Ok(rx.recv().await)
105    }
106
107    async fn close(&mut self) -> Result<(), TransportError> {
108        let mut stream = self.stream.lock().await;
109        stream
110            .close(None)
111            .await
112            .map_err(|e| TransportError::Protocol(format!("Failed to close WebSocket: {}", e)))
113    }
114}
115
116/// WebSocket client transport
117pub struct WebSocketClient {
118    url: String,
119}
120
121impl WebSocketClient {
122    /// Create a new WebSocket client
123    pub fn new(url: impl Into<String>) -> Self {
124        Self { url: url.into() }
125    }
126
127    /// Connect to the WebSocket server
128    pub async fn connect(
129        &self,
130    ) -> Result<
131        WebSocketTransport<
132            impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Sync + Unpin + 'static,
133        >,
134        TransportError,
135    > {
136        let (stream, _) = tokio_tungstenite::connect_async(&self.url)
137            .await
138            .map_err(|e| TransportError::Protocol(format!("Failed to connect: {}", e)))?;
139
140        Ok(WebSocketTransport::new(stream))
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use tokio::net::TcpStream;
147    use tokio_tungstenite::MaybeTlsStream;
148
149    #[tokio::test]
150    async fn test_websocket_transport() {
151        // For now, just test the WebSocket transport can be created
152        // Full integration test would require more complex setup
153
154        // This is a basic smoke test to ensure compilation
155        let _unused_import: Option<MaybeTlsStream<TcpStream>> = None;
156    }
157}
capnweb_transport/websocket.rs

capnweb_transport/
websocket.rs
