diff --git a/aquatic_ws/src/lib/config.rs b/aquatic_ws/src/lib/config.rs index cf208af..153832e 100644 --- a/aquatic_ws/src/lib/config.rs +++ b/aquatic_ws/src/lib/config.rs @@ -33,6 +33,8 @@ pub struct NetworkConfig { pub peer_announce_interval: usize, // FIXME: should this really be in NetworkConfig? pub poll_event_capacity: usize, pub poll_timeout_milliseconds: u64, + pub websocket_max_message_size: usize, + pub websocket_max_frame_size: usize, } @@ -96,6 +98,8 @@ impl Default for NetworkConfig { peer_announce_interval: 120, poll_event_capacity: 4096, poll_timeout_milliseconds: 50, + websocket_max_message_size: 64 * 1024, + websocket_max_frame_size: 16 * 1024, } } } diff --git a/aquatic_ws/src/lib/network/connection.rs b/aquatic_ws/src/lib/network/connection.rs index c77533e..befcbf9 100644 --- a/aquatic_ws/src/lib/network/connection.rs +++ b/aquatic_ws/src/lib/network/connection.rs @@ -10,6 +10,7 @@ use native_tls::{TlsAcceptor, TlsStream, MidHandshakeTlsStream}; use tungstenite::WebSocket; use tungstenite::handshake::{MidHandshake, HandshakeError, server::NoCallback}; use tungstenite::server::{ServerHandshake}; +use tungstenite::protocol::WebSocketConfig; use crate::common::*; @@ -102,6 +103,7 @@ impl HandshakeMachine { #[inline] fn advance( self, + ws_config: WebSocketConfig, opt_tls_acceptor: &Option, // If set, run TLS ) -> (Option>, bool) { // bool = stop looping match self { @@ -111,8 +113,9 @@ impl HandshakeMachine { tls_acceptor.accept(stream) ) } else { - let handshake_result = ::tungstenite::server::accept( + let handshake_result = ::tungstenite::server::accept_with_config( Stream::TcpStream(stream), + Some(ws_config) ); Self::handle_ws_handshake_result(handshake_result) @@ -188,6 +191,7 @@ pub struct EstablishedWs { pub struct Connection { + ws_config: WebSocketConfig, pub valid_until: ValidUntil, inner: Either, } @@ -198,10 +202,12 @@ pub struct Connection { impl Connection { #[inline] pub fn new( + ws_config: WebSocketConfig, valid_until: ValidUntil, tcp_stream: TcpStream, ) -> Self { Self { + ws_config, valid_until, inner: Either::Right(HandshakeMachine::new(tcp_stream)) } @@ -224,9 +230,15 @@ impl Connection { match self.inner { Either::Left(_) => (Some(self), false), Either::Right(machine) => { - let (opt_inner, stop_loop) = machine.advance(opt_tls_acceptor); + let ws_config = self.ws_config; + + let (opt_inner, stop_loop) = machine.advance( + ws_config, + opt_tls_acceptor + ); let opt_new_self = opt_inner.map(|inner| Self { + ws_config, valid_until, inner }); diff --git a/aquatic_ws/src/lib/network/mod.rs b/aquatic_ws/src/lib/network/mod.rs index e65116a..8e92829 100644 --- a/aquatic_ws/src/lib/network/mod.rs +++ b/aquatic_ws/src/lib/network/mod.rs @@ -6,6 +6,7 @@ use log::info; use native_tls::TlsAcceptor; use mio::{Events, Poll, Interest, Token}; use mio::net::TcpListener; +use tungstenite::protocol::WebSocketConfig; use crate::common::*; use crate::config::Config; @@ -54,6 +55,11 @@ pub fn run_poll_loop( let poll_timeout = Duration::from_millis( config.network.poll_timeout_milliseconds ); + let ws_config = WebSocketConfig { + max_message_size: Some(config.network.websocket_max_message_size), + max_frame_size: Some(config.network.websocket_max_frame_size), + max_send_queue: None, + }; let mut listener = TcpListener::from_std(listener); let mut poll = Poll::new().expect("create poll"); @@ -79,6 +85,7 @@ pub fn run_poll_loop( if token.0 == 0 { accept_new_streams( + ws_config, &mut listener, &mut poll, &mut connections, @@ -113,6 +120,7 @@ pub fn run_poll_loop( fn accept_new_streams( + ws_config: WebSocketConfig, listener: &mut TcpListener, poll: &mut Poll, connections: &mut ConnectionMap, @@ -136,7 +144,7 @@ fn accept_new_streams( .register(&mut stream, token, Interest::READABLE) .unwrap(); - let connection = Connection::new(valid_until, stream); + let connection = Connection::new(ws_config, valid_until, stream); connections.insert(token, connection); },