fix: Remove PostComputeError wrapper (#9)

Natchica · web-flow · commit 7b695a39be3f · 2025-05-19T15:31:51.000+02:00
diff --git a/src/api/worker_api.rs b/src/api/worker_api.rs
@@ -29,7 +29,6 @@ use serde::Serialize;
 /// ```
 #[derive(Serialize, Debug)]
 pub struct ExitMessage<'a> {
-    #[serde(rename = "cause")]
     pub cause: &'a ReplicateStatusCause,
 }
 
@@ -171,15 +170,15 @@ mod tests {
         let causes = [
             (
                 ReplicateStatusCause::PostComputeInvalidTeeSignature,
-                "PostComputeInvalidTeeSignature",
+                "POST_COMPUTE_INVALID_TEE_SIGNATURE",
             ),
             (
                 ReplicateStatusCause::PostComputeWorkerAddressMissing,
-                "PostComputeWorkerAddressMissing",
+                "POST_COMPUTE_WORKER_ADDRESS_MISSING",
             ),
             (
                 ReplicateStatusCause::PostComputeFailedUnknownIssue,
-                "PostComputeFailedUnknownIssue",
+                "POST_COMPUTE_FAILED_UNKNOWN_ISSUE",
             ),
         ];
 
diff --git a/src/compute/app_runner.rs b/src/compute/app_runner.rs
@@ -1,6 +1,6 @@
 use crate::api::worker_api::{ExitMessage, WorkerApiClient};
 use crate::compute::{
-    errors::{PostComputeError, ReplicateStatusCause},
+    errors::ReplicateStatusCause,
     signer::get_challenge,
     utils::env_utils::{TeeSessionEnvironmentVariable, get_env_var_or_error},
 };
@@ -14,7 +14,7 @@ use std::error::Error;
 /// the post-compute workflow.
 pub trait PostComputeRunnerInterface {
     fn run_post_compute(&self, chain_task_id: &str) -> Result<(), Box<dyn Error>>;
-    fn get_challenge(&self, chain_task_id: &str) -> Result<String, PostComputeError>;
+    fn get_challenge(&self, chain_task_id: &str) -> Result<String, ReplicateStatusCause>;
     fn send_exit_cause(
         &self,
         authorization: &str,
@@ -41,11 +41,11 @@ impl DefaultPostComputeRunner {
 }
 
 impl PostComputeRunnerInterface for DefaultPostComputeRunner {
-    fn run_post_compute(&self, chain_task_id: &str) -> Result<(), Box<dyn Error>> {
+    fn run_post_compute(&self, _chain_task_id: &str) -> Result<(), Box<dyn Error>> {
         Err("run_post_compute not implemented yet".into())
     }
 
-    fn get_challenge(&self, chain_task_id: &str) -> Result<String, PostComputeError> {
+    fn get_challenge(&self, chain_task_id: &str) -> Result<String, ReplicateStatusCause> {
         get_challenge(chain_task_id)
     }
 
@@ -95,7 +95,7 @@ pub fn start_with_runner<R: PostComputeRunnerInterface>(runner: &R) -> i32 {
     println!("Tee worker post-compute started");
     let chain_task_id: String = match get_env_var_or_error(
         TeeSessionEnvironmentVariable::IEXEC_TASK_ID,
-        ReplicateStatusCause::PostComputeChainTaskIdMissing,
+        ReplicateStatusCause::PostComputeTaskIdMissing,
     ) {
         Ok(id) => id,
         Err(e) => {
@@ -114,9 +114,9 @@ pub fn start_with_runner<R: PostComputeRunnerInterface>(runner: &R) -> i32 {
         }
         Err(error) => {
             let exit_cause: &ReplicateStatusCause;
-            match error.downcast_ref::<PostComputeError>() {
+            match error.downcast_ref::<ReplicateStatusCause>() {
                 Some(post_compute_error) => {
-                    exit_cause = post_compute_error.exit_cause();
+                    exit_cause = post_compute_error;
                     error!(
                         "TEE post-compute failed with exit cause [errorMessage:{}]",
                         &exit_cause
@@ -224,19 +224,17 @@ mod tests {
             if self.run_post_compute_success {
                 Ok(())
             } else if let Some(cause) = &self.error_cause {
-                Err(Box::new(PostComputeError::new(cause.clone())))
+                Err(Box::new(cause.clone()))
             } else {
                 Err("Mock error".into())
             }
         }
 
-        fn get_challenge(&self, _chain_task_id: &str) -> Result<String, PostComputeError> {
+        fn get_challenge(&self, _chain_task_id: &str) -> Result<String, ReplicateStatusCause> {
             if self.get_challenge_success {
                 Ok("mock_challenge".to_string())
             } else {
-                Err(PostComputeError::new(
-                    ReplicateStatusCause::PostComputeTeeChallengePrivateKeyMissing,
-                ))
+                Err(ReplicateStatusCause::PostComputeTeeChallengePrivateKeyMissing)
             }
         }
 
diff --git a/src/compute/errors.rs b/src/compute/errors.rs
@@ -2,33 +2,19 @@ use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 #[derive(Debug, PartialEq, Clone, Error, Serialize, Deserialize)]
+#[serde(rename_all(serialize = "SCREAMING_SNAKE_CASE"))]
+#[allow(clippy::enum_variant_names)]
 pub enum ReplicateStatusCause {
-    #[error("Failed to verify TeeEnclaveChallenge signature (exiting)")]
-    PostComputeInvalidTeeSignature,
+    #[error("Task ID related environment variable is missing")]
+    PostComputeTaskIdMissing,
+    #[error("Unexpected error occurred")]
+    PostComputeFailedUnknownIssue,
     #[error("Invalid enclave challenge private key")]
     PostComputeInvalidEnclaveChallengePrivateKey,
-    #[error("Worker address related environment variable is missing")]
-    PostComputeWorkerAddressMissing,
+    #[error("Invalid TEE signature")]
+    PostComputeInvalidTeeSignature,
     #[error("Tee challenge private key related environment variable is missing")]
     PostComputeTeeChallengePrivateKeyMissing,
-    #[error("Chain task ID related environment variable is missing")]
-    PostComputeChainTaskIdMissing,
-    #[error("Unexpected error occured")]
-    PostComputeFailedUnknownIssue,
-}
-
-#[derive(Debug, Error, Clone)]
-#[error("PostCompute failed: {exit_cause}")]
-pub struct PostComputeError {
-    pub exit_cause: ReplicateStatusCause,
-}
-
-impl PostComputeError {
-    pub fn new(cause: ReplicateStatusCause) -> Self {
-        Self { exit_cause: cause }
-    }
-
-    pub fn exit_cause(&self) -> &ReplicateStatusCause {
-        &self.exit_cause
-    }
+    #[error("Worker address related environment variable is missing")]
+    PostComputeWorkerAddressMissing,
 }
diff --git a/src/compute/signer.rs b/src/compute/signer.rs
@@ -1,5 +1,5 @@
 use crate::compute::{
-    errors::{PostComputeError, ReplicateStatusCause::*},
+    errors::ReplicateStatusCause,
     utils::{
         env_utils::{TeeSessionEnvironmentVariable, get_env_var_or_error},
         hash_utils::{concatenate_and_hash, hex_string_to_byte_array},
@@ -22,7 +22,7 @@ use alloy_signer_local::PrivateKeySigner;
 /// # Returns
 ///
 /// * `Ok(String)` - The signature as a hexadecimal string if successful
-/// * `Err(PostComputeError)` - An error if the private key is invalid or if signing fails
+/// * `Err(ReplicateStatusCause)` - An error if the private key is invalid or if signing fails
 ///
 /// # Errors
 ///
@@ -44,14 +44,14 @@ use alloy_signer_local::PrivateKeySigner;
 pub fn sign_enclave_challenge(
     message_hash: &str,
     enclave_challenge_private_key: &str,
-) -> Result<String, PostComputeError> {
+) -> Result<String, ReplicateStatusCause> {
     let signer: PrivateKeySigner = enclave_challenge_private_key
         .parse::<PrivateKeySigner>()
-        .map_err(|_| PostComputeError::new(PostComputeInvalidEnclaveChallengePrivateKey))?;
+        .map_err(|_| ReplicateStatusCause::PostComputeInvalidEnclaveChallengePrivateKey)?;
 
     let signature: Signature = signer
         .sign_message_sync(&hex_string_to_byte_array(message_hash))
-        .map_err(|_| PostComputeError::new(PostComputeInvalidTeeSignature))?;
+        .map_err(|_| ReplicateStatusCause::PostComputeInvalidTeeSignature)?;
 
     Ok(signature.to_string())
 }
@@ -69,7 +69,7 @@ pub fn sign_enclave_challenge(
 /// # Returns
 ///
 /// * `Ok(String)` - The challenge signature as a hexadecimal string if successful
-/// * `Err(PostComputeError)` - An error if required environment variables are missing or if signing fails
+/// * `Err(ReplicateStatusCause)` - An error if required environment variables are missing or if signing fails
 ///
 /// # Errors
 ///
@@ -97,14 +97,14 @@ pub fn sign_enclave_challenge(
 ///     Err(e) => eprintln!("Error generating challenge: {:?}", e),
 /// }
 /// ```
-pub fn get_challenge(chain_task_id: &str) -> Result<String, PostComputeError> {
+pub fn get_challenge(chain_task_id: &str) -> Result<String, ReplicateStatusCause> {
     let worker_address: String = get_env_var_or_error(
         TeeSessionEnvironmentVariable::SIGN_WORKER_ADDRESS,
-        PostComputeWorkerAddressMissing,
+        ReplicateStatusCause::PostComputeWorkerAddressMissing,
     )?;
     let tee_challenge_private_key: String = get_env_var_or_error(
         TeeSessionEnvironmentVariable::SIGN_TEE_CHALLENGE_PRIVATE_KEY,
-        PostComputeTeeChallengePrivateKeyMissing,
+        ReplicateStatusCause::PostComputeTeeChallengePrivateKeyMissing,
     )?;
     let message_hash: String = concatenate_and_hash(&[chain_task_id, &worker_address]);
     sign_enclave_challenge(&message_hash, &tee_challenge_private_key)
@@ -113,7 +113,6 @@ pub fn get_challenge(chain_task_id: &str) -> Result<String, PostComputeError> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::compute::utils::env_utils::TeeSessionEnvironmentVariable::*;
     use temp_env::with_vars;
 
     const CHAIN_TASK_ID: &str = "0x123456789abcdef";
@@ -141,7 +140,7 @@ mod tests {
         assert!(
             matches!(
                 result,
-                Err(ref err) if err.exit_cause == PostComputeInvalidEnclaveChallengePrivateKey
+                Err(err) if err == ReplicateStatusCause::PostComputeInvalidEnclaveChallengePrivateKey
             ),
             "Should return missing TEE challenge private key error"
         );
@@ -151,9 +150,12 @@ mod tests {
     fn should_get_challenge() {
         with_vars(
             vec![
-                (SIGN_WORKER_ADDRESS.name(), Some(WORKER_ADDRESS)),
                 (
-                    SIGN_TEE_CHALLENGE_PRIVATE_KEY.name(),
+                    TeeSessionEnvironmentVariable::SIGN_WORKER_ADDRESS.name(),
+                    Some(WORKER_ADDRESS),
+                ),
+                (
+                    TeeSessionEnvironmentVariable::SIGN_TEE_CHALLENGE_PRIVATE_KEY.name(),
                     Some(ENCLAVE_CHALLENGE_PRIVATE_KEY),
                 ),
             ],
@@ -181,9 +183,12 @@ mod tests {
     fn should_fail_on_missing_worker_address_env_var() {
         with_vars(
             vec![
-                (SIGN_WORKER_ADDRESS.name(), None),
                 (
-                    SIGN_TEE_CHALLENGE_PRIVATE_KEY.name(),
+                    TeeSessionEnvironmentVariable::SIGN_WORKER_ADDRESS.name(),
+                    None,
+                ),
+                (
+                    TeeSessionEnvironmentVariable::SIGN_TEE_CHALLENGE_PRIVATE_KEY.name(),
                     Some(ENCLAVE_CHALLENGE_PRIVATE_KEY),
                 ),
             ],
@@ -192,7 +197,7 @@ mod tests {
                 assert!(
                     matches!(
                         result,
-                        Err(ref err) if err.exit_cause == PostComputeWorkerAddressMissing
+                        Err(err) if err == ReplicateStatusCause::PostComputeWorkerAddressMissing
                     ),
                     "Should return missing worker address error"
                 );
@@ -204,15 +209,21 @@ mod tests {
     fn should_fail_on_missing_private_key_env_var() {
         with_vars(
             vec![
-                (SIGN_WORKER_ADDRESS.name(), Some(WORKER_ADDRESS)),
-                (SIGN_TEE_CHALLENGE_PRIVATE_KEY.name(), None),
+                (
+                    TeeSessionEnvironmentVariable::SIGN_WORKER_ADDRESS.name(),
+                    Some(WORKER_ADDRESS),
+                ),
+                (
+                    TeeSessionEnvironmentVariable::SIGN_TEE_CHALLENGE_PRIVATE_KEY.name(),
+                    None,
+                ),
             ],
             || {
                 let result = get_challenge(CHAIN_TASK_ID);
                 assert!(
                     matches!(
                         result,
-                        Err(ref err) if err.exit_cause == PostComputeTeeChallengePrivateKeyMissing
+                        Err(err) if err == ReplicateStatusCause::PostComputeTeeChallengePrivateKeyMissing
                     ),
                     "Should return missing private key error"
                 );
diff --git a/src/compute/utils/env_utils.rs b/src/compute/utils/env_utils.rs
@@ -1,4 +1,4 @@
-use crate::compute::errors::{PostComputeError, ReplicateStatusCause};
+use crate::compute::errors::ReplicateStatusCause;
 use std::env;
 
 pub enum TeeSessionEnvironmentVariable {
@@ -24,9 +24,9 @@ impl TeeSessionEnvironmentVariable {
 pub fn get_env_var_or_error(
     env_var: TeeSessionEnvironmentVariable,
     status_cause_if_missing: ReplicateStatusCause,
-) -> Result<String, PostComputeError> {
+) -> Result<String, ReplicateStatusCause> {
     match env::var(env_var.name()) {
         Ok(value) if !value.is_empty() => Ok(value),
-        _ => Err(PostComputeError::new(status_cause_if_missing)),
+        _ => Err(status_cause_if_missing),
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,6 @@ use serde::Serialize;`
`29`	`29`	/// ```
`30`	`30`	`#[derive(Serialize, Debug)]`
`31`	`31`	`pub struct ExitMessage<'a> {`
`32`		`- #[serde(rename = "cause")]`
`33`	`32`	`pub cause: &'a ReplicateStatusCause,`
`34`	`33`	`}`
`35`	`34`
`@@ -171,15 +170,15 @@ mod tests {`
`171`	`170`	`let causes = [`
`172`	`171`	`(`
`173`	`172`	`ReplicateStatusCause::PostComputeInvalidTeeSignature,`
`174`		`- "PostComputeInvalidTeeSignature",`
	`173`	`+ "POST_COMPUTE_INVALID_TEE_SIGNATURE",`
`175`	`174`	`),`
`176`	`175`	`(`
`177`	`176`	`ReplicateStatusCause::PostComputeWorkerAddressMissing,`
`178`		`- "PostComputeWorkerAddressMissing",`
	`177`	`+ "POST_COMPUTE_WORKER_ADDRESS_MISSING",`
`179`	`178`	`),`
`180`	`179`	`(`
`181`	`180`	`ReplicateStatusCause::PostComputeFailedUnknownIssue,`
`182`		`- "PostComputeFailedUnknownIssue",`
	`181`	`+ "POST_COMPUTE_FAILED_UNKNOWN_ISSUE",`
`183`	`182`	`),`
`184`	`183`	`];`
`185`	`184`