1pub use http_kit::ws::*;
2
3use http_kit::{HttpError, StatusCode};
4use serde::Serialize;
5
6#[derive(Debug, thiserror::Error)]
8pub enum WebSocketError {
9 #[error("Fail to encode payload: {0}")]
11 FailToEncodePayload(serde_json::Error),
12
13 #[error("Unsupported websocket scheme: {0}")]
15 UnsupportedScheme(String),
16
17 #[error("Invalid URI: {0}")]
19 InvalidUri(#[from] url::ParseError),
20
21 #[error("Invalid websocket target: {0}")]
23 InvalidTarget(&'static str),
24
25 #[error("{0}")]
27 Transport(#[source] Box<crate::Error>),
28
29 #[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
44impl 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#[derive(Clone, Debug)]
68#[non_exhaustive]
69pub struct WebSocketConfig {
70 pub max_message_size: Option<usize>,
73
74 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 #[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 #[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 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 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 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 pub async fn connect(uri: impl AsRef<str>) -> Result<WebSocket, WebSocketError> {
208 connect_with(uri, &Transport::system(), WebSocketConfig::default()).await
209 }
210
211 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 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 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 pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<(), WebSocketError> {
299 self.sender.send_text(text).await
300 }
301
302 pub async fn send_binary(&self, bytes: impl Into<Bytes>) -> Result<(), WebSocketError> {
308 self.sender.send_binary(bytes).await
309 }
310
311 pub async fn recv(&self) -> Result<Option<WebSocketMessage>, WebSocketError> {
317 self.receiver.recv().await
318 }
319
320 pub async fn close(self) -> Result<(), WebSocketError> {
326 self.sender.close().await
327 }
328
329 #[must_use]
331 pub fn split(self) -> (WebSocketSender, WebSocketReceiver) {
332 (self.sender, self.receiver)
333 }
334 }
335
336 impl WebSocketSender {
337 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 pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<(), WebSocketError> {
357 self.send_message(WebSocketMessage::text(text)).await
358 }
359
360 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 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 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 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 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 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 pub async fn connect(uri: impl AsRef<str>) -> Result<WebSocket> {
532 connect_with(uri, &Transport::system(), WebSocketConfig::default()).await
533 }
534
535 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 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 pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<()> {
672 self.sender.send_text(text).await
673 }
674
675 pub async fn send_binary(&self, bytes: impl Into<Bytes>) -> Result<()> {
681 self.sender.send_binary(bytes).await
682 }
683
684 pub async fn recv(&self) -> Result<Option<WebSocketMessage>> {
690 self.receiver.recv().await
691 }
692
693 pub async fn close(self) -> Result<()> {
699 self.sender.close().await
700 }
701
702 #[must_use]
704 pub fn split(self) -> (WebSocketSender, WebSocketReceiver) {
705 (self.sender, self.receiver)
706 }
707 }
708
709 impl WebSocketSender {
710 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 pub async fn send_text(&self, text: impl Into<ByteStr>) -> Result<()> {
729 self.send_message(WebSocketMessage::text(text)).await
730 }
731
732 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 }
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 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 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};