feat: configure request body limit

Author: enigbe

rust/server/src/main.rs

@@ -17,7 +17,7 @@ use tokio::signal::unix::SignalKind;
 use hyper::server::conn::http1;
 use hyper_util::rt::TokioIo;
 
-use crate::vss_service::VssService;
+use crate::vss_service::{VssService, VssServiceConfig};
 use api::auth::{Authorizer, NoopAuthorizer};
 use api::kv_store::KvStore;
 use impls::postgres_store::{Certificate, PostgresPlaintextBackend, PostgresTlsBackend};
@@ -118,7 +118,10 @@ fn main() {
 					match res {
 						Ok((stream, _)) => {
 							let io_stream = TokioIo::new(stream);
-							let vss_service = VssService::new(Arc::clone(&store), Arc::clone(&authorizer));
+							let vss_service_config = if let Some(req_body_size) = config.server_config.maximum_request_body_size {
+								VssServiceConfig::new(req_body_size)
+							} else {VssServiceConfig::default()};
+							let vss_service = VssService::new(Arc::clone(&store), Arc::clone(&authorizer), vss_service_config);
 							runtime.spawn(async move {
 								if let Err(err) = http1::Builder::new().serve_connection(io_stream, vss_service).await {
 									eprintln!("Failed to serve connection: {}", err);

rust/server/src/util/config.rs

@@ -10,6 +10,7 @@ pub(crate) struct Config {
 pub(crate) struct ServerConfig {
 	pub(crate) host: String,
 	pub(crate) port: u16,
+	pub(crate) maximum_request_body_size: Option<usize>,
 }
 
 #[derive(Deserialize)]

rust/server/src/vss_service.rs

@@ -18,17 +18,45 @@ use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
 
-const MAXIMUM_REQUEST_BODY_SIZE: u16 = 65_535;
+const MAXIMUM_REQUEST_BODY_SIZE: usize = 20 * 1024 * 1024;
+const DEFAULT_REQUEST_BODY_SIZE: usize = 10 * 1024 * 1024;
+
+#[derive(Clone)]
+pub(crate) struct VssServiceConfig {
+	maximum_request_body_size: usize,
+}
+
+impl VssServiceConfig {
+	pub fn new(maximum_request_body_size: usize) -> Self {
+		let capped = maximum_request_body_size.min(MAXIMUM_REQUEST_BODY_SIZE);
+		if capped < maximum_request_body_size {
+			eprintln!(
+				"Warning: maximum_request_body_size {} exceeds limit, capped to {}.",
+				maximum_request_body_size, MAXIMUM_REQUEST_BODY_SIZE
+			);
+		}
+		Self { maximum_request_body_size: capped }
+	}
+}
+
+impl Default for VssServiceConfig {
+	fn default() -> Self {
+		Self { maximum_request_body_size: DEFAULT_REQUEST_BODY_SIZE }
+	}
+}
 
 #[derive(Clone)]
 pub struct VssService {
 	store: Arc<dyn KvStore>,
 	authorizer: Arc<dyn Authorizer>,
+	config: VssServiceConfig,
 }
 
 impl VssService {
-	pub(crate) fn new(store: Arc<dyn KvStore>, authorizer: Arc<dyn Authorizer>) -> Self {
-		Self { store, authorizer }
+	pub(crate) fn new(
+		store: Arc<dyn KvStore>, authorizer: Arc<dyn Authorizer>, config: VssServiceConfig,
+	) -> Self {
+		Self { store, authorizer, config }
 	}
 }
 
@@ -43,22 +71,51 @@ impl Service<Request<Incoming>> for VssService {
 		let store = Arc::clone(&self.store);
 		let authorizer = Arc::clone(&self.authorizer);
 		let path = req.uri().path().to_owned();
+		let maximum_request_body_size = self.config.maximum_request_body_size;
 
 		Box::pin(async move {
 			let prefix_stripped_path = path.strip_prefix(BASE_PATH_PREFIX).unwrap_or_default();
 
 			match prefix_stripped_path {
 				"/getObject" => {
-					handle_request(store, authorizer, req, handle_get_object_request).await
+					handle_request(
+						store,
+						authorizer,
+						req,
+						maximum_request_body_size,
+						handle_get_object_request,
+					)
+					.await
 				},
 				"/putObjects" => {
-					handle_request(store, authorizer, req, handle_put_object_request).await
+					handle_request(
+						store,
+						authorizer,
+						req,
+						maximum_request_body_size,
+						handle_put_object_request,
+					)
+					.await
 				},
 				"/deleteObject" => {
-					handle_request(store, authorizer, req, handle_delete_object_request).await
+					handle_request(
+						store,
+						authorizer,
+						req,
+						maximum_request_body_size,
+						handle_delete_object_request,
+					)
+					.await
 				},
 				"/listKeyVersions" => {
-					handle_request(store, authorizer, req, handle_list_object_request).await
+					handle_request(
+						store,
+						authorizer,
+						req,
+						maximum_request_body_size,
+						handle_list_object_request,
+					)
+					.await
 				},
 				_ => {
 					let error_msg = "Invalid request path.".as_bytes();
@@ -99,7 +156,7 @@ async fn handle_request<
 	Fut: Future<Output = Result<R, VssError>> + Send,
 >(
 	store: Arc<dyn KvStore>, authorizer: Arc<dyn Authorizer>, request: Request<Incoming>,
-	handler: F,
+	maximum_request_body_size: usize, handler: F,
 ) -> Result<<VssService as Service<Request<Incoming>>>::Response, hyper::Error> {
 	let (parts, body) = request.into_parts();
 	let headers_map = parts
@@ -113,7 +170,7 @@ async fn handle_request<
 		Err(e) => return Ok(build_error_response(e)),
 	};
 
-	let limited_body = Limited::new(body, MAXIMUM_REQUEST_BODY_SIZE.into());
+	let limited_body = Limited::new(body, maximum_request_body_size);
 	let bytes = match limited_body.collect().await {
 		Ok(body) => body.to_bytes(),
 		Err(_) => {

rust/server/vss-server-config.toml

@@ -1,6 +1,7 @@
 [server_config]
 host = "127.0.0.1"
 port = 8080
+# maximum_request_body_size = 10485760       # Optional:  maximum request body size in bytes capped at 20971520 (20 MB).
 
 [postgresql_config]
 username = "postgres"  # Optional in TOML, can be overridden by env var `VSS_POSTGRESQL_USERNAME`