Add WebSocket::config method to set the WebSocket configuration.

poem/CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+# [3.1.1] 2024-10-02
+
+- Add `WebSocket::config` method to set the WebSocket configuration.
+
 # [3.1.0] 2024-09-08
 
 - build(deps): update nix requirement from 0.28.0 to 0.29.0 [#851](https://github.com/poem-web/poem/pull/851)

poem/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "poem"
-version = "3.1.0"
+version = "3.1.1"
 authors.workspace = true
 edition.workspace = true
 license.workspace = true

poem/src/web/websocket/extractor.rs

@@ -2,7 +2,7 @@ use std::{borrow::Cow, future::Future};
 
 use futures_util::{future::BoxFuture, FutureExt};
 use headers::HeaderMapExt;
-use tokio_tungstenite::tungstenite::protocol::Role;
+use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig};
 
 use super::{utils::sign, WebSocketStream};
 use crate::{
@@ -24,6 +24,7 @@ pub struct WebSocket {
     on_upgrade: OnUpgrade,
     protocols: Option<Box<[Cow<'static, str>]>>,
     sec_websocket_protocol: Option<HeaderValue>,
+    config: Option<WebSocketConfig>,
 }
 
 impl WebSocket {
@@ -62,6 +63,7 @@ impl WebSocket {
             on_upgrade: req.take_upgrade()?,
             protocols: None,
             sec_websocket_protocol,
+            config: None,
         })
     }
 }
@@ -109,6 +111,14 @@ impl WebSocket {
         self
     }
 
+    /// Set the WebSocket configuration.
+    pub fn config(self, config: WebSocketConfig) -> Self {
+        Self {
+            config: Some(config),
+            ..self
+        }
+    }
+
     /// Finalize upgrading the connection and call the provided `callback` with
     /// the stream.
     ///
@@ -198,9 +208,12 @@ where
                 Err(_) => return,
             };
 
-            let stream =
-                tokio_tungstenite::WebSocketStream::from_raw_socket(upgraded, Role::Server, None)
-                    .await;
+            let stream = tokio_tungstenite::WebSocketStream::from_raw_socket(
+                upgraded,
+                Role::Server,
+                self.websocket.config,
+            )
+            .await;
             (self.callback)(WebSocketStream::new(stream)).await;
         });
 

poem/src/web/websocket/mod.rs

@@ -30,6 +30,7 @@ mod utils;
 pub use extractor::{BoxWebSocketUpgraded, WebSocket, WebSocketUpgraded};
 pub use message::{CloseCode, Message};
 pub use stream::WebSocketStream;
+pub use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
 
 #[cfg(test)]
 mod tests {