aquatic_ws: add settings for max ws message sizes, set low defaults

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,
         }
     }
 }

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<TlsAcceptor>, // If set, run TLS
     ) -> (Option<Either<EstablishedWs, Self>>, 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<EstablishedWs, HandshakeMachine>,
 }
@@ -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
                 });

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);
             },