feat: bound incoming request and add postgres service

rust/api/src/kv_store_tests.rs

@@ -44,6 +44,7 @@ macro_rules! define_kv_store_tests {
 		create_test!(put_should_fail_when_global_version_mismatched);
 		create_test!(put_should_succeed_when_no_global_version_is_given);
 		create_test!(put_and_delete_should_succeed_as_atomic_transaction);
+		create_test!(put_should_succeed_with_maximum_supported_value_size);
 		create_test!(delete_should_succeed_when_item_exists);
 		create_test!(delete_should_succeed_when_item_does_not_exist);
 		create_test!(delete_should_be_idempotent);
@@ -266,6 +267,29 @@ pub trait KvStoreTestSuite {
 		Ok(())
 	}
 
+	async fn put_should_succeed_with_maximum_supported_value_size() -> Result<(), VssError> {
+		const MAXIMUM_SUPPORTED_VALUE_SIZE: usize = 1024 * 1024 * 1024;
+		const PROTOCOL_OVERHEAD_MARGIN: usize = 150;
+		let kv_store = Self::create_store().await;
+		let ctx = TestContext::new(&kv_store);
+
+		// Construct entry that's for a field that's the maximum size of a non-"large_object" object
+		let large_value = vec![0u8; MAXIMUM_SUPPORTED_VALUE_SIZE - PROTOCOL_OVERHEAD_MARGIN];
+		let kv = KeyValue { key: "k1".into(), version: 0, value: Bytes::from(large_value) };
+
+		// Put succeeds
+		ctx.put_objects(None, vec![kv]).await?;
+
+		// Retrieval succeeds
+		let result = ctx.get_object("k1").await?;
+		assert_eq!(result.value.len(), MAXIMUM_SUPPORTED_VALUE_SIZE - PROTOCOL_OVERHEAD_MARGIN);
+		assert!(result.value.iter().all(|&b| b == 0));
+
+		ctx.delete_object(result).await?;
+
+		Ok(())
+	}
+
 	async fn delete_should_succeed_when_item_exists() -> Result<(), VssError> {
 		let kv_store = Self::create_store().await;
 		let ctx = TestContext::new(&kv_store);

rust/docker-compose.yml

@@ -0,0 +1,21 @@
+version: '3.8'
+services:
+  postgres:
+    image: postgres:15
+    environment:
+      POSTGRES_DB: postgres
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: postgres
+    volumes:
+      - postgres-data:/var/lib/postgresql/data
+    ports:
+      - "5432:5432"
+    networks:
+      - app-network
+
+volumes:
+  postgres-data:
+
+networks:
+  app-network:
+    driver: bridge

rust/server/src/main.rs

@@ -23,7 +23,7 @@ use api::kv_store::KvStore;
 use auth_impls::JWTAuthorizer;
 use impls::postgres_store::{Certificate, PostgresPlaintextBackend, PostgresTlsBackend};
 use util::config::{Config, ServerConfig};
-use vss_service::VssService;
+use vss_service::{VssService, VssServiceConfig};
 
 mod util;
 mod vss_service;
@@ -35,14 +35,17 @@ fn main() {
 		std::process::exit(1);
 	}
 
-	let Config { server_config: ServerConfig { host, port }, jwt_auth_config, postgresql_config } =
-		match util::config::load_config(&args[1]) {
-			Ok(cfg) => cfg,
-			Err(e) => {
-				eprintln!("Failed to load configuration: {}", e);
-				std::process::exit(1);
-			},
-		};
+	let Config {
+		server_config: ServerConfig { host, port, maximum_request_body_size },
+		jwt_auth_config,
+		postgresql_config,
+	} = match util::config::load_config(&args[1]) {
+		Ok(cfg) => cfg,
+		Err(e) => {
+			eprintln!("Failed to load configuration: {}", e);
+			std::process::exit(1);
+		},
+	};
 	let addr: SocketAddr = match format!("{}:{}", host, port).parse() {
 		Ok(addr) => addr,
 		Err(e) => {
@@ -144,7 +147,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) = 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

@@ -11,6 +11,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

@@ -1,4 +1,4 @@
-use http_body_util::{BodyExt, Full};
+use http_body_util::{BodyExt, Full, Limited};
 use hyper::body::{Bytes, Incoming};
 use hyper::service::Service;
 use hyper::{Request, Response, StatusCode};
@@ -18,15 +18,37 @@ use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
 
+const MAXIMUM_REQUEST_BODY_SIZE: usize = 1024 * 1024 * 1024;
+
+#[derive(Clone)]
+pub(crate) struct VssServiceConfig {
+	maximum_request_body_size: usize,
+}
+
+impl VssServiceConfig {
+	pub fn new(maximum_request_body_size: usize) -> Self {
+		Self { maximum_request_body_size: maximum_request_body_size.min(MAXIMUM_REQUEST_BODY_SIZE) }
+	}
+}
+
+impl Default for VssServiceConfig {
+	fn default() -> Self {
+		Self { maximum_request_body_size: MAXIMUM_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 }
 	}
 }
 
@@ -41,22 +63,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();
@@ -97,7 +148,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
@@ -110,8 +161,17 @@ async fn handle_request<
 		Ok(auth_response) => auth_response.user_token,
 		Err(e) => return Ok(build_error_response(e)),
 	};
-	// TODO: we should bound the amount of data we read to avoid allocating too much memory.
-	let bytes = body.collect().await?.to_bytes();
+
+	let limited_body = Limited::new(body, maximum_request_body_size);
+	let bytes = match limited_body.collect().await {
+		Ok(body) => body.to_bytes(),
+		Err(_) => {
+			return Ok(Response::builder()
+				.status(StatusCode::PAYLOAD_TOO_LARGE)
+				.body(Full::new(Bytes::from("Request body too large")))
+				.unwrap());
+		},
+	};
 	match T::decode(bytes) {
 		Ok(request) => match handler(store.clone(), user_token, request).await {
 			Ok(response) => Ok(Response::builder()

rust/server/vss-server-config.toml

@@ -1,6 +1,7 @@
 [server_config]
 host = "127.0.0.1"
 port = 8080
+# maximum_request_body_size = 1073741824    # Optional:  maximum request body size in bytes capped at 1 GB.
 
 # Uncomment the table below to verify JWT tokens in the HTTP Authorization header against the given RSA public key,
 # can be overridden by env var `VSS_JWT_RSA_PEM`