refactor(websocket): Refactor websocket implementation

0x676e67 · Jan 28, 2025 · fee5458 · fee5458
1 parent 4eefefd
commit fee5458
Show file tree

Hide file tree

Showing 8 changed files with 467 additions and 309 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -31,7 +31,7 @@ full = [
     "deflate",
 ]
 
-websocket = ["dep:async-tungstenite", "tokio-util/compat"]
+websocket = ["dep:tokio-tungstenite"]
 
 charset = ["dep:encoding_rs"]
 
@@ -147,7 +147,7 @@ tokio-util = { version = "0.7.0", default-features = false, features = [
 tokio-socks = { version = "0.5.2", optional = true }
 
 ## websocket
-async-tungstenite = { version = "0.28.0", optional = true }
+tokio-tungstenite = { version = "0.26.0", default-features = false, features = ["handshake"], optional = true }
 
 ## hickory-dns
 hickory-resolver = { version = "0.24", optional = true }

diff --git a/examples/http2_websocket.rs b/examples/http2_websocket.rs
@@ -7,7 +7,7 @@
 
 use futures_util::{SinkExt, StreamExt, TryStreamExt};
 use http::header;
-use rquest::{Client, Impersonate, Message, RequestBuilder};
+use rquest::{Client, Impersonate, Message, RequestBuilder, Utf8Bytes};
 use std::time::Duration;
 
 #[tokio::main]
@@ -26,7 +26,7 @@ async fn main() -> Result<(), rquest::Error> {
     let websocket = client
         .websocket("wss://127.0.0.1:3000/ws")
         .configure_request(configure_request)
-        .http2_only()
+        .use_http2()
         .send()
         .await?;
 
@@ -36,16 +36,20 @@ async fn main() -> Result<(), rquest::Error> {
 
     tokio::spawn(async move {
         for i in 1..11 {
-            tx.send(Message::Text(format!("Hello, World! #{i}")))
+            if let Err(err) = tx
+                .send(Message::Text(Utf8Bytes::from(format!(
+                    "Hello, World! #{i}"
+                ))))
                 .await
-                .unwrap();
+            {
+                eprintln!("failed to send message: {err}");
+            }
         }
     });
 
     while let Some(message) = rx.try_next().await? {
-        match message {
-            Message::Text(text) => println!("received: {text}"),
-            _ => {}
+        if let Message::Text(text) = message {
+            println!("received: {text}");
         }
     }
 

diff --git a/examples/websocket.rs b/examples/websocket.rs
@@ -7,7 +7,7 @@
 
 use futures_util::{SinkExt, StreamExt, TryStreamExt};
 use http::header;
-use rquest::{Client, Impersonate, Message, RequestBuilder};
+use rquest::{Client, Impersonate, Message, RequestBuilder, Utf8Bytes};
 use std::time::Duration;
 
 #[tokio::main]
@@ -35,16 +35,20 @@ async fn main() -> Result<(), rquest::Error> {
 
     tokio::spawn(async move {
         for i in 1..11 {
-            tx.send(Message::Text(format!("Hello, World! #{i}")))
+            if let Err(err) = tx
+                .send(Message::Text(Utf8Bytes::from(format!(
+                    "Hello, World! #{i}"
+                ))))
                 .await
-                .unwrap();
+            {
+                eprintln!("failed to send message: {err}");
+            }
         }
     });
 
     while let Some(message) = rx.try_next().await? {
-        match message {
-            Message::Text(text) => println!("received: {text}"),
-            _ => {}
+        if let Message::Text(text) = message {
+            println!("received: {text}");
         }
     }
 

diff --git a/src/client/websocket/json.rs b/src/client/websocket/json.rs
@@ -1,6 +1,9 @@
 use crate::{Error, Message};
+use bytes::Bytes;
 use serde::{de::DeserializeOwned, Serialize};
 
+use super::Utf8Bytes;
+
 impl Message {
     /// Tries to serialize the JSON as a [`Message::Text`].
     ///
@@ -15,6 +18,7 @@ impl Message {
     #[cfg_attr(docsrs, doc(cfg(feature = "json")))]
     pub fn text_from_json<T: Serialize + ?Sized>(json: &T) -> Result<Self, Error> {
         serde_json::to_string(json)
+            .map(Utf8Bytes::from)
             .map(Message::Text)
             .map_err(Into::into)
     }
@@ -32,6 +36,7 @@ impl Message {
     #[cfg_attr(docsrs, doc(cfg(feature = "json")))]
     pub fn binary_from_json<T: Serialize + ?Sized>(json: &T) -> Result<Self, Error> {
         serde_json::to_vec(json)
+            .map(Bytes::from)
             .map(Message::Binary)
             .map_err(Into::into)
     }