From c1cced3662b698022bd02627c2499a04dd9f3e45 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Drwi=C4=99ga?= <tomusdrw@users.noreply.github.com>
Date: Mon, 2 Apr 2018 12:33:09 +0200
Subject: [PATCH] Allow customization of max WS connections. (#8257)

* Allow customization of max WS connections.

* remove superflous line.

* Add test for CLI argument parsing.
---
 parity/cli/mod.rs       |  7 +++++++
 parity/configuration.rs | 15 ++++++++++++++-
 parity/rpc.rs           |  3 +++
 rpc/src/lib.rs          |  2 ++
 rpc/src/tests/ws.rs     |  1 +
 5 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/parity/cli/mod.rs b/parity/cli/mod.rs
index 691bdd614..04a4a21e1 100644
--- a/parity/cli/mod.rs
+++ b/parity/cli/mod.rs
@@ -499,6 +499,10 @@ usage! {
 			"--ws-hosts=[HOSTS]",
 			"List of allowed Host header values. This option will validate the Host header sent by the browser, it is additional security against some attack vectors. Special options: \"all\", \"none\".",
 
+			ARG arg_ws_max_connections: (usize) = 100usize, or |c: &Config| c.websockets.as_ref()?.max_connections,
+			"--ws-max-connections=[CONN]",
+			"Maximal number of allowed concurrent WS connections.",
+
 		["API and console options – IPC"]
 			FLAG flag_no_ipc: (bool) = false, or |c: &Config| c.ipc.as_ref()?.disable.clone(),
 			"--no-ipc",
@@ -1087,6 +1091,7 @@ struct Ws {
 	apis: Option<Vec<String>>,
 	origins: Option<Vec<String>>,
 	hosts: Option<Vec<String>>,
+	max_connections: Option<usize>,
 }
 
 #[derive(Default, Debug, PartialEq, Deserialize)]
@@ -1520,6 +1525,7 @@ mod tests {
 			arg_ws_apis: "web3,eth,net,parity,traces,rpc,secretstore".into(),
 			arg_ws_origins: "none".into(),
 			arg_ws_hosts: "none".into(),
+			arg_ws_max_connections: 100,
 
 			// IPC
 			flag_no_ipc: false,
@@ -1759,6 +1765,7 @@ mod tests {
 				apis: None,
 				origins: Some(vec!["none".into()]),
 				hosts: None,
+				max_connections: None,
 			}),
 			rpc: Some(Rpc {
 				disable: Some(true),
diff --git a/parity/configuration.rs b/parity/configuration.rs
index bff839039..2661e90f6 100644
--- a/parity/configuration.rs
+++ b/parity/configuration.rs
@@ -922,6 +922,7 @@ impl Configuration {
 			support_token_api,
 			ui_address: ui.address(),
 			dapps_address: http.address(),
+			max_connections: self.args.arg_ws_max_connections,
 		};
 
 		Ok(conf)
@@ -1361,7 +1362,8 @@ mod tests {
 			signer_path: expected.into(),
 			ui_address: None,
 			dapps_address: Some("127.0.0.1:8545".into()),
-			support_token_api: true
+			support_token_api: true,
+			max_connections: 100,
 		}, UiConfiguration {
 			enabled: false,
 			interface: "127.0.0.1".into(),
@@ -1374,6 +1376,17 @@ mod tests {
         } ));
 	}
 
+	#[test]
+	fn test_ws_max_connections() {
+		let args = vec!["parity", "--ws-max-connections", "1"];
+		let conf = parse(&args);
+
+		assert_eq!(conf.ws_config().unwrap(), WsConfiguration {
+			max_connections: 1,
+			..Default::default()
+		});
+	}
+
 	#[test]
 	fn test_run_cmd() {
 		let args = vec!["parity"];
diff --git a/parity/rpc.rs b/parity/rpc.rs
index 41b10ba87..a1b2271c4 100644
--- a/parity/rpc.rs
+++ b/parity/rpc.rs
@@ -146,6 +146,7 @@ pub struct WsConfiguration {
 	pub interface: String,
 	pub port: u16,
 	pub apis: ApiSet,
+	pub max_connections: usize,
 	pub origins: Option<Vec<String>>,
 	pub hosts: Option<Vec<String>>,
 	pub signer_path: PathBuf,
@@ -162,6 +163,7 @@ impl Default for WsConfiguration {
 			interface: "127.0.0.1".into(),
 			port: 8546,
 			apis: ApiSet::UnsafeContext,
+			max_connections: 100,
 			origins: Some(vec!["parity://*".into(),"chrome-extension://*".into(), "moz-extension://*".into()]),
 			hosts: Some(Vec::new()),
 			signer_path: replace_home(&data_dir, "$BASE/signer").into(),
@@ -240,6 +242,7 @@ pub fn new_ws<D: rpc_apis::Dependencies>(
 		remote.clone(),
 		allowed_origins,
 		allowed_hosts,
+		conf.max_connections,
 		rpc::WsExtractor::new(path.clone()),
 		rpc::WsExtractor::new(path.clone()),
 		rpc::WsStats::new(deps.stats.clone()),
diff --git a/rpc/src/lib.rs b/rpc/src/lib.rs
index 95a0b424b..a08a1a88a 100644
--- a/rpc/src/lib.rs
+++ b/rpc/src/lib.rs
@@ -175,6 +175,7 @@ pub fn start_ws<M, S, H, T, U, V>(
 	remote: tokio_core::reactor::Remote,
 	allowed_origins: ws::DomainsValidation<ws::Origin>,
 	allowed_hosts: ws::DomainsValidation<ws::Host>,
+	max_connections: usize,
 	extractor: T,
 	middleware: V,
 	stats: U,
@@ -191,6 +192,7 @@ pub fn start_ws<M, S, H, T, U, V>(
 		.request_middleware(middleware)
 		.allowed_origins(allowed_origins)
 		.allowed_hosts(allowed_hosts)
+		.max_connections(max_connections)
 		.session_stats(stats)
 		.start(addr)
 }
diff --git a/rpc/src/tests/ws.rs b/rpc/src/tests/ws.rs
index 6d4874845..1db7a91e9 100644
--- a/rpc/src/tests/ws.rs
+++ b/rpc/src/tests/ws.rs
@@ -40,6 +40,7 @@ pub fn serve() -> (Server<ws::Server>, usize, GuardedAuthCodes) {
 		remote,
 		ws::DomainsValidation::Disabled,
 		ws::DomainsValidation::Disabled,
+		5,
 		extractors::WsExtractor::new(Some(&authcodes.path)),
 		extractors::WsExtractor::new(Some(&authcodes.path)),
 		extractors::WsStats::new(stats),