f: Add protocol version response header

Author: tankyleo

rust/server/src/vss_service.rs

@@ -23,6 +23,8 @@ use log::{debug, trace};
 use crate::util::KeyValueVecKeyPrinter;
 
 const MAXIMUM_REQUEST_BODY_SIZE: usize = 1024 * 1024 * 1024;
+const PROTOCOL_VERSION_HEADER: &str = "vss-protocol-version";
+const PROTOCOL_VERSION: Version = Version::V1;
 
 #[derive(Clone, Copy)]
 pub(crate) struct VssServiceConfig {
@@ -81,12 +83,8 @@ impl Service<Request<Incoming>> for VssService {
 
 			match prefix_stripped_path {
 				"/version" => {
-					let response = VersionResponse { version: Version::V1.into() };
-					let response = Response::builder()
-						.body(Full::new(Bytes::from(response.encode_to_vec())))
-						// unwrap safety: body only errors when previous chained calls failed.
-						.unwrap();
-					Ok(response)
+					let response = VersionResponse { version: PROTOCOL_VERSION.into() };
+					Ok(build_response(StatusCode::OK, Bytes::from(response.encode_to_vec())))
 				},
 				"/getObject" => {
 					handle_request(
@@ -128,13 +126,10 @@ impl Service<Request<Incoming>> for VssService {
 					)
 					.await
 				},
-				_ => {
-					let error_msg = "Invalid request path.".as_bytes();
-					Ok(Response::builder()
-						.status(StatusCode::BAD_REQUEST)
-						.body(Full::new(Bytes::from(error_msg)))
-						.unwrap())
-				},
+				_ => Ok(build_response(
+					StatusCode::BAD_REQUEST,
+					Bytes::from_static(b"Invalid request path."),
+				)),
 			}
 		})
 	}
@@ -226,29 +221,35 @@ async fn handle_request<
 	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 safety: body only errors when previous chained calls failed.
-				.unwrap());
+			return Ok(build_response(
+				StatusCode::PAYLOAD_TOO_LARGE,
+				Bytes::from_static(b"Request body too large"),
+			));
 		},
 	};
 	match T::decode(bytes) {
 		Ok(request) => match handler(store.clone(), user_token, request).await {
-			Ok(response) => Ok(Response::builder()
-				.body(Full::new(Bytes::from(response.encode_to_vec())))
-				// unwrap safety: body only errors when previous chained calls failed.
-				.unwrap()),
+			Ok(response) => {
+				Ok(build_response(StatusCode::OK, Bytes::from(response.encode_to_vec())))
+			},
 			Err(e) => Ok(build_error_response(e)),
 		},
-		Err(_) => Ok(Response::builder()
-			.status(StatusCode::BAD_REQUEST)
-			.body(Full::new(Bytes::from(b"Error parsing request".to_vec())))
-			// unwrap safety: body only errors when previous chained calls failed.
-			.unwrap()),
+		Err(_) => Ok(build_response(
+			StatusCode::BAD_REQUEST,
+			Bytes::from_static(b"Error parsing request"),
+		)),
 	}
 }
 
+fn build_response(status_code: StatusCode, body: Bytes) -> Response<Full<Bytes>> {
+	Response::builder()
+		.status(status_code)
+		.header(PROTOCOL_VERSION_HEADER, PROTOCOL_VERSION.as_str_name())
+		.body(Full::new(body))
+		// unwrap safety: body only errors when previous chained calls failed.
+		.unwrap()
+}
+
 fn build_error_response(e: VssError) -> Response<Full<Bytes>> {
 	let (status_code, error_response) = match e {
 		VssError::NoSuchKeyError(msg) => {
@@ -292,9 +293,36 @@ fn build_error_response(e: VssError) -> Response<Full<Bytes>> {
 			(status, error)
 		},
 	};
-	Response::builder()
-		.status(status_code)
-		.body(Full::new(Bytes::from(error_response.encode_to_vec())))
-		// unwrap safety: body only errors when previous chained calls failed.
-		.unwrap()
+	build_response(status_code, Bytes::from(error_response.encode_to_vec()))
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn build_response_adds_protocol_version_header() {
+		let response = build_response(StatusCode::OK, Bytes::new());
+
+		assert_eq!(
+			Version::from_str_name(
+				response.headers().get(PROTOCOL_VERSION_HEADER).unwrap().to_str().unwrap()
+			)
+			.unwrap(),
+			PROTOCOL_VERSION,
+		);
+	}
+
+	#[test]
+	fn build_error_response_adds_protocol_version_header() {
+		let response = build_error_response(VssError::InvalidRequestError("bad request".into()));
+
+		assert_eq!(
+			Version::from_str_name(
+				response.headers().get(PROTOCOL_VERSION_HEADER).unwrap().to_str().unwrap()
+			)
+			.unwrap(),
+			PROTOCOL_VERSION,
+		);
+	}
 }