From 3b9709fea5e98c15da2adeae5922507261a6ece6 Mon Sep 17 00:00:00 2001
From: Gursharan Singh <3442979+G8XSU@users.noreply.github.com>
Date: Tue, 30 Apr 2024 13:07:28 -0700
Subject: [PATCH] Add AuthException

It will be used when an unauthorized request is made or
 authentication from headers fails.
---
 build.rs       |  2 +-
 src/error.rs   |  7 +++++++
 src/types.rs   |  4 ++++
 tests/tests.rs | 55 +++++++++++++++++++++++++++++++++++++++++++++++++-
 4 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/build.rs b/build.rs
index bd1af3e..12db0a7 100644
--- a/build.rs
+++ b/build.rs
@@ -14,7 +14,7 @@ fn main() {
 #[cfg(genproto)]
 fn generate_protos() {
 	download_file(
-				"https://raw.githubusercontent.com/lightningdevkit/vss-server/cb1159c3b1835c66a857b25b114f15d18d2a4297/app/src/main/proto/vss.proto",
+				"https://raw.githubusercontent.com/lightningdevkit/vss-server/7f492fcac0c561b212f49ca40f7d16075822440f/app/src/main/proto/vss.proto",
 				"src/proto/vss.proto",
 		).unwrap();
 
diff --git a/src/error.rs b/src/error.rs
index 07aa32d..63065c0 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -19,6 +19,9 @@ pub enum VssError {
 	/// Please refer to [`ErrorCode::ConflictException`].
 	ConflictError(String),
 
+	/// Please refer to [`ErrorCode::AuthException`].
+	AuthError(String),
+
 	/// Please refer to [`ErrorCode::InternalServerException`].
 	InternalServerError(String),
 
@@ -53,6 +56,9 @@ impl Display for VssError {
 			VssError::ConflictError(message) => {
 				write!(f, "Potential version conflict in write operation: {}", message)
 			}
+			VssError::AuthError(message) => {
+				write!(f, "Authentication or Authorization failure: {}", message)
+			}
 			VssError::InternalServerError(message) => {
 				write!(f, "InternalServerError: {}", message)
 			}
@@ -71,6 +77,7 @@ impl From<ErrorResponse> for VssError {
 			ErrorCode::NoSuchKeyException => VssError::NoSuchKeyError(error_response.message),
 			ErrorCode::InvalidRequestException => VssError::InvalidRequestError(error_response.message),
 			ErrorCode::ConflictException => VssError::ConflictError(error_response.message),
+			ErrorCode::AuthException => VssError::AuthError(error_response.message),
 			ErrorCode::InternalServerException => VssError::InternalServerError(error_response.message),
 			_ => VssError::InternalError(format!(
 				"VSS responded with an unknown error code: {}, message: {}",
diff --git a/src/types.rs b/src/types.rs
index 6ea4e66..ace954c 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -348,6 +348,8 @@ pub enum ErrorCode {
 	InternalServerException = 3,
 	/// Used when the specified `key` in a `GetObjectRequest` does not exist.
 	NoSuchKeyException = 4,
+	/// Used when authentication fails or in case of an unauthorized request.
+	AuthException = 5,
 }
 impl ErrorCode {
 	/// String value of the enum field names used in the ProtoBuf definition.
@@ -361,6 +363,7 @@ impl ErrorCode {
 			ErrorCode::InvalidRequestException => "INVALID_REQUEST_EXCEPTION",
 			ErrorCode::InternalServerException => "INTERNAL_SERVER_EXCEPTION",
 			ErrorCode::NoSuchKeyException => "NO_SUCH_KEY_EXCEPTION",
+			ErrorCode::AuthException => "AUTH_EXCEPTION",
 		}
 	}
 	/// Creates an enum from field names used in the ProtoBuf definition.
@@ -371,6 +374,7 @@ impl ErrorCode {
 			"INVALID_REQUEST_EXCEPTION" => Some(Self::InvalidRequestException),
 			"INTERNAL_SERVER_EXCEPTION" => Some(Self::InternalServerException),
 			"NO_SUCH_KEY_EXCEPTION" => Some(Self::NoSuchKeyException),
+			"AUTH_EXCEPTION" => Some(Self::AuthException),
 			_ => None,
 		}
 	}
diff --git a/tests/tests.rs b/tests/tests.rs
index 96f5a47..f86dea0 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -255,6 +255,56 @@ mod tests {
 		mock_server.expect(4).assert();
 	}
 
+	#[tokio::test]
+	async fn test_auth_err_handling() {
+		let base_url = mockito::server_url();
+		let vss_client = VssClient::new(&base_url, retry_policy());
+
+		// Invalid Request Error
+		let error_response =
+			ErrorResponse { error_code: ErrorCode::AuthException.into(), message: "AuthException".to_string() };
+		let mock_server = mockito::mock("POST", Matcher::Any)
+			.with_status(401)
+			.with_body(&error_response.encode_to_vec())
+			.create();
+
+		let get_result = vss_client
+			.get_object(&GetObjectRequest { store_id: "store".to_string(), key: "k1".to_string() })
+			.await;
+		assert!(matches!(get_result.unwrap_err(), VssError::AuthError { .. }));
+
+		let put_result = vss_client
+			.put_object(&PutObjectRequest {
+				store_id: "store".to_string(),
+				global_version: Some(4),
+				transaction_items: vec![KeyValue { key: "k1".to_string(), version: 2, value: b"k1v3".to_vec() }],
+				delete_items: vec![],
+			})
+			.await;
+		assert!(matches!(put_result.unwrap_err(), VssError::AuthError { .. }));
+
+		let delete_result = vss_client
+			.delete_object(&DeleteObjectRequest {
+				store_id: "store".to_string(),
+				key_value: Some(KeyValue { key: "k1".to_string(), version: 2, value: b"k1v3".to_vec() }),
+			})
+			.await;
+		assert!(matches!(delete_result.unwrap_err(), VssError::AuthError { .. }));
+
+		let list_result = vss_client
+			.list_key_versions(&ListKeyVersionsRequest {
+				store_id: "store".to_string(),
+				page_size: Some(5),
+				page_token: None,
+				key_prefix: Some("k".into()),
+			})
+			.await;
+		assert!(matches!(list_result.unwrap_err(), VssError::AuthError { .. }));
+
+		// Verify 4 requests hit the server
+		mock_server.expect(4).assert();
+	}
+
 	#[tokio::test]
 	async fn test_conflict_err_handling() {
 		let base_url = mockito::server_url();
@@ -401,7 +451,10 @@ mod tests {
 			.skip_retry_on_error(|e| {
 				matches!(
 					e,
-					VssError::NoSuchKeyError(..) | VssError::InvalidRequestError(..) | VssError::ConflictError(..)
+					VssError::NoSuchKeyError(..)
+						| VssError::InvalidRequestError(..)
+						| VssError::ConflictError(..)
+						| VssError::AuthError(..)
 				)
 			})
 	}