websockets patch feature: add support for protocol config

trillium-rs · Mar 7, 2022 · 5b67d90 · 5b67d90
1 parent a5662ad
commit 5b67d90
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 13 deletions.
diff --git a/websockets/src/json.rs b/websockets/src/json.rs
@@ -217,10 +217,7 @@ where
     [`JsonWebSocketHandler`]
      */
     pub fn new_json(handler: T) -> Self {
-        Self {
-            handler: JsonHandler::new(handler),
-            protocols: Default::default(),
-        }
+        Self::new(JsonHandler::new(handler))
     }
 }
 

diff --git a/websockets/src/lib.rs b/websockets/src/lib.rs
@@ -72,7 +72,7 @@ use trillium::{
 
 pub use async_tungstenite::{
     self,
-    tungstenite::{self, Error, Message},
+    tungstenite::{self, protocol::WebSocketConfig, Error, Message},
 };
 pub use websocket_connection::WebSocketConn;
 pub use websocket_handler::WebSocketHandler;
@@ -96,6 +96,7 @@ See crate-level docs for example usage.
 pub struct WebSocket<H> {
     handler: H,
     protocols: Vec<String>,
+    config: Option<WebSocketConfig>,
 }
 
 impl<H> Deref for WebSocket<H> {
@@ -133,6 +134,7 @@ where
         Self {
             handler,
             protocols: Default::default(),
+            config: None,
         }
     }
 
@@ -145,6 +147,14 @@ where
             ..self
         }
     }
+
+    /// configure the websocket protocol
+    pub fn with_protocol_config(self, config: WebSocketConfig) -> Self {
+        Self {
+            config: Some(config),
+            ..self
+        }
+    }
 }
 
 struct IsWebsocket;
@@ -202,7 +212,7 @@ where
     }
 
     async fn upgrade(&self, upgrade: Upgrade) {
-        let conn = WebSocketConn::new(upgrade).await;
+        let conn = WebSocketConn::new(upgrade, self.config).await;
         let (mut conn, outbound) = unwrap_or_return!(self.handler.connect(conn).await);
         let inbound = conn.take_inbound_stream();
 

diff --git a/websockets/src/websocket_connection.rs b/websockets/src/websocket_connection.rs
@@ -1,6 +1,6 @@
-use crate::Result;
+use crate::{Result, WebSocketConfig};
 use async_tungstenite::{
-    tungstenite::{protocol::Role, Message},
+    tungstenite::{self, protocol::Role, Message},
     WebSocketStream,
 };
 use futures_util::{
@@ -59,11 +59,11 @@ impl WebSocketConn {
     }
 
     /// Sends a [`Message`] to the client
-    pub async fn send(&mut self, message: Message) -> async_tungstenite::tungstenite::Result<()> {
+    pub async fn send(&mut self, message: Message) -> tungstenite::Result<()> {
         self.sink.send(message).await
     }
 
-    pub(crate) async fn new(upgrade: Upgrade) -> Self {
+    pub(crate) async fn new(upgrade: Upgrade, config: Option<WebSocketConfig>) -> Self {
         let Upgrade {
             request_headers,
             path,
@@ -75,9 +75,9 @@ impl WebSocketConn {
         } = upgrade;
 
         let wss = if let Some(vec) = buffer {
-            WebSocketStream::from_partially_read(transport, vec, Role::Server, None).await
+            WebSocketStream::from_partially_read(transport, vec, Role::Server, config).await
         } else {
-            WebSocketStream::from_raw_socket(transport, Role::Server, None).await
+            WebSocketStream::from_raw_socket(transport, Role::Server, config).await
         };
 
         let (sink, stream) = wss.split();
@@ -102,7 +102,7 @@ impl WebSocketConn {
     }
 
     /// close the websocket connection gracefully
-    pub async fn close(&mut self) -> async_tungstenite::tungstenite::Result<()> {
+    pub async fn close(&mut self) -> tungstenite::Result<()> {
         self.send(Message::Close(None)).await
     }