Skip to content

Commit

Permalink
Expose direct configuration of max in and out buffer capacity (#594)
Browse files Browse the repository at this point in the history
* Expose direct configuration of max in and out buffer capacity

* Correct argument duplicate typo"

* Add server builder tests

* rustfmt ws/src/server_builder.rs
  • Loading branch information
emostov authored Oct 5, 2020
1 parent 07a002a commit 2cb9752
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 2 deletions.
6 changes: 4 additions & 2 deletions ws/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@ impl Server {
executor: UninitializedExecutor,
max_connections: usize,
max_payload_bytes: usize,
max_in_buffer_capacity: usize,
max_out_buffer_capacity: usize,
) -> Result<Server>
where
S::Future: Unpin,
Expand All @@ -68,8 +70,8 @@ impl Server {
config.max_connections = max_connections;
// don't accept super large requests
config.max_fragment_size = max_payload_bytes;
config.max_in_buffer_capacity = max_payload_bytes;
config.max_out_buffer_capacity = max_payload_bytes;
config.max_in_buffer_capacity = max_in_buffer_capacity;
config.max_out_buffer_capacity = max_out_buffer_capacity;
// don't grow non-final fragments (to prevent DOS)
config.fragments_grow = false;
config.fragments_capacity = cmp::max(1, max_payload_bytes / config.fragment_size);
Expand Down
63 changes: 63 additions & 0 deletions ws/src/server_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ pub struct ServerBuilder<M: core::Metadata, S: core::Middleware<M>> {
executor: UninitializedExecutor,
max_connections: usize,
max_payload_bytes: usize,
max_in_buffer_capacity: usize,
max_out_buffer_capacity: usize,
}

impl<M: core::Metadata + Default, S: core::Middleware<M>> ServerBuilder<M, S>
Expand Down Expand Up @@ -61,6 +63,8 @@ where
executor: UninitializedExecutor::Unspawned,
max_connections: 100,
max_payload_bytes: 5 * 1024 * 1024,
max_in_buffer_capacity: 10 * 1024 * 1024,
max_out_buffer_capacity: 10 * 1024 * 1024,
}
}

Expand Down Expand Up @@ -115,6 +119,20 @@ where
self
}

/// The maximum size to which the incoming buffer can grow.
/// Default: 10,485,760
pub fn max_in_buffer_capacity(mut self, max_in_buffer_capacity: usize) -> Self {
self.max_in_buffer_capacity = max_in_buffer_capacity;
self
}

/// The maximum size to which the outgoing buffer can grow.
/// Default: 10,485,760
pub fn max_out_buffer_capacity(mut self, max_out_buffer_capacity: usize) -> Self {
self.max_out_buffer_capacity = max_out_buffer_capacity;
self
}

/// Starts a new `WebSocket` server in separate thread.
/// Returns a `Server` handle which closes the server when droped.
pub fn start(self, addr: &SocketAddr) -> Result<Server> {
Expand All @@ -129,6 +147,51 @@ where
self.executor,
self.max_connections,
self.max_payload_bytes,
self.max_in_buffer_capacity,
self.max_out_buffer_capacity,
)
}
}

#[cfg(test)]
mod tests {
use super::*;

fn basic_server_builder() -> ServerBuilder<(), jsonrpc_core::middleware::Noop> {
let io = core::IoHandler::default();
ServerBuilder::new(io)
}
#[test]
fn config_usize_vals_have_correct_defaults() {
let server = basic_server_builder();

assert_eq!(server.max_connections, 100);
assert_eq!(server.max_payload_bytes, 5 * 1024 * 1024);
assert_eq!(server.max_in_buffer_capacity, 10 * 1024 * 1024);
assert_eq!(server.max_out_buffer_capacity, 10 * 1024 * 1024);
}

#[test]
fn config_usize_vals_can_be_set() {
let server = basic_server_builder();

// We can set them individually
let server = server.max_connections(10);
assert_eq!(server.max_connections, 10);

let server = server.max_payload(29);
assert_eq!(server.max_payload_bytes, 29);

let server = server.max_in_buffer_capacity(38);
assert_eq!(server.max_in_buffer_capacity, 38);

let server = server.max_out_buffer_capacity(47);
assert_eq!(server.max_out_buffer_capacity, 47);

// Setting values consecutively does not impact other values
assert_eq!(server.max_connections, 10);
assert_eq!(server.max_payload_bytes, 29);
assert_eq!(server.max_in_buffer_capacity, 38);
assert_eq!(server.max_out_buffer_capacity, 47);
}
}

0 comments on commit 2cb9752

Please sign in to comment.