feat(servers): add max_response_size

http-server/src/server.rs

@@ -56,6 +56,7 @@ pub struct Builder<M = ()> {
 	access_control: AccessControl,
 	resources: Resources,
 	max_request_body_size: u32,
+	max_response_body_size: u32,
 	keep_alive: bool,
 	/// Custom tokio runtime to run the server on.
 	tokio_runtime: Option<tokio::runtime::Handle>,
@@ -66,6 +67,7 @@ impl Default for Builder {
 	fn default() -> Self {
 		Self {
 			max_request_body_size: TEN_MB_SIZE_BYTES,
+			max_response_body_size: TEN_MB_SIZE_BYTES,
 			resources: Resources::default(),
 			access_control: AccessControl::default(),
 			keep_alive: true,
@@ -111,6 +113,7 @@ impl<M> Builder<M> {
 	pub fn set_middleware<T: Middleware>(self, middleware: T) -> Builder<T> {
 		Builder {
 			max_request_body_size: self.max_request_body_size,
+			max_response_body_size: self.max_response_body_size,
 			resources: self.resources,
 			access_control: self.access_control,
 			keep_alive: self.keep_alive,
@@ -125,6 +128,12 @@ impl<M> Builder<M> {
 		self
 	}
 
+	/// Sets the maximum size of a response body in bytes (default is 10 MiB).
+	pub fn max_response_body_size(mut self, size: u32) -> Self {
+		self.max_response_body_size = size;
+		self
+	}
+
 	/// Sets access control settings.
 	pub fn set_access_control(mut self, acl: AccessControl) -> Self {
 		self.access_control = acl;
@@ -190,6 +199,7 @@ impl<M> Builder<M> {
 				local_addr,
 				access_control: self.access_control,
 				max_request_body_size: self.max_request_body_size,
+				max_response_body_size: self.max_response_body_size,
 				resources: self.resources,
 				tokio_runtime: self.tokio_runtime,
 				middleware: self.middleware,
@@ -261,6 +271,8 @@ pub struct Server<M = ()> {
 	local_addr: Option<SocketAddr>,
 	/// Max request body size.
 	max_request_body_size: u32,
+	/// Max response body size.
+	max_response_body_size: u32,
 	/// Access control
 	access_control: AccessControl,
 	/// Tracker for currently used resources on the server
@@ -279,6 +291,7 @@ impl<M: Middleware> Server<M> {
 	/// Start the server.
 	pub fn start(mut self, methods: impl Into<Methods>) -> Result<ServerHandle, Error> {
 		let max_request_body_size = self.max_request_body_size;
+		let max_response_body_size = self.max_response_body_size;
 		let access_control = self.access_control;
 		let (tx, mut rx) = mpsc::channel(1);
 		let listener = self.listener;
@@ -341,6 +354,7 @@ impl<M: Middleware> Server<M> {
 									methods,
 									resources,
 									max_request_body_size,
+									max_response_body_size,
 								)
 								.await?;
 
@@ -429,6 +443,7 @@ async fn process_validated_request(
 	methods: Methods,
 	resources: Resources,
 	max_request_body_size: u32,
+	max_response_body_size: u32,
 ) -> Result<hyper::Response<hyper::Body>, HyperError> {
 	let (parts, body) = request.into_parts();
 
@@ -446,7 +461,7 @@ async fn process_validated_request(
 
 	// NOTE(niklasad1): it's a channel because it's needed for batch requests.
 	let (tx, mut rx) = mpsc::unbounded::<String>();
-	let sink = MethodSink::new_with_limit(tx, max_request_body_size);
+	let sink = MethodSink::new_with_limit(tx, max_response_body_size);
 
 	type Notif<'a> = Notification<'a, Option<&'a RawValue>>;
 

http-server/src/tests.rs

@@ -405,3 +405,46 @@ async fn run_forever() {
 	server_handle.handle.take();
 	server_handle.with_timeout(TIMEOUT).await.unwrap();
 }
+
+#[tokio::test]
+async fn can_set_the_max_request_body_size() {
+	let addr = "127.0.0.1:0";
+	// Rejects all requests larger than 100 bytes
+	let server = HttpServerBuilder::default().max_request_body_size(100).build(addr).unwrap();
+	let mut module = RpcModule::new(());
+	module.register_method("anything", |_p, _cx| Ok("a".repeat(100))).unwrap();
+	let addr = server.local_addr().unwrap();
+	let uri = to_http_uri(addr);
+	let handle = server.start(module).unwrap();
+
+	// Invalid: too long
+	let req = format!(r#"{{"jsonrpc":"2.0", "method":{}, "id":1}}"#, "a".repeat(100));
+	let response = http_request(req.into(), uri.clone()).with_default_timeout().await.unwrap().unwrap();
+	assert_eq!(response.body, oversized_request());
+
+	// Max request body size should not override the max response size
+	let req = r#"{"jsonrpc":"2.0", "method":"anything", "id":1}"#;
+	let response = http_request(req.into(), uri.clone()).with_default_timeout().await.unwrap().unwrap();
+	assert_eq!(response.body, ok_response(JsonValue::String("a".repeat(100)), Id::Num(1)));
+
+	handle.stop().unwrap();
+}
+
+#[tokio::test]
+async fn can_set_the_max_response_size() {
+	let addr = "127.0.0.1:0";
+	// Set the max response size to 100 bytes
+	let server = HttpServerBuilder::default().max_response_body_size(100).build(addr).unwrap();
+	let mut module = RpcModule::new(());
+	module.register_method("anything", |_p, _cx| Ok("a".repeat(101))).unwrap();
+	let addr = server.local_addr().unwrap();
+	let uri = to_http_uri(addr);
+	let handle = server.start(module).unwrap();
+
+	// Oversized response.
+	let req = r#"{"jsonrpc":"2.0", "method":"anything", "id":1}"#;
+	let response = http_request(req.into(), uri.clone()).with_default_timeout().await.unwrap().unwrap();
+	assert_eq!(response.body, oversized_response(Id::Num(1), 100));
+
+	handle.stop().unwrap();
+}

ws-server/src/server.rs

@@ -270,6 +270,7 @@ where
 				methods.clone(),
 				resources.clone(),
 				cfg.max_request_body_size,
+				cfg.max_response_body_size,
 				stop_monitor.clone(),
 				middleware,
 				id_provider,
@@ -290,6 +291,7 @@ async fn background_task(
 	methods: Methods,
 	resources: Resources,
 	max_request_body_size: u32,
+	max_response_body_size: u32,
 	stop_server: StopMonitor,
 	middleware: impl Middleware,
 	id_provider: Arc<dyn IdProvider>,
@@ -303,7 +305,7 @@ async fn background_task(
 	let close_notify_server_stop = close_notify.clone();
 
 	let stop_server2 = stop_server.clone();
-	let sink = MethodSink::new_with_limit(tx, max_request_body_size);
+	let sink = MethodSink::new_with_limit(tx, max_response_body_size);
 
 	middleware.on_connect();
 
@@ -476,7 +478,7 @@ async fn background_task(
 					// request in the batch and read the results off of a new channel, `rx_batch`, and then send the
 					// complete batch response back to the client over `tx`.
 					let (tx_batch, mut rx_batch) = mpsc::unbounded();
-					let sink_batch = MethodSink::new_with_limit(tx_batch, max_request_body_size);
+					let sink_batch = MethodSink::new_with_limit(tx_batch, max_response_body_size);
 					if let Ok(batch) = serde_json::from_slice::<Vec<Request>>(&d) {
 						tracing::debug!("recv batch len={}", batch.len());
 						tracing::trace!("recv: batch={:?}", batch);
@@ -624,6 +626,8 @@ impl AllowedValue {
 struct Settings {
 	/// Maximum size in bytes of a request.
 	max_request_body_size: u32,
+	/// Maximum size in bytes of a response.
+	max_response_body_size: u32,
 	/// Maximum number of incoming connections allowed.
 	max_connections: u64,
 	/// Policy by which to accept or deny incoming requests based on the `Origin` header.
@@ -638,6 +642,7 @@ impl Default for Settings {
 	fn default() -> Self {
 		Self {
 			max_request_body_size: TEN_MB_SIZE_BYTES,
+			max_response_body_size: TEN_MB_SIZE_BYTES,
 			max_connections: MAX_CONNECTIONS,
 			allowed_origins: AllowedValue::Any,
 			allowed_hosts: AllowedValue::Any,
@@ -680,6 +685,12 @@ impl<M> Builder<M> {
 		self
 	}
 
+	/// Set the maximum size of a response body in bytes. Default is 10 MiB.
+	pub fn max_response_body_size(mut self, size: u32) -> Self {
+		self.settings.max_response_body_size = size;
+		self
+	}
+
 	/// Set the maximum number of connections allowed. Default is 100.
 	pub fn max_connections(mut self, max: u64) -> Self {
 		self.settings.max_connections = max;

ws-server/src/tests.rs

@@ -183,7 +183,7 @@ async fn can_set_the_max_request_body_size() {
 	init_logger();
 
 	let addr = "127.0.0.1:0";
-	// Rejects all requests larger than 10 bytes
+	// Rejects all requests larger than 100 bytes
 	let server = WsServerBuilder::default().max_request_body_size(100).build(addr).await.unwrap();
 	let mut module = RpcModule::new(());
 	module.register_method("anything", |_p, _cx| Ok("a".repeat(100))).unwrap();
@@ -197,6 +197,28 @@ async fn can_set_the_max_request_body_size() {
 	let response = client.send_request_text(req).await.unwrap();
 	assert_eq!(response, oversized_request());
 
+	// Max request body size should not override the max response body size
+	let req = r#"{"jsonrpc":"2.0", "method":"anything", "id":1}"#;
+	let response = client.send_request_text(req).await.unwrap();
+	assert_eq!(response, ok_response(JsonValue::String("a".repeat(100)), Id::Num(1)));
+
+	handle.stop().unwrap();
+}
+
+#[tokio::test]
+async fn can_set_the_max_response_body_size() {
+	init_logger();
+
+	let addr = "127.0.0.1:0";
+	// Set the max response body size to 100 bytes
+	let server = WsServerBuilder::default().max_response_body_size(100).build(addr).await.unwrap();
+	let mut module = RpcModule::new(());
+	module.register_method("anything", |_p, _cx| Ok("a".repeat(101))).unwrap();
+	let addr = server.local_addr().unwrap();
+	let handle = server.start(module).unwrap();
+
+	let mut client = WebSocketTestClient::new(addr).await.unwrap();
+
 	// Oversized response.
 	let req = r#"{"jsonrpc":"2.0", "method":"anything", "id":1}"#;
 	let response = client.send_request_text(req).await.unwrap();