Skip to main content

zenwave/
websocket.rs

1pub use http_kit::ws::*;
2
3use http_kit::{HttpError, StatusCode};
4use serde::Serialize;
5
6/// Errors returned by websocket operations.
7#[derive(Debug, thiserror::Error)]
8pub enum WebSocketError {
9    /// Failed to encode a payload for transmission.
10    #[error("Fail to encode payload: {0}")]
11    FailToEncodePayload(serde_json::Error),
12
13    /// Unsupported websocket URI scheme encountered.
14    #[error("Unsupported websocket scheme: {0}")]
15    UnsupportedScheme(String),
16
17    /// Provided websocket URI was invalid.
18    #[error("Invalid URI: {0}")]
19    InvalidUri(#[from] url::ParseError),
20
21    /// The URI parsed but names no reachable target.
22    #[error("Invalid websocket target: {0}")]
23    InvalidTarget(&'static str),
24
25    /// Connecting through the [`crate::Transport`] failed (DNS, TCP, TLS, proxy).
26    #[error("{0}")]
27    Transport(#[source] Box<crate::Error>),
28
29    /// Underlying websocket connection failed.
30    #[error("Connection failed: {0}")]
31    ConnectionFailed(#[source] Box<dyn std::error::Error + Send + Sync>),
32}
33
34impl HttpError for WebSocketError {
35    fn status(&self) -> StatusCode {
36        match self {
37            Self::ConnectionFailed(_) => StatusCode::BAD_GATEWAY,
38            Self::Transport(error) => error.status(),
39            _ => StatusCode::INTERNAL_SERVER_ERROR,
40        }
41    }
42}
43
44// Convert WebSocketError to unified zenwave::Error
45impl From<WebSocketError> for crate::Error {
46    fn from(err: WebSocketError) -> Self {
47        use crate::error::WebSocketErrorKind;
48
49        match err {
50            WebSocketError::FailToEncodePayload(e) => {
51                Self::WebSocket(WebSocketErrorKind::EncodeFailed(e))
52            }
53            WebSocketError::UnsupportedScheme(s) => {
54                Self::WebSocket(WebSocketErrorKind::UnsupportedScheme(s))
55            }
56            WebSocketError::InvalidUri(e) => Self::InvalidUri(e.to_string()),
57            WebSocketError::InvalidTarget(message) => Self::InvalidUri(message.to_string()),
58            WebSocketError::Transport(error) => *error,
59            WebSocketError::ConnectionFailed(e) => {
60                Self::WebSocket(WebSocketErrorKind::ConnectionFailed(e.to_string()))
61            }
62        }
63    }
64}
65
66/// Configuration applied when establishing a websocket connection.
67#[derive(Clone, Debug)]
68#[non_exhaustive]
69pub struct WebSocketConfig {
70    /// Maximum incoming websocket message size in bytes.
71    /// `None` means no limit.
72    pub max_message_size: Option<usize>,
73
74    /// Maximum incoming websocket frame size in bytes.
75    /// `None` means no limit.
76    pub max_frame_size: Option<usize>,
77}
78
79const DEFAULT_MAX_MESSAGE_SIZE: Option<usize> = Some(64 << 20);
80const DEFAULT_MAX_FRAME_SIZE: Option<usize> = Some(16 << 20);
81
82impl Default for WebSocketConfig {
83    fn default() -> Self {
84        Self {
85            max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
86            max_frame_size: DEFAULT_MAX_FRAME_SIZE,
87        }
88    }
89}
90
91impl WebSocketConfig {
92    /// Override the maximum incoming websocket message size in bytes.
93    ///
94    /// `None` means no limit.
95    ///
96    /// Defaults to 64 MiB.
97    #[must_use]
98    pub const fn with_max_message_size(mut self, max_message_size: Option<usize>) -> Self {
99        self.max_message_size = max_message_size;
100        self
101    }
102
103    /// Override the maximum incoming websocket frame size in bytes.
104    ///
105    /// `None` means no limit.
106    ///
107    /// Defaults to 16 MiB.
108    #[must_use]
109    pub const fn with_max_frame_size(mut self, max_frame_size: Option<usize>) -> Self {
110        self.max_frame_size = max_frame_size;
111        self
112    }
113}
114
115#[allow(clippy::result_large_err)]
116fn serialize_payload<T>(value: &T) -> Result<String, WebSocketError>
117where
118    T: Serialize,
119{
120    serde_json::to_string(value).map_err(WebSocketError::FailToEncodePayload)
121}
122
123#[cfg(not(target_arch = "wasm32"))]
124mod native {
125    use async_lock::Mutex;
126    use async_tungstenite::{
127        WebSocketReceiver as AsyncReceiver, WebSocketSender as AsyncSender, WebSocketStream,
128        client_async_with_config,
129        tungstenite::{
130            Message as TungsteniteMessage, Utf8Bytes,
131            protocol::WebSocketConfig as TungsteniteConfig,
132        },
133    };
134    use futures_util::StreamExt;
135    use http_kit::utils::{ByteStr, Bytes};
136    use std::{fmt, sync::Arc};
137    use url::Url;
138
139    use crate::{
140        Transport,
141        transport::{
142            connect::{Connection, Protocols, Target, Via, connect as connect_stream},
143            stream::Stream,
144        },
145    };
146
147    use super::{WebSocketConfig, WebSocketError, WebSocketMessage, serialize_payload};
148
149    type NativeSocket = WebSocketStream<Stream>;
150    type NativeSender = AsyncSender<Stream>;
151    type NativeReceiver = AsyncReceiver<Stream>;
152
153    #[derive(Debug)]
154    struct SharedSocket {
155        sender: Mutex<NativeSender>,
156        receiver: Mutex<NativeReceiver>,
157    }
158
159    /// A websocket connection backed by async-io + Tungstenite.
160    pub struct WebSocket {
161        sender: WebSocketSender,
162        receiver: WebSocketReceiver,
163    }
164
165    impl fmt::Debug for WebSocket {
166        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
167            f.debug_struct("WebSocket").finish()
168        }
169    }
170
171    /// Sending half of a websocket connection.
172    pub struct WebSocketSender {
173        inner: Arc<SharedSocket>,
174    }
175
176    impl fmt::Debug for WebSocketSender {
177        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178            f.debug_struct("WebSocketSender").finish()
179        }
180    }
181
182    impl Clone for WebSocketSender {
183        fn clone(&self) -> Self {
184            Self {
185                inner: Arc::clone(&self.inner),
186            }
187        }
188    }
189
190    /// Receiving half of a websocket connection.
191    pub struct WebSocketReceiver {
192        inner: Arc<SharedSocket>,
193    }
194
195    impl fmt::Debug for WebSocketReceiver {
196        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
197            f.debug_struct("WebSocketReceiver").finish()
198        }
199    }
200
201    /// Establish a websocket connection to the provided URI through the
202    /// system transport with default options.
203    ///
204    /// # Errors
205    ///
206    /// Returns an error if the URI is invalid or the connection attempt fails.
207    pub async fn connect(uri: impl AsRef<str>) -> Result<WebSocket, WebSocketError> {
208        connect_with(uri, &Transport::system(), WebSocketConfig::default()).await
209    }
210
211    /// Establish a websocket connection to the provided URI through `transport`.
212    ///
213    /// # Errors
214    ///
215    /// Returns an error if the URI is invalid or the connection attempt fails.
216    pub async fn connect_with(
217        uri: impl AsRef<str>,
218        transport: &Transport,
219        websocket_config: WebSocketConfig,
220    ) -> Result<WebSocket, WebSocketError> {
221        let url = Url::parse(uri.as_ref())?;
222        let tls = match url.scheme() {
223            "ws" => false,
224            "wss" => true,
225            other => return Err(WebSocketError::UnsupportedScheme(other.to_string())),
226        };
227        let host = url
228            .host_str()
229            .ok_or_else(|| WebSocketError::InvalidTarget("websocket URI is missing a host"))?;
230        let port = url
231            .port_or_known_default()
232            .ok_or_else(|| WebSocketError::InvalidTarget("websocket URI does not imply a port"))?;
233        // The websocket handshake is an HTTP/1.1 upgrade, so the TLS layer
234        // must not offer h2.
235        let Connection { stream, via, .. } = connect_stream(
236            transport,
237            Target {
238                host,
239                port,
240                tls,
241                tunnel_plaintext: true,
242                protocols: Protocols::Http1,
243            },
244        )
245        .await
246        .map_err(|error| WebSocketError::Transport(Box::new(error)))?;
247        debug_assert!(
248            matches!(via, Via::Direct),
249            "websocket targets are always tunnelled, never proxied in absolute form"
250        );
251
252        let mut config = TungsteniteConfig::default();
253        config.max_message_size = websocket_config.max_message_size;
254        config.max_frame_size = websocket_config.max_frame_size;
255        let request: String = url.into();
256        let (ws_stream, _) = client_async_with_config(request, stream, Some(config))
257            .await
258            .map_err(|e| WebSocketError::ConnectionFailed(Box::new(e)))?;
259
260        Ok(WebSocket::from_socket(ws_stream))
261    }
262
263    impl WebSocket {
264        fn from_socket(socket: NativeSocket) -> Self {
265            let (sender, receiver) = socket.split();
266            let shared = Arc::new(SharedSocket {
267                sender: Mutex::new(sender),
268                receiver: Mutex::new(receiver),
269            });
270
271            Self {
272                sender: WebSocketSender {
273                    inner: Arc::clone(&shared),
274                },
275                receiver: WebSocketReceiver { inner: shared },
276            }
277        }
278
279        /// Send a websocket message serialized as JSON.
280        ///
281        /// # Errors
282        ///
283        /// Returns an error if serialization fails or when the underlying socket cannot
284        /// write the resulting frame.
285        pub async fn send<T>(&self, value: T) -> Result<(), WebSocketError>
286        where
287            T: serde::Serialize,
288        {
289            let payload = serialize_payload(&value)?;
290            self.send_text(payload).await
291        }
292
293        /// Send a text websocket message.
294        ///
295        /// # Errors
296        ///
297        /// Returns an error when the underlying socket cannot write the frame.
298        pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<(), WebSocketError> {
299            self.sender.send_text(text).await
300        }
301
302        /// Send a binary websocket message.
303        ///
304        /// # Errors
305        ///
306        /// Returns an error when the underlying socket cannot write the frame.
307        pub async fn send_binary(&self, bytes: impl Into<Bytes>) -> Result<(), WebSocketError> {
308            self.sender.send_binary(bytes).await
309        }
310
311        /// Receive the next websocket message.
312        ///
313        /// # Errors
314        ///
315        /// Returns an error when the underlying socket cannot read the next frame.
316        pub async fn recv(&self) -> Result<Option<WebSocketMessage>, WebSocketError> {
317            self.receiver.recv().await
318        }
319
320        /// Close the websocket connection gracefully.
321        ///
322        /// # Errors
323        ///
324        /// Returns an error when the close frame cannot be sent.
325        pub async fn close(self) -> Result<(), WebSocketError> {
326            self.sender.close().await
327        }
328
329        /// Split the websocket into sending and receiving halves.
330        #[must_use]
331        pub fn split(self) -> (WebSocketSender, WebSocketReceiver) {
332            (self.sender, self.receiver)
333        }
334    }
335
336    impl WebSocketSender {
337        /// Send a websocket message serialized as JSON.
338        ///
339        /// # Errors
340        ///
341        /// Returns an error if serialization fails or when the underlying socket cannot
342        /// write the resulting frame.
343        pub async fn send<T>(&self, value: T) -> Result<(), WebSocketError>
344        where
345            T: serde::Serialize,
346        {
347            let payload = serialize_payload(&value)?;
348            self.send_text(payload).await
349        }
350
351        /// Send a text websocket message.
352        ///
353        /// # Errors
354        ///
355        /// Returns an error when the underlying socket cannot write the frame.
356        pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<(), WebSocketError> {
357            self.send_message(WebSocketMessage::text(text)).await
358        }
359
360        /// Send a binary websocket message.
361        ///
362        /// # Errors
363        ///
364        /// Returns an error when the underlying socket cannot write the frame.
365        pub async fn send_binary(&self, bytes: impl Into<Bytes>) -> Result<(), WebSocketError> {
366            self.send_message(WebSocketMessage::binary(bytes)).await
367        }
368
369        async fn send_message(&self, message: WebSocketMessage) -> Result<(), WebSocketError> {
370            let mut sender = self.inner.sender.lock().await;
371            sender
372                .send(to_tungstenite_message(message))
373                .await
374                .map_err(|e| WebSocketError::ConnectionFailed(Box::new(e)))
375        }
376
377        /// Close the websocket connection gracefully.
378        ///
379        /// # Errors
380        ///
381        /// Returns an error when the close frame cannot be sent.
382        pub async fn close(&self) -> Result<(), WebSocketError> {
383            let mut sender = self.inner.sender.lock().await;
384            sender
385                .close(None)
386                .await
387                .map_err(|e| WebSocketError::ConnectionFailed(Box::new(e)))
388        }
389    }
390
391    impl WebSocketReceiver {
392        /// Receive the next websocket message.
393        ///
394        /// # Errors
395        ///
396        /// Returns an error when the underlying socket cannot read the next frame.
397        pub async fn recv(&self) -> Result<Option<WebSocketMessage>, WebSocketError> {
398            loop {
399                let message = {
400                    let mut receiver = self.inner.receiver.lock().await;
401                    receiver.next().await
402                };
403
404                let Some(message) = message else {
405                    return Ok(None);
406                };
407
408                let message = message.map_err(|e| WebSocketError::ConnectionFailed(Box::new(e)))?;
409
410                match message {
411                    TungsteniteMessage::Text(text) => {
412                        return Ok(Some(WebSocketMessage::Text(unsafe {
413                            ByteStr::from_utf8_unchecked(text.into())
414                        })));
415                    }
416                    TungsteniteMessage::Binary(bytes) => {
417                        return Ok(Some(WebSocketMessage::Binary(bytes)));
418                    }
419                    TungsteniteMessage::Close(_) => return Ok(None),
420                    TungsteniteMessage::Ping(payload) => {
421                        self.respond_pong(payload).await?;
422                    }
423                    TungsteniteMessage::Pong(_) | TungsteniteMessage::Frame(_) => {}
424                }
425            }
426        }
427
428        async fn respond_pong(&self, payload: Bytes) -> Result<(), WebSocketError> {
429            let mut sender = self.inner.sender.lock().await;
430            sender
431                .send(TungsteniteMessage::Pong(payload))
432                .await
433                .map_err(|e| WebSocketError::ConnectionFailed(Box::new(e)))
434        }
435    }
436
437    fn to_tungstenite_message(value: WebSocketMessage) -> TungsteniteMessage {
438        match value {
439            WebSocketMessage::Text(text) => TungsteniteMessage::Text(unsafe {
440                Utf8Bytes::from_bytes_unchecked(text.into_bytes())
441            }),
442            WebSocketMessage::Binary(bytes) => TungsteniteMessage::Binary(bytes),
443            WebSocketMessage::Ping(bytes) => TungsteniteMessage::Ping(bytes),
444            WebSocketMessage::Pong(bytes) => TungsteniteMessage::Pong(bytes),
445            WebSocketMessage::Close => TungsteniteMessage::Close(None),
446        }
447    }
448}
449
450#[cfg(target_arch = "wasm32")]
451mod wasm {
452    use std::{cell::RefCell, fmt, rc::Rc, sync::Arc};
453
454    use async_lock::Mutex;
455    use futures_channel::{mpsc, oneshot};
456    use futures_util::StreamExt;
457    use http_kit::utils::{ByteStr, Bytes};
458    use std::io;
459    use wasm_bindgen::{JsCast, JsValue, closure::Closure};
460    use web_sys::{
461        BinaryType, CloseEvent, ErrorEvent, MessageEvent, WebSocket as BrowserWebSocket,
462    };
463
464    use super::{WebSocketConfig, WebSocketError, WebSocketMessage, serialize_payload};
465    use crate::Transport;
466
467    type Result<T> = core::result::Result<T, WebSocketError>;
468
469    enum WsEvent {
470        Message(WebSocketMessage),
471        Error(String),
472        Closed,
473    }
474
475    /// Browser/wasm websocket connection backed by `web_sys`.
476    pub struct WebSocket {
477        sender: WebSocketSender,
478        receiver: WebSocketReceiver,
479    }
480
481    impl fmt::Debug for WebSocket {
482        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
483            f.debug_struct("WebSocket").finish()
484        }
485    }
486
487    #[derive(Debug)]
488    struct SharedSocket {
489        socket: BrowserWebSocket,
490        receiver: Mutex<mpsc::UnboundedReceiver<WsEvent>>,
491        _on_message: Closure<dyn FnMut(MessageEvent)>,
492        _on_error: Closure<dyn FnMut(ErrorEvent)>,
493        _on_close: Closure<dyn FnMut(CloseEvent)>,
494    }
495
496    /// Sending half of a websocket connection.
497    pub struct WebSocketSender {
498        inner: Arc<SharedSocket>,
499    }
500
501    impl fmt::Debug for WebSocketSender {
502        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
503            f.debug_struct("WebSocketSender").finish()
504        }
505    }
506
507    impl Clone for WebSocketSender {
508        fn clone(&self) -> Self {
509            Self {
510                inner: Arc::clone(&self.inner),
511            }
512        }
513    }
514
515    /// Receiving half of a websocket connection.
516    pub struct WebSocketReceiver {
517        inner: Arc<SharedSocket>,
518    }
519
520    impl fmt::Debug for WebSocketReceiver {
521        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
522            f.debug_struct("WebSocketReceiver").finish()
523        }
524    }
525
526    /// Establish a websocket connection from the browser environment.
527    ///
528    /// # Errors
529    ///
530    /// Returns an error if the browser reports an error or the connection fails.
531    pub async fn connect(uri: impl AsRef<str>) -> Result<WebSocket> {
532        connect_with(uri, &Transport::system(), WebSocketConfig::default()).await
533    }
534
535    /// Establish a websocket connection from the browser environment.
536    ///
537    /// The browser owns proxying and trust, so `transport` carries nothing
538    /// here; the parameter exists so cross-platform code has one signature.
539    ///
540    /// # Errors
541    ///
542    /// Returns an error if the browser reports an error or the connection fails.
543    pub async fn connect_with(
544        uri: impl AsRef<str>,
545        _transport: &Transport,
546        _config: WebSocketConfig,
547    ) -> Result<WebSocket> {
548        let socket = BrowserWebSocket::new(uri.as_ref())
549            .map_err(|e| connection_failed(format_js_value(&e)))?;
550        socket.set_binary_type(BinaryType::Arraybuffer);
551
552        let (event_tx, event_rx) = mpsc::unbounded::<WsEvent>();
553        let (ready_tx, ready_rx) = oneshot::channel::<core::result::Result<(), String>>();
554        let pending = Rc::new(RefCell::new(Some(ready_tx)));
555
556        let onopen_pending = Rc::clone(&pending);
557        let on_open = Closure::wrap(Box::new(move || {
558            if let Some(sender) = onopen_pending.borrow_mut().take() {
559                let _ = sender.send(Ok(()));
560            }
561        }) as Box<dyn FnMut()>);
562        socket.set_onopen(Some(on_open.as_ref().unchecked_ref()));
563
564        let on_message_tx = event_tx.clone();
565        let on_message = Closure::wrap(Box::new(move |event: MessageEvent| {
566            let data = event.data();
567            if let Some(text) = data.as_string() {
568                let _ =
569                    on_message_tx.unbounded_send(WsEvent::Message(WebSocketMessage::from(text)));
570                return;
571            }
572
573            if let Ok(array) = data.clone().dyn_into::<js_sys::ArrayBuffer>() {
574                let view = js_sys::Uint8Array::new(&array);
575                let mut bytes = vec![0; view.length() as usize];
576                view.copy_to(&mut bytes[..]);
577                let _ =
578                    on_message_tx.unbounded_send(WsEvent::Message(WebSocketMessage::from(bytes)));
579                return;
580            }
581
582            if let Ok(view) = data.dyn_into::<js_sys::Uint8Array>() {
583                let mut bytes = vec![0; view.length() as usize];
584                view.copy_to(&mut bytes[..]);
585                let _ =
586                    on_message_tx.unbounded_send(WsEvent::Message(WebSocketMessage::from(bytes)));
587                return;
588            }
589
590            let _ = on_message_tx.unbounded_send(WsEvent::Error(
591                "Unsupported websocket message type".to_string(),
592            ));
593        }) as Box<dyn FnMut(MessageEvent)>);
594        socket.set_onmessage(Some(on_message.as_ref().unchecked_ref()));
595
596        let on_error_pending = Rc::clone(&pending);
597        let on_error_tx = event_tx.clone();
598        let on_error = Closure::wrap(Box::new(move |event: ErrorEvent| {
599            let message = event.message();
600            if let Some(sender) = on_error_pending.borrow_mut().take() {
601                let _ = sender.send(Err(message.clone()));
602            }
603            let _ = on_error_tx.unbounded_send(WsEvent::Error(message));
604        }) as Box<dyn FnMut(ErrorEvent)>);
605        socket.set_onerror(Some(on_error.as_ref().unchecked_ref()));
606
607        let on_close_pending = Rc::clone(&pending);
608        let on_close_tx = event_tx.clone();
609        let on_close = Closure::wrap(Box::new(move |event: CloseEvent| {
610            if let Some(sender) = on_close_pending.borrow_mut().take() {
611                let reason = event.reason();
612                let message = if reason.is_empty() {
613                    format!("Connection closed (code {})", event.code())
614                } else {
615                    reason
616                };
617                let _ = sender.send(Err(message));
618            }
619            let _ = on_close_tx.unbounded_send(WsEvent::Closed);
620        }) as Box<dyn FnMut(CloseEvent)>);
621        socket.set_onclose(Some(on_close.as_ref().unchecked_ref()));
622
623        match ready_rx.await {
624            Ok(Ok(())) => {
625                socket.set_onopen(None);
626                drop(on_open);
627            }
628            Ok(Err(message)) => {
629                return Err(connection_failed(message));
630            }
631            Err(_) => {
632                return Err(connection_failed("WebSocket connection cancelled"));
633            }
634        }
635
636        let shared = Arc::new(SharedSocket {
637            socket,
638            receiver: Mutex::new(event_rx),
639            _on_message: on_message,
640            _on_error: on_error,
641            _on_close: on_close,
642        });
643
644        Ok(WebSocket {
645            sender: WebSocketSender {
646                inner: Arc::clone(&shared),
647            },
648            receiver: WebSocketReceiver { inner: shared },
649        })
650    }
651
652    impl WebSocket {
653        /// Send a websocket message serialized as JSON.
654        ///
655        /// # Errors
656        ///
657        /// Returns an error if serialization fails or the browser cannot queue the frame.
658        pub async fn send<T>(&self, value: T) -> Result<()>
659        where
660            T: serde::Serialize,
661        {
662            let payload = serialize_payload(&value)?;
663            self.send_text(payload).await
664        }
665
666        /// Send a text websocket message.
667        ///
668        /// # Errors
669        ///
670        /// Returns an error if the browser fails to queue the frame.
671        pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<()> {
672            self.sender.send_text(text).await
673        }
674
675        /// Send a binary websocket message.
676        ///
677        /// # Errors
678        ///
679        /// Returns an error if the browser fails to queue the frame.
680        pub async fn send_binary(&self, bytes: impl Into<Bytes>) -> Result<()> {
681            self.sender.send_binary(bytes).await
682        }
683
684        /// Receive the next websocket message.
685        ///
686        /// # Errors
687        ///
688        /// Returns an error if the websocket reports an error event.
689        pub async fn recv(&self) -> Result<Option<WebSocketMessage>> {
690            self.receiver.recv().await
691        }
692
693        /// Close the websocket connection gracefully.
694        ///
695        /// # Errors
696        ///
697        /// Returns an error if the browser refuses to close the socket.
698        pub async fn close(self) -> Result<()> {
699            self.sender.close().await
700        }
701
702        /// Split the websocket into sending and receiving halves.
703        #[must_use]
704        pub fn split(self) -> (WebSocketSender, WebSocketReceiver) {
705            (self.sender, self.receiver)
706        }
707    }
708
709    impl WebSocketSender {
710        /// Send a websocket message serialized as JSON.
711        ///
712        /// # Errors
713        ///
714        /// Returns an error if serialization fails or the browser cannot queue the frame.
715        pub async fn send<T>(&self, value: T) -> Result<()>
716        where
717            T: serde::Serialize,
718        {
719            let payload = serialize_payload(&value)?;
720            self.send_text(payload).await
721        }
722
723        /// Send a text websocket message.
724        ///
725        /// # Errors
726        ///
727        /// Returns an error if the browser fails to queue the frame.
728        pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<()> {
729            self.send_message(WebSocketMessage::text(text)).await
730        }
731
732        /// Send a binary websocket message.
733        ///
734        /// # Errors
735        ///
736        /// Returns an error if the browser fails to queue the frame.
737        pub async fn send_binary(&self, bytes: impl Into<Bytes>) -> Result<()> {
738            self.send_message(WebSocketMessage::binary(bytes)).await
739        }
740
741        async fn send_message(&self, message: WebSocketMessage) -> Result<()> {
742            match message {
743                WebSocketMessage::Text(text) => self
744                    .inner
745                    .socket
746                    .send_with_str(&text)
747                    .map_err(|e| connection_failed(format_js_value(&e)))?,
748                WebSocketMessage::Binary(bytes) => self
749                    .inner
750                    .socket
751                    .send_with_u8_array(&bytes)
752                    .map_err(|e| connection_failed(format_js_value(&e)))?,
753                WebSocketMessage::Ping(_) | WebSocketMessage::Pong(_) => {
754                    // Browser WebSocket API doesn't expose ping/pong frames
755                    // They are handled automatically by the browser
756                }
757                WebSocketMessage::Close => {
758                    self.inner
759                        .socket
760                        .close()
761                        .map_err(|e| connection_failed(format_js_value(&e)))?;
762                }
763            }
764            Ok(())
765        }
766
767        /// Close the websocket connection gracefully.
768        ///
769        /// # Errors
770        ///
771        /// Returns an error if the browser refuses to close the socket.
772        pub async fn close(&self) -> Result<()> {
773            self.inner
774                .socket
775                .close()
776                .map_err(|e| connection_failed(format_js_value(&e)))
777        }
778    }
779
780    impl WebSocketReceiver {
781        /// Receive the next websocket message.
782        ///
783        /// # Errors
784        ///
785        /// Returns an error if the websocket reports an error event.
786        pub async fn recv(&self) -> Result<Option<WebSocketMessage>> {
787            let mut receiver = self.inner.receiver.lock().await;
788            match receiver.next().await {
789                Some(WsEvent::Message(message)) => Ok(Some(message)),
790                Some(WsEvent::Closed) | None => Ok(None),
791                Some(WsEvent::Error(message)) => Err(connection_failed(message)),
792            }
793        }
794    }
795
796    fn connection_failed(message: impl Into<ByteStr>) -> WebSocketError {
797        WebSocketError::ConnectionFailed(Box::new(io::Error::new(
798            io::ErrorKind::Other,
799            message.into().to_string(),
800        )))
801    }
802
803    fn format_js_value(value: &JsValue) -> String {
804        value.as_string().unwrap_or_else(|| format!("{value:?}"))
805    }
806}
807
808#[cfg(not(target_arch = "wasm32"))]
809pub use native::{WebSocket, WebSocketReceiver, WebSocketSender, connect, connect_with};
810
811#[cfg(target_arch = "wasm32")]
812pub use wasm::{WebSocket, WebSocketReceiver, WebSocketSender, connect, connect_with};