feat: use dataset address for improved error handling

nabil-Tounarti · nabil-Tounarti · commit 70d4c7b045d8 · 2025-10-20T15:53:07.000+02:00
diff --git a/pre-compute/src/api/worker_api.rs b/pre-compute/src/api/worker_api.rs
@@ -154,12 +154,12 @@ mod tests {
                 r#"{"cause":"PRE_COMPUTE_WORKER_ADDRESS_MISSING","message":"Worker address related environment variable is missing"}"#,
             ),
             (
-                ReplicateStatusCause::PreComputeDatasetUrlMissing(2),
-                r#"{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 2"}"#,
+                ReplicateStatusCause::PreComputeDatasetUrlMissing("0xDataset2".to_string()),
+                r#"{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 0xDataset2"}"#,
             ),
             (
-                ReplicateStatusCause::PreComputeInvalidDatasetChecksum(1),
-                r#"{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 1"}"#,
+                ReplicateStatusCause::PreComputeInvalidDatasetChecksum("0xDataset1".to_string()),
+                r#"{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 0xDataset1"}"#,
             ),
         ];
 
@@ -172,12 +172,12 @@ mod tests {
     #[test]
     fn serialize_vec_of_causes_succeeds_when_multiple_causes() {
         let causes = vec![
-            ReplicateStatusCause::PreComputeDatasetUrlMissing(0),
-            ReplicateStatusCause::PreComputeInvalidDatasetChecksum(1),
+            ReplicateStatusCause::PreComputeDatasetUrlMissing("0xDatasetA".to_string()),
+            ReplicateStatusCause::PreComputeInvalidDatasetChecksum("0xDatasetB".to_string()),
         ];
 
         let serialized = to_string(&causes).expect("Failed to serialize");
-        let expected = r#"[{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 0"},{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 1"}]"#;
+        let expected = r#"[{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 0xDatasetA"},{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 0xDatasetB"}]"#;
         assert_eq!(serialized, expected);
     }
     // endregion
diff --git a/pre-compute/src/compute/dataset.rs b/pre-compute/src/compute/dataset.rs
@@ -27,15 +27,23 @@ const AES_IV_LENGTH: usize = 16;
 #[cfg_attr(test, derive(Debug))]
 #[derive(Clone, Default)]
 pub struct Dataset {
+    pub address: String,
     pub url: String,
     pub checksum: String,
     pub filename: String,
     pub key: String,
 }
 
 impl Dataset {
-    pub fn new(url: String, checksum: String, filename: String, key: String) -> Self {
+    pub fn new(
+        address: String,
+        url: String,
+        checksum: String,
+        filename: String,
+        key: String,
+    ) -> Self {
         Dataset {
+            address,
             url,
             checksum,
             filename,
@@ -48,7 +56,6 @@ impl Dataset {
     /// # Arguments
     ///
     /// * `chain_task_id` - The chain task ID for logging
-    /// * `index` - The dataset index for error reporting
     ///
     /// # Returns
     ///
@@ -58,7 +65,6 @@ impl Dataset {
     pub fn download_encrypted_dataset(
         &self,
         chain_task_id: &str,
-        index: usize,
     ) -> Result<Vec<u8>, ReplicateStatusCause> {
         info!(
             "Downloading encrypted dataset file [chainTaskId:{chain_task_id}, url:{}]",
@@ -81,7 +87,9 @@ impl Dataset {
         } else {
             download_from_url(&self.url)
         }
-        .ok_or(ReplicateStatusCause::PreComputeDatasetDownloadFailed(index))?;
+        .ok_or(ReplicateStatusCause::PreComputeDatasetDownloadFailed(
+            self.address.clone(),
+        ))?;
 
         info!("Checking encrypted dataset checksum [chainTaskId:{chain_task_id}]");
         let actual_checksum = sha256_from_bytes(&encrypted_content);
@@ -92,7 +100,7 @@ impl Dataset {
                 self.checksum
             );
             return Err(ReplicateStatusCause::PreComputeInvalidDatasetChecksum(
-                index,
+                self.address.clone(),
             ));
         }
 
@@ -108,7 +116,6 @@ impl Dataset {
     /// # Arguments
     ///
     /// * `encrypted_content` - Full encrypted dataset, including the IV prefix.
-    /// * `index` - The dataset index for error reporting
     ///
     /// # Returns
     ///
@@ -117,15 +124,14 @@ impl Dataset {
     pub fn decrypt_dataset(
         &self,
         encrypted_content: &[u8],
-        index: usize,
     ) -> Result<Vec<u8>, ReplicateStatusCause> {
-        let key = general_purpose::STANDARD
-            .decode(&self.key)
-            .map_err(|_| ReplicateStatusCause::PreComputeDatasetDecryptionFailed(index))?;
+        let key = general_purpose::STANDARD.decode(&self.key).map_err(|_| {
+            ReplicateStatusCause::PreComputeDatasetDecryptionFailed(self.address.clone())
+        })?;
 
         if encrypted_content.len() < AES_IV_LENGTH || key.len() != AES_KEY_LENGTH {
             return Err(ReplicateStatusCause::PreComputeDatasetDecryptionFailed(
-                index,
+                self.address.clone(),
             ));
         }
 
@@ -135,7 +141,9 @@ impl Dataset {
 
         Aes256CbcDec::new(key_slice.into(), iv_slice.into())
             .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
-            .map_err(|_| ReplicateStatusCause::PreComputeDatasetDecryptionFailed(index))
+            .map_err(|_| {
+                ReplicateStatusCause::PreComputeDatasetDecryptionFailed(self.address.clone())
+            })
     }
 }
 
@@ -157,6 +165,7 @@ mod tests {
 
     fn get_test_dataset() -> Dataset {
         Dataset::new(
+            "0xDatasetAddress".to_string(),
             HTTP_DATASET_URL.to_string(),
             DATASET_CHECKSUM.to_string(),
             PLAIN_DATA_FILE.to_string(),
@@ -168,18 +177,21 @@ mod tests {
     #[test]
     fn download_encrypted_dataset_success() {
         let dataset = get_test_dataset();
-        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID, 0);
+        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID);
         assert!(actual_content.is_ok());
     }
 
     #[test]
     fn download_encrypted_dataset_failure_with_invalid_dataset_url() {
         let mut dataset = get_test_dataset();
         dataset.url = "http://bad-url".to_string();
-        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID, 10);
+        dataset.address = "0xbaddataset".to_string();
+        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID);
         assert_eq!(
             actual_content,
-            Err(ReplicateStatusCause::PreComputeDatasetDownloadFailed(10))
+            Err(ReplicateStatusCause::PreComputeDatasetDownloadFailed(
+                "0xbaddataset".to_string()
+            ))
         );
     }
 
@@ -189,7 +201,7 @@ mod tests {
         dataset.url = IPFS_DATASET_URL.to_string();
         dataset.checksum =
             "0x323b1637c7999942fbebfe5d42fe15dbfe93737577663afa0181938d7ad4a2ac".to_string();
-        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID, 0);
+        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID);
         let expected_content = Ok("hello world !\n".as_bytes().to_vec());
         assert_eq!(actual_content, expected_content);
     }
@@ -198,17 +210,23 @@ mod tests {
     fn download_encrypted_dataset_failure_with_invalid_gateway() {
         let mut dataset = get_test_dataset();
         dataset.url = "/ipfs/INVALID_IPFS_DATASET_URL".to_string();
-        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID, 1);
-        let expected_content = Err(ReplicateStatusCause::PreComputeDatasetDownloadFailed(1));
+        dataset.address = "0xinvalidgateway".to_string();
+        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID);
+        let expected_content = Err(ReplicateStatusCause::PreComputeDatasetDownloadFailed(
+            "0xinvalidgateway".to_string(),
+        ));
         assert_eq!(actual_content, expected_content);
     }
 
     #[test]
     fn download_encrypted_dataset_failure_with_invalid_dataset_checksum() {
         let mut dataset = get_test_dataset();
         dataset.checksum = "invalid_dataset_checksum".to_string();
-        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID, 9999);
-        let expected_content = Err(ReplicateStatusCause::PreComputeInvalidDatasetChecksum(9999));
+        dataset.address = "0xinvalidchecksum".to_string();
+        let actual_content = dataset.download_encrypted_dataset(CHAIN_TASK_ID);
+        let expected_content = Err(ReplicateStatusCause::PreComputeInvalidDatasetChecksum(
+            "0xinvalidchecksum".to_string(),
+        ));
         assert_eq!(actual_content, expected_content);
     }
     // endregion
@@ -218,11 +236,9 @@ mod tests {
     fn decrypt_dataset_success_with_valid_dataset() {
         let dataset = get_test_dataset();
 
-        let encrypted_data = dataset
-            .download_encrypted_dataset(CHAIN_TASK_ID, 3)
-            .unwrap();
+        let encrypted_data = dataset.download_encrypted_dataset(CHAIN_TASK_ID).unwrap();
         let expected_plain_data = Ok("Some very useful data.".as_bytes().to_vec());
-        let actual_plain_data = dataset.decrypt_dataset(&encrypted_data, 3);
+        let actual_plain_data = dataset.decrypt_dataset(&encrypted_data);
 
         assert_eq!(actual_plain_data, expected_plain_data);
     }
@@ -231,14 +247,15 @@ mod tests {
     fn decrypt_dataset_failure_with_bad_key() {
         let mut dataset = get_test_dataset();
         dataset.key = "bad_key".to_string();
-        let encrypted_data = dataset
-            .download_encrypted_dataset(CHAIN_TASK_ID, 55)
-            .unwrap();
-        let actual_plain_data = dataset.decrypt_dataset(&encrypted_data, 55);
+        dataset.address = "0xbadkey".to_string();
+        let encrypted_data = dataset.download_encrypted_dataset(CHAIN_TASK_ID).unwrap();
+        let actual_plain_data = dataset.decrypt_dataset(&encrypted_data);
 
         assert_eq!(
             actual_plain_data,
-            Err(ReplicateStatusCause::PreComputeDatasetDecryptionFailed(55))
+            Err(ReplicateStatusCause::PreComputeDatasetDecryptionFailed(
+                "0xbadkey".to_string()
+            ))
         );
     }
     // endregion
diff --git a/pre-compute/src/compute/errors.rs b/pre-compute/src/compute/errors.rs
@@ -10,17 +10,17 @@ pub enum ReplicateStatusCause {
     #[error("input file URL {0} is missing")]
     PreComputeAtLeastOneInputFileUrlMissing(usize),
     #[error("Dataset checksum related environment variable is missing for dataset {0}")]
-    PreComputeDatasetChecksumMissing(usize),
+    PreComputeDatasetChecksumMissing(String),
     #[error("Failed to decrypt dataset {0}")]
-    PreComputeDatasetDecryptionFailed(usize),
+    PreComputeDatasetDecryptionFailed(String),
     #[error("Failed to download encrypted dataset file for dataset {0}")]
-    PreComputeDatasetDownloadFailed(usize),
+    PreComputeDatasetDownloadFailed(String),
     #[error("Dataset filename related environment variable is missing for dataset {0}")]
-    PreComputeDatasetFilenameMissing(usize),
+    PreComputeDatasetFilenameMissing(String),
     #[error("Dataset key related environment variable is missing for dataset {0}")]
-    PreComputeDatasetKeyMissing(usize),
+    PreComputeDatasetKeyMissing(String),
     #[error("Dataset URL related environment variable is missing for dataset {0}")]
-    PreComputeDatasetUrlMissing(usize),
+    PreComputeDatasetUrlMissing(String),
     #[error("Unexpected error occurred")]
     PreComputeFailedUnknownIssue,
     #[error("Invalid TEE signature")]
@@ -32,7 +32,7 @@ pub enum ReplicateStatusCause {
     #[error("Input files number related environment variable is missing")]
     PreComputeInputFilesNumberMissing,
     #[error("Invalid dataset checksum for dataset {0}")]
-    PreComputeInvalidDatasetChecksum(usize),
+    PreComputeInvalidDatasetChecksum(String),
     #[error("Output folder related environment variable is missing")]
     PreComputeOutputFolderNotFound,
     #[error("Output path related environment variable is missing")]
@@ -64,13 +64,15 @@ mod tests {
     use super::*;
     use serde_json::to_string;
 
+    const DATASET_ADDRESS: &str = "0xDatasetAddress";
+
     #[test]
     fn serialize_produces_correct_json_when_error_has_dataset_index() {
-        let cause = ReplicateStatusCause::PreComputeDatasetUrlMissing(2);
+        let cause = ReplicateStatusCause::PreComputeDatasetUrlMissing(DATASET_ADDRESS.to_string());
         let serialized = to_string(&cause).unwrap();
         assert_eq!(
             serialized,
-            r#"{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 2"}"#
+            r#"{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 0xDatasetAddress"}"#
         );
     }
 
@@ -92,20 +94,22 @@ mod tests {
                 r#"{"cause":"PRE_COMPUTE_AT_LEAST_ONE_INPUT_FILE_URL_MISSING","message":"input file URL 1 is missing"}"#,
             ),
             (
-                ReplicateStatusCause::PreComputeDatasetChecksumMissing(3),
-                r#"{"cause":"PRE_COMPUTE_DATASET_CHECKSUM_MISSING","message":"Dataset checksum related environment variable is missing for dataset 3"}"#,
+                ReplicateStatusCause::PreComputeDatasetChecksumMissing(DATASET_ADDRESS.to_string()),
+                r#"{"cause":"PRE_COMPUTE_DATASET_CHECKSUM_MISSING","message":"Dataset checksum related environment variable is missing for dataset 0xDatasetAddress"}"#,
             ),
             (
-                ReplicateStatusCause::PreComputeDatasetDecryptionFailed(0),
-                r#"{"cause":"PRE_COMPUTE_DATASET_DECRYPTION_FAILED","message":"Failed to decrypt dataset 0"}"#,
+                ReplicateStatusCause::PreComputeDatasetDecryptionFailed(
+                    DATASET_ADDRESS.to_string(),
+                ),
+                r#"{"cause":"PRE_COMPUTE_DATASET_DECRYPTION_FAILED","message":"Failed to decrypt dataset 0xDatasetAddress"}"#,
             ),
             (
-                ReplicateStatusCause::PreComputeDatasetDownloadFailed(5),
-                r#"{"cause":"PRE_COMPUTE_DATASET_DOWNLOAD_FAILED","message":"Failed to download encrypted dataset file for dataset 5"}"#,
+                ReplicateStatusCause::PreComputeDatasetDownloadFailed(DATASET_ADDRESS.to_string()),
+                r#"{"cause":"PRE_COMPUTE_DATASET_DOWNLOAD_FAILED","message":"Failed to download encrypted dataset file for dataset 0xDatasetAddress"}"#,
             ),
             (
-                ReplicateStatusCause::PreComputeInvalidDatasetChecksum(2),
-                r#"{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 2"}"#,
+                ReplicateStatusCause::PreComputeInvalidDatasetChecksum(DATASET_ADDRESS.to_string()),
+                r#"{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 0xDatasetAddress"}"#,
             ),
         ];
 
@@ -118,12 +122,12 @@ mod tests {
     #[test]
     fn serialize_produces_correct_json_when_vector_of_multiple_errors() {
         let causes = vec![
-            ReplicateStatusCause::PreComputeDatasetUrlMissing(5),
-            ReplicateStatusCause::PreComputeInvalidDatasetChecksum(99),
+            ReplicateStatusCause::PreComputeDatasetUrlMissing(DATASET_ADDRESS.to_string()),
+            ReplicateStatusCause::PreComputeInvalidDatasetChecksum("0xAnotherDataset".to_string()),
         ];
 
         let serialized = to_string(&causes).unwrap();
-        let expected = r#"[{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 5"},{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 99"}]"#;
+        let expected = r#"[{"cause":"PRE_COMPUTE_DATASET_URL_MISSING","message":"Dataset URL related environment variable is missing for dataset 0xDatasetAddress"},{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 0xAnotherDataset"}]"#;
         assert_eq!(serialized, expected);
     }
 }
diff --git a/pre-compute/src/compute/pre_compute_app.rs b/pre-compute/src/compute/pre_compute_app.rs
@@ -59,10 +59,9 @@ impl PreComputeAppTrait for PreComputeApp {
         // TODO: Collect all errors instead of propagating immediately, and return the list of errors
         self.pre_compute_args = PreComputeArgs::read_args()?;
         self.check_output_folder()?;
-        for (index, dataset) in self.pre_compute_args.datasets.iter().enumerate() {
-            let encrypted_content =
-                dataset.download_encrypted_dataset(&self.chain_task_id, index)?;
-            let plain_content = dataset.decrypt_dataset(&encrypted_content, index)?;
+        for dataset in self.pre_compute_args.datasets.iter() {
+            let encrypted_content = dataset.download_encrypted_dataset(&self.chain_task_id)?;
+            let plain_content = dataset.decrypt_dataset(&encrypted_content)?;
             self.save_plain_dataset_file(&plain_content, &dataset.filename)?;
         }
         self.download_input_files()?;
@@ -190,6 +189,7 @@ mod tests {
                 is_dataset_required: true,
                 iexec_bulk_slice_size: 0,
                 datasets: vec![Dataset {
+                    address: "0xDatasetAddress".to_string(),
                     url: HTTP_DATASET_URL.to_string(),
                     checksum: DATASET_CHECKSUM.to_string(),
                     filename: PLAIN_DATA_FILE.to_string(),
diff --git a/pre-compute/src/compute/pre_compute_args.rs b/pre-compute/src/compute/pre_compute_args.rs
diff --git a/pre-compute/src/compute/utils/env_utils.rs b/pre-compute/src/compute/utils/env_utils.rs
