client: distinstict APIs to configure max request and response sizes (#…

…967) * client: distinstict APIs to configure max size * Update client/http-client/src/transport.rs Co-authored-by: Alexandru Vasile <60601340+lexnv@users.noreply.github.com> * fix tests Co-authored-by: Alexandru Vasile <60601340+lexnv@users.noreply.github.com>
paritytech · Jan 10, 2023 · 86a3eb3 · 86a3eb3
1 parent c7d6873
commit 86a3eb3
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 41 deletions.
diff --git a/client/http-client/src/client.rs b/client/http-client/src/client.rs
@@ -71,7 +71,8 @@ use tracing::instrument;
 /// ```
 #[derive(Debug)]
 pub struct HttpClientBuilder {
-	max_request_body_size: u32,
+	max_request_size: u32,
+	max_response_size: u32,
 	request_timeout: Duration,
 	max_concurrent_requests: usize,
 	certificate_store: CertificateStore,
@@ -81,9 +82,15 @@ pub struct HttpClientBuilder {
 }
 
 impl HttpClientBuilder {
-	/// Sets the maximum size of a request body in bytes (default is 10 MiB).
-	pub fn max_request_body_size(mut self, size: u32) -> Self {
-		self.max_request_body_size = size;
+	/// Set the maximum size of a request body in bytes. Default is 10 MiB.
+	pub fn max_request_size(mut self, size: u32) -> Self {
+		self.max_request_size = size;
+		self
+	}
+
+	/// Set the maximum size of a response in bytes. Default is 10 MiB.
+	pub fn max_response_size(mut self, size: u32) -> Self {
+		self.max_response_size = size;
 		self
 	}
 
@@ -130,7 +137,8 @@ impl HttpClientBuilder {
 	/// Build the HTTP client with target to connect to.
 	pub fn build(self, target: impl AsRef<str>) -> Result<HttpClient, Error> {
 		let Self {
-			max_request_body_size,
+			max_request_size,
+			max_response_size,
 			max_concurrent_requests,
 			request_timeout,
 			certificate_store,
@@ -139,9 +147,15 @@ impl HttpClientBuilder {
 			max_log_length,
 		} = self;
 
-		let transport =
-			HttpTransportClient::new(target, max_request_body_size, certificate_store, max_log_length, headers)
-				.map_err(|e| Error::Transport(e.into()))?;
+		let transport = HttpTransportClient::new(
+			max_request_size,
+			target,
+			max_response_size,
+			certificate_store,
+			max_log_length,
+			headers,
+		)
+		.map_err(|e| Error::Transport(e.into()))?;
 		Ok(HttpClient {
 			transport,
 			id_manager: Arc::new(RequestIdManager::new(max_concurrent_requests, id_kind)),
@@ -153,7 +167,8 @@ impl HttpClientBuilder {
 impl Default for HttpClientBuilder {
 	fn default() -> Self {
 		Self {
-			max_request_body_size: TEN_MB_SIZE_BYTES,
+			max_request_size: TEN_MB_SIZE_BYTES,
+			max_response_size: TEN_MB_SIZE_BYTES,
 			request_timeout: Duration::from_secs(60),
 			max_concurrent_requests: 256,
 			certificate_store: CertificateStore::Native,

diff --git a/client/http-client/src/transport.rs b/client/http-client/src/transport.rs
@@ -44,7 +44,9 @@ pub struct HttpTransportClient {
 	/// HTTP client
 	client: HyperClient,
 	/// Configurable max request body size
-	max_request_body_size: u32,
+	max_request_size: u32,
+	/// Configurable max response body size
+	max_response_size: u32,
 	/// Max length for logging for requests and responses
 	///
 	/// Logs bigger than this limit will be truncated.
@@ -56,8 +58,9 @@ pub struct HttpTransportClient {
 impl HttpTransportClient {
 	/// Initializes a new HTTP client.
 	pub(crate) fn new(
+		max_request_size: u32,
 		target: impl AsRef<str>,
-		max_request_body_size: u32,
+		max_response_size: u32,
 		cert_store: CertificateStore,
 		max_log_length: u32,
 		headers: HeaderMap,
@@ -107,13 +110,13 @@ impl HttpTransportClient {
 			}
 		}
 
-		Ok(Self { target, client, max_request_body_size, max_log_length, headers: cached_headers })
+		Ok(Self { target, client, max_request_size, max_response_size, max_log_length, headers: cached_headers })
 	}
 
 	async fn inner_send(&self, body: String) -> Result<hyper::Response<hyper::Body>, Error> {
 		tx_log_from_str(&body, self.max_log_length);
 
-		if body.len() > self.max_request_body_size as usize {
+		if body.len() > self.max_request_size as usize {
 			return Err(Error::RequestTooLarge);
 		}
 
@@ -135,7 +138,7 @@ impl HttpTransportClient {
 	pub(crate) async fn send_and_read_body(&self, body: String) -> Result<Vec<u8>, Error> {
 		let response = self.inner_send(body).await?;
 		let (parts, body) = response.into_parts();
-		let (body, _) = http_helpers::read_body(&parts.headers, body, self.max_request_body_size).await?;
+		let (body, _) = http_helpers::read_body(&parts.headers, body, self.max_response_size).await?;
 
 		rx_log_from_bytes(&body, self.max_log_length);
 
@@ -210,21 +213,22 @@ mod tests {
 		assert_eq!(client.target.path_and_query().map(|pq| pq.as_str()), Some(path_and_query));
 		assert_eq!(client.target.host(), Some(host));
 		assert_eq!(client.target.port_u16(), Some(port));
-		assert_eq!(client.max_request_body_size, max_request_size);
+		assert_eq!(client.max_request_size, max_request_size);
 	}
 
 	#[test]
 	fn invalid_http_url_rejected() {
-		let err = HttpTransportClient::new("ws://localhost:9933", 80, CertificateStore::Native, 80, HeaderMap::new())
-			.unwrap_err();
+		let err =
+			HttpTransportClient::new(80, "ws://localhost:9933", 80, CertificateStore::Native, 80, HeaderMap::new())
+				.unwrap_err();
 		assert!(matches!(err, Error::Url(_)));
 	}
 
 	#[cfg(feature = "tls")]
 	#[test]
 	fn https_works() {
 		let client =
-			HttpTransportClient::new("https://localhost:9933", 80, CertificateStore::Native, 80, HeaderMap::new())
+			HttpTransportClient::new(80, "https://localhost:9933", 80, CertificateStore::Native, 80, HeaderMap::new())
 				.unwrap();
 		assert_target(&client, "localhost", "https", "/", 9933, 80);
 	}
@@ -233,25 +237,27 @@ mod tests {
 	#[test]
 	fn https_fails_without_tls_feature() {
 		let err =
-			HttpTransportClient::new("https://localhost:9933", 80, CertificateStore::Native, 80, HeaderMap::new())
+			HttpTransportClient::new(80, "https://localhost:9933", 80, CertificateStore::Native, 80, HeaderMap::new())
 				.unwrap_err();
 		assert!(matches!(err, Error::Url(_)));
 	}
 
 	#[test]
 	fn faulty_port() {
-		let err = HttpTransportClient::new("http://localhost:-43", 80, CertificateStore::Native, 80, HeaderMap::new())
-			.unwrap_err();
+		let err =
+			HttpTransportClient::new(80, "http://localhost:-43", 80, CertificateStore::Native, 80, HeaderMap::new())
+				.unwrap_err();
 		assert!(matches!(err, Error::Url(_)));
 		let err =
-			HttpTransportClient::new("http://localhost:-99999", 80, CertificateStore::Native, 80, HeaderMap::new())
+			HttpTransportClient::new(80, "http://localhost:-99999", 80, CertificateStore::Native, 80, HeaderMap::new())
 				.unwrap_err();
 		assert!(matches!(err, Error::Url(_)));
 	}
 
 	#[test]
 	fn url_with_path_works() {
 		let client = HttpTransportClient::new(
+			1337,
 			"http://localhost:9944/my-special-path",
 			1337,
 			CertificateStore::Native,
@@ -265,6 +271,7 @@ mod tests {
 	#[test]
 	fn url_with_query_works() {
 		let client = HttpTransportClient::new(
+			u32::MAX,
 			"http://127.0.0.1:9999/my?name1=value1&name2=value2",
 			u32::MAX,
 			CertificateStore::WebPki,
@@ -278,6 +285,7 @@ mod tests {
 	#[test]
 	fn url_with_fragment_is_ignored() {
 		let client = HttpTransportClient::new(
+			999,
 			"http://127.0.0.1:9944/my.htm#ignore",
 			999,
 			CertificateStore::Native,
@@ -291,10 +299,19 @@ mod tests {
 	#[tokio::test]
 	async fn request_limit_works() {
 		let eighty_bytes_limit = 80;
-		let client =
-			HttpTransportClient::new("http://localhost:9933", 80, CertificateStore::WebPki, 99, HeaderMap::new())
-				.unwrap();
-		assert_eq!(client.max_request_body_size, eighty_bytes_limit);
+		let fifty_bytes_limit = 50;
+
+		let client = HttpTransportClient::new(
+			eighty_bytes_limit,
+			"http://localhost:9933",
+			fifty_bytes_limit,
+			CertificateStore::WebPki,
+			99,
+			HeaderMap::new(),
+		)
+		.unwrap();
+		assert_eq!(client.max_request_size, eighty_bytes_limit);
+		assert_eq!(client.max_response_size, fifty_bytes_limit);
 
 		let body = "a".repeat(81);
 		assert_eq!(body.len(), 81);

diff --git a/client/transport/src/ws/mod.rs b/client/transport/src/ws/mod.rs
@@ -49,6 +49,7 @@ pub use soketto::handshake::client::Header;
 #[derive(Debug)]
 pub struct Sender {
 	inner: connection::Sender<BufReader<BufWriter<EitherStream>>>,
+	max_request_size: u32,
 }
 
 /// Receiving end of WebSocket transport.
@@ -66,8 +67,10 @@ pub struct WsTransportClientBuilder {
 	pub connection_timeout: Duration,
 	/// Custom headers to pass during the HTTP handshake.
 	pub headers: http::HeaderMap,
-	/// Max payload size
-	pub max_request_body_size: u32,
+	/// Max request payload size
+	pub max_request_size: u32,
+	/// Max response payload size
+	pub max_response_size: u32,
 	/// Max number of redirections.
 	pub max_redirections: usize,
 }
@@ -76,7 +79,8 @@ impl Default for WsTransportClientBuilder {
 	fn default() -> Self {
 		Self {
 			certificate_store: CertificateStore::Native,
-			max_request_body_size: TEN_MB_SIZE_BYTES,
+			max_request_size: TEN_MB_SIZE_BYTES,
+			max_response_size: TEN_MB_SIZE_BYTES,
 			connection_timeout: Duration::from_secs(10),
 			headers: http::HeaderMap::new(),
 			max_redirections: 5,
@@ -91,9 +95,15 @@ impl WsTransportClientBuilder {
 		self
 	}
 
-	/// Set max request body size (default is 10 MB).
-	pub fn max_request_body_size(mut self, size: u32) -> Self {
-		self.max_request_body_size = size;
+	/// Set the maximum size of a request in bytes. Default is 10 MiB.
+	pub fn max_request_size(mut self, size: u32) -> Self {
+		self.max_request_size = size;
+		self
+	}
+
+	/// Set the maximum size of a response in bytes. Default is 10 MiB.
+	pub fn max_response_size(mut self, size: u32) -> Self {
+		self.max_response_size = size;
 		self
 	}
 
@@ -181,6 +191,9 @@ pub enum WsError {
 	/// Error in the WebSocket connection.
 	#[error("WebSocket connection error: {0}")]
 	Connection(#[source] soketto::connection::Error),
+	/// Message was too large.
+	#[error("The message was too large")]
+	MessageTooLarge,
 }
 
 #[async_trait]
@@ -190,6 +203,10 @@ impl TransportSenderT for Sender {
 	/// Sends out a request. Returns a `Future` that finishes when the request has been
 	/// successfully sent.
 	async fn send(&mut self, body: String) -> Result<(), Self::Error> {
+		if body.len() > self.max_request_size as usize {
+			return Err(WsError::MessageTooLarge);
+		}
+
 		tracing::trace!("send: {}", body);
 		self.inner.send_text(body).await?;
 		self.inner.flush().await?;
@@ -300,9 +317,12 @@ impl WsTransportClientBuilder {
 					Ok(ServerResponse::Accepted { .. }) => {
 						tracing::debug!("Connection established to target: {:?}", target);
 						let mut builder = client.into_builder();
-						builder.set_max_message_size(self.max_request_body_size as usize);
+						builder.set_max_message_size(self.max_response_size as usize);
 						let (sender, receiver) = builder.finish();
-						return Ok((Sender { inner: sender }, Receiver { inner: receiver }));
+						return Ok((
+							Sender { inner: sender, max_request_size: self.max_request_size },
+							Receiver { inner: receiver },
+						));
 					}
 
 					Ok(ServerResponse::Rejected { status_code }) => {

diff --git a/client/ws-client/src/lib.rs b/client/ws-client/src/lib.rs
@@ -76,7 +76,8 @@ use jsonrpsee_core::{Error, TEN_MB_SIZE_BYTES};
 #[derive(Clone, Debug)]
 pub struct WsClientBuilder {
 	certificate_store: CertificateStore,
-	max_request_body_size: u32,
+	max_request_size: u32,
+	max_response_size: u32,
 	request_timeout: Duration,
 	connection_timeout: Duration,
 	ping_interval: Option<Duration>,
@@ -92,7 +93,8 @@ impl Default for WsClientBuilder {
 	fn default() -> Self {
 		Self {
 			certificate_store: CertificateStore::Native,
-			max_request_body_size: TEN_MB_SIZE_BYTES,
+			max_request_size: TEN_MB_SIZE_BYTES,
+			max_response_size: TEN_MB_SIZE_BYTES,
 			request_timeout: Duration::from_secs(60),
 			connection_timeout: Duration::from_secs(10),
 			ping_interval: None,
@@ -113,9 +115,15 @@ impl WsClientBuilder {
 		self
 	}
 
-	/// See documentation [`WsTransportClientBuilder::max_request_body_size`] (default is 10 MB).
-	pub fn max_request_body_size(mut self, size: u32) -> Self {
-		self.max_request_body_size = size;
+	/// See documentation [`WsTransportClientBuilder::max_request_size`] (default is 10 MB).
+	pub fn max_request_size(mut self, size: u32) -> Self {
+		self.max_request_size = size;
+		self
+	}
+
+	/// See documentation [`WsTransportClientBuilder::max_response_size`] (default is 10 MB).
+	pub fn max_response_size(mut self, size: u32) -> Self {
+		self.max_response_size = size;
 		self
 	}
 
@@ -185,7 +193,8 @@ impl WsClientBuilder {
 		let Self {
 			certificate_store,
 			max_concurrent_requests,
-			max_request_body_size,
+			max_request_size,
+			max_response_size,
 			request_timeout,
 			connection_timeout,
 			ping_interval,
@@ -200,7 +209,8 @@ impl WsClientBuilder {
 			certificate_store,
 			connection_timeout,
 			headers,
-			max_request_body_size,
+			max_request_size,
+			max_response_size,
 			max_redirections,
 		};