From 78055feb19dc95f53f6145aba606a253efe9f82a Mon Sep 17 00:00:00 2001 From: Quake Wang Date: Tue, 22 Mar 2022 18:21:28 +0900 Subject: [PATCH] fix: max_request_body_size setting should not override max_response_size (#711) * fix: max_request_body_size setting should not override max_response_size * chore: apply review comment --- http-server/src/server.rs | 17 +++++++++++++++- http-server/src/tests.rs | 43 +++++++++++++++++++++++++++++++++++++++ ws-server/src/server.rs | 15 ++++++++++++-- ws-server/src/tests.rs | 24 +++++++++++++++++++++- 4 files changed, 95 insertions(+), 4 deletions(-) diff --git a/http-server/src/server.rs b/http-server/src/server.rs index 71c8d0f6e1..b2a2ba643c 100644 --- a/http-server/src/server.rs +++ b/http-server/src/server.rs @@ -56,6 +56,7 @@ pub struct Builder { 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, @@ -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 Builder { pub fn set_middleware(self, middleware: T) -> Builder { 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 Builder { 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 Builder { 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 { local_addr: Option, /// 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 Server { /// Start the server. pub fn start(mut self, methods: impl Into) -> Result { 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 Server { 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, 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::(); - 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>>; diff --git a/http-server/src/tests.rs b/http-server/src/tests.rs index 40d043f335..8b13334e0f 100644 --- a/http-server/src/tests.rs +++ b/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(); +} diff --git a/ws-server/src/server.rs b/ws-server/src/server.rs index 29da616eb3..600e402138 100644 --- a/ws-server/src/server.rs +++ b/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, @@ -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::>(&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 Builder { 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; diff --git a/ws-server/src/tests.rs b/ws-server/src/tests.rs index 2a1f2edfc9..1664cdc84b 100644 --- a/ws-server/src/tests.rs +++ b/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();