Skip to content

Commit 54af079

Browse files
nabil-Tounartinabil-Tounarti
authored
refactor(pre-compute-app): remove unnecessary Option fields (#31)
1 parent 239ba9e commit 54af079

File tree

3 files changed

+71
-101
lines changed

3 files changed

+71
-101
lines changed

‎src/compute/app_runner.rs

Lines changed: 40 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -27,33 +27,24 @@ pub enum ExitMode {
2727
/// It uses the provided app to execute core operations and handles all the
2828
/// workflow states and transitions.
2929
///
30-
/// # Arguments
31-
///
32-
/// * `pre_compute_app` - An implementation of [`PreComputeAppTrait`] that will be used to execute the pre-compute operations.
33-
///
3430
/// # Example
3531
///
3632
/// ```
3733
/// use crate::app_runner::start;
3834
/// use crate::pre_compute_app::PreComputeApp;
3935
///
40-
/// let pre_compute_app = PreComputeApp::new();
41-
/// let exit_code = start_with_app(pre_compute_app);
36+
/// let chain_task_id = "0x123456789abcdef".to_string();
37+
/// let mut pre_compute_app = PreComputeApp::new(chain_task_id.clone());
38+
///
39+
/// let exit_code = start_with_app(&pre_compute_app, &chain_task_id)
4240
/// ```
43-
pub fn start_with_app<A: PreComputeAppTrait>(pre_compute_app: &mut A) -> ExitMode {
44-
info!("TEE pre-compute started");
45-
41+
pub fn start_with_app<A: PreComputeAppTrait>(
42+
pre_compute_app: &mut A,
43+
chain_task_id: &str,
44+
) -> ExitMode {
4645
let exit_cause = ReplicateStatusCause::PreComputeFailedUnknownIssue;
47-
let chain_task_id =
48-
match get_env_var_or_error(IexecTaskId, ReplicateStatusCause::PreComputeTaskIdMissing) {
49-
Ok(id) => id,
50-
Err(e) => {
51-
error!("TEE pre-compute cannot proceed without taskID context: {e:?}");
52-
return ExitMode::InitializationFailure;
53-
}
54-
};
5546

56-
match pre_compute_app.run(&chain_task_id) {
47+
match pre_compute_app.run() {
5748
Ok(_) => {
5849
info!("TEE pre-compute completed");
5950
return ExitMode::Success;
@@ -63,7 +54,7 @@ pub fn start_with_app<A: PreComputeAppTrait>(pre_compute_app: &mut A) -> ExitMod
6354
}
6455
}
6556

66-
let authorization = match get_challenge(&chain_task_id) {
57+
let authorization = match get_challenge(chain_task_id) {
6758
Ok(auth) => auth,
6859
Err(_) => {
6960
error!("Failed to sign exitCause message [{exit_cause:?}]");
@@ -77,7 +68,7 @@ pub fn start_with_app<A: PreComputeAppTrait>(pre_compute_app: &mut A) -> ExitMod
7768

7869
match WorkerApiClient::from_env().send_exit_cause_for_pre_compute_stage(
7970
&authorization,
80-
&chain_task_id,
71+
chain_task_id,
8172
&exit_message,
8273
) {
8374
Ok(_) => ExitMode::ReportedFailure,
@@ -102,8 +93,19 @@ pub fn start_with_app<A: PreComputeAppTrait>(pre_compute_app: &mut A) -> ExitMod
10293
/// std::process::exit(exit_code);
10394
/// ```
10495
pub fn start() -> ExitMode {
105-
let mut pre_compute_app = PreComputeApp::new();
106-
start_with_app(&mut pre_compute_app)
96+
info!("TEE pre-compute started");
97+
98+
let chain_task_id =
99+
match get_env_var_or_error(IexecTaskId, ReplicateStatusCause::PreComputeTaskIdMissing) {
100+
Ok(id) => id,
101+
Err(e) => {
102+
error!("TEE pre-compute cannot proceed without taskID context: {e:?}");
103+
return ExitMode::InitializationFailure;
104+
}
105+
};
106+
let mut pre_compute_app = PreComputeApp::new(chain_task_id.clone());
107+
108+
start_with_app(&mut pre_compute_app, &chain_task_id)
107109
}
108110

109111
#[cfg(test)]
@@ -116,13 +118,17 @@ mod pre_compute_start_with_app_tests {
116118
use wiremock::{Mock, MockServer, ResponseTemplate};
117119

118120
const CHAIN_TASK_ID: &str = "0x123456789abcdef";
119-
const WORKER_ADDRESS: &str = "0xabcdef123456789";
120121
const ENCLAVE_CHALLENGE_PRIVATE_KEY: &str =
121122
"0xdd3b993ec21c71c1f6d63a5240850e0d4d8dd83ff70d29e49247958548c1d479";
122123
const ENV_IEXEC_TASK_ID: &str = "IEXEC_TASK_ID";
123124
const ENV_SIGN_WORKER_ADDRESS: &str = "SIGN_WORKER_ADDRESS";
124125
const ENV_SIGN_TEE_CHALLENGE_PRIVATE_KEY: &str = "SIGN_TEE_CHALLENGE_PRIVATE_KEY";
125126
const ENV_WORKER_HOST: &str = "WORKER_HOST_ENV_VAR";
127+
const IEXEC_INPUT_FILES_NUMBER: &str = "IEXEC_INPUT_FILES_NUMBER";
128+
const IEXEC_PRE_COMPUTE_OUT: &str = "IEXEC_PRE_COMPUTE_OUT";
129+
const IS_DATASET_REQUIRED: &str = "IS_DATASET_REQUIRED";
130+
const OUTPUT_DIR: &str = "/iexec_out";
131+
const WORKER_ADDRESS: &str = "0xabcdef123456789";
126132

127133
#[test]
128134
fn start_fails_when_task_id_missing() {
@@ -148,13 +154,12 @@ mod pre_compute_start_with_app_tests {
148154

149155
let mut mock = MockPreComputeAppTrait::new();
150156
mock.expect_run()
151-
.withf(|chain_task_id| chain_task_id == CHAIN_TASK_ID)
152-
.returning(|_| Err(ReplicateStatusCause::PreComputeWorkerAddressMissing));
157+
.returning(|| Err(ReplicateStatusCause::PreComputeWorkerAddressMissing));
153158

154159
temp_env::with_vars(env_vars_to_set, || {
155160
temp_env::with_vars_unset(env_vars_to_unset, || {
156161
assert_eq!(
157-
start_with_app(&mut mock),
162+
start_with_app(&mut mock, CHAIN_TASK_ID),
158163
ExitMode::UnreportedFailure,
159164
"Should return 2 if get_challenge fails due to missing signer address"
160165
);
@@ -172,13 +177,12 @@ mod pre_compute_start_with_app_tests {
172177

173178
let mut mock = MockPreComputeAppTrait::new();
174179
mock.expect_run()
175-
.withf(|chain_task_id| chain_task_id == CHAIN_TASK_ID)
176-
.returning(|_| Err(ReplicateStatusCause::PreComputeTeeChallengePrivateKeyMissing));
180+
.returning(|| Err(ReplicateStatusCause::PreComputeTeeChallengePrivateKeyMissing));
177181

178182
temp_env::with_vars(env_vars_to_set, || {
179183
temp_env::with_vars_unset(env_vars_to_unset, || {
180184
assert_eq!(
181-
start_with_app(&mut mock),
185+
start_with_app(&mut mock, CHAIN_TASK_ID),
182186
ExitMode::UnreportedFailure,
183187
"Should return 2 if get_challenge fails due to missing private key"
184188
);
@@ -200,8 +204,7 @@ mod pre_compute_start_with_app_tests {
200204

201205
let mut mock = MockPreComputeAppTrait::new();
202206
mock.expect_run()
203-
.withf(|chain_task_id| chain_task_id == CHAIN_TASK_ID)
204-
.returning(|_| Err(ReplicateStatusCause::PreComputeTeeChallengePrivateKeyMissing));
207+
.returning(|| Err(ReplicateStatusCause::PreComputeTeeChallengePrivateKeyMissing));
205208

206209
let result_code = tokio::task::spawn_blocking(move || {
207210
let env_vars = vec![
@@ -214,7 +217,7 @@ mod pre_compute_start_with_app_tests {
214217
(ENV_WORKER_HOST, Some(mock_server_addr_string.as_str())),
215218
];
216219

217-
temp_env::with_vars(env_vars, || start_with_app(&mut mock))
220+
temp_env::with_vars(env_vars, || start_with_app(&mut mock, CHAIN_TASK_ID))
218221
})
219222
.await
220223
.expect("Blocking task panicked");
@@ -247,8 +250,7 @@ mod pre_compute_start_with_app_tests {
247250

248251
let mut mock = MockPreComputeAppTrait::new();
249252
mock.expect_run()
250-
.withf(|chain_task_id| chain_task_id == CHAIN_TASK_ID)
251-
.returning(|_| Err(ReplicateStatusCause::PreComputeTeeChallengePrivateKeyMissing));
253+
.returning(|| Err(ReplicateStatusCause::PreComputeTeeChallengePrivateKeyMissing));
252254

253255
// Move the blocking operations into spawn_blocking
254256
let result_code = tokio::task::spawn_blocking(move || {
@@ -260,9 +262,12 @@ mod pre_compute_start_with_app_tests {
260262
Some(ENCLAVE_CHALLENGE_PRIVATE_KEY),
261263
),
262264
(ENV_WORKER_HOST, Some(mock_server_addr_string.as_str())),
265+
(IEXEC_INPUT_FILES_NUMBER, Some("0")),
266+
(IEXEC_PRE_COMPUTE_OUT, Some(OUTPUT_DIR)),
267+
(IS_DATASET_REQUIRED, Some("false")),
263268
];
264269

265-
temp_env::with_vars(env_vars, || start_with_app(&mut mock))
270+
temp_env::with_vars(env_vars, start)
266271
})
267272
.await
268273
.expect("Blocking task panicked");

‎src/compute/pre_compute_app.rs

Lines changed: 29 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ const AES_IV_LENGTH: usize = 16;
2626

2727
#[cfg_attr(test, automock)]
2828
pub trait PreComputeAppTrait {
29-
fn run(&mut self, chain_task_id: &str) -> Result<(), ReplicateStatusCause>;
29+
fn run(&mut self) -> Result<(), ReplicateStatusCause>;
3030
fn check_output_folder(&self) -> Result<(), ReplicateStatusCause>;
3131
fn download_input_files(&self) -> Result<(), ReplicateStatusCause>;
3232
fn download_encrypted_dataset(&self) -> Result<Vec<u8>, ReplicateStatusCause>;
@@ -35,25 +35,24 @@ pub trait PreComputeAppTrait {
3535
}
3636

3737
pub struct PreComputeApp {
38-
chain_task_id: Option<String>,
39-
pre_compute_args: Option<PreComputeArgs>,
38+
chain_task_id: String,
39+
pre_compute_args: PreComputeArgs,
4040
}
4141

4242
impl PreComputeApp {
43-
pub fn new() -> Self {
43+
pub fn new(chain_task_id: String) -> Self {
4444
PreComputeApp {
45-
chain_task_id: None,
46-
pre_compute_args: None,
45+
chain_task_id,
46+
pre_compute_args: PreComputeArgs::default(),
4747
}
4848
}
4949
}
5050

5151
impl PreComputeAppTrait for PreComputeApp {
52-
fn run(&mut self, chain_task_id: &str) -> Result<(), ReplicateStatusCause> {
53-
self.chain_task_id = Some(chain_task_id.to_string());
54-
self.pre_compute_args = Some(PreComputeArgs::read_args()?);
52+
fn run(&mut self) -> Result<(), ReplicateStatusCause> {
53+
self.pre_compute_args = PreComputeArgs::read_args()?;
5554
self.check_output_folder()?;
56-
if self.pre_compute_args.as_ref().unwrap().is_dataset_required {
55+
if self.pre_compute_args.is_dataset_required {
5756
let encrypted_content = self.download_encrypted_dataset()?;
5857
let plain_content = self.decrypt_dataset(&encrypted_content)?;
5958
self.save_plain_dataset_file(&plain_content)?;
@@ -82,14 +81,8 @@ impl PreComputeAppTrait for PreComputeApp {
8281
/// pre_compute_app.check_output_folder()?;
8382
/// ```
8483
fn check_output_folder(&self) -> Result<(), ReplicateStatusCause> {
85-
let output_dir = self
86-
.pre_compute_args
87-
.as_ref()
88-
.ok_or(ReplicateStatusCause::PreComputeOutputFolderNotFound)?
89-
.output_dir
90-
.clone();
91-
92-
let chain_task_id = self.chain_task_id.as_deref().unwrap_or("unknown");
84+
let output_dir: &str = &self.pre_compute_args.output_dir;
85+
let chain_task_id: &str = &self.chain_task_id;
9386

9487
info!("Checking output folder [chainTaskId:{chain_task_id}, path:{output_dir}]");
9588

@@ -130,8 +123,8 @@ impl PreComputeAppTrait for PreComputeApp {
130123
/// pre_compute_app.download_input_files()?;
131124
/// ```
132125
fn download_input_files(&self) -> Result<(), ReplicateStatusCause> {
133-
let args = self.pre_compute_args.as_ref().unwrap();
134-
let chain_task_id = self.chain_task_id.as_ref().unwrap();
126+
let args = &self.pre_compute_args;
127+
let chain_task_id: &str = &self.chain_task_id;
135128

136129
for url in &args.input_files {
137130
info!("Downloading input file [chainTaskId:{chain_task_id}, url:{url}]");
@@ -162,8 +155,8 @@ impl PreComputeAppTrait for PreComputeApp {
162155
/// app.download_encrypted_dataset()?;
163156
/// ```
164157
fn download_encrypted_dataset(&self) -> Result<Vec<u8>, ReplicateStatusCause> {
165-
let args = self.pre_compute_args.as_ref().unwrap();
166-
let chain_task_id: &str = self.chain_task_id.as_ref().unwrap();
158+
let args = &self.pre_compute_args;
159+
let chain_task_id = &self.chain_task_id;
167160
let encrypted_dataset_url: &str = &args.encrypted_dataset_url;
168161

169162
info!(
@@ -228,11 +221,7 @@ impl PreComputeAppTrait for PreComputeApp {
228221
/// let decrypted = app.decrypt_dataset(&encrypted)?;
229222
/// ```
230223
fn decrypt_dataset(&self, encrypted_content: &[u8]) -> Result<Vec<u8>, ReplicateStatusCause> {
231-
let base64_key: &str = &self
232-
.pre_compute_args
233-
.as_ref()
234-
.unwrap()
235-
.encrypted_dataset_base64_key;
224+
let base64_key: &str = &self.pre_compute_args.encrypted_dataset_base64_key;
236225

237226
let key = general_purpose::STANDARD
238227
.decode(base64_key)
@@ -275,8 +264,8 @@ impl PreComputeAppTrait for PreComputeApp {
275264
/// app.save_plain_dataset_file(&plain_data)?;
276265
/// ```
277266
fn save_plain_dataset_file(&self, plain_dataset: &[u8]) -> Result<(), ReplicateStatusCause> {
278-
let chain_task_id: &str = self.chain_task_id.as_ref().unwrap();
279-
let args = self.pre_compute_args.as_ref().unwrap();
267+
let chain_task_id: &str = &self.chain_task_id;
268+
let args = &self.pre_compute_args;
280269
let output_dir: &str = &args.output_dir;
281270
let plain_dataset_filename: &str = &args.plain_dataset_filename;
282271

@@ -325,16 +314,16 @@ mod tests {
325314
output_dir: &str,
326315
) -> PreComputeApp {
327316
PreComputeApp {
328-
chain_task_id: Some(chain_task_id.to_string()),
329-
pre_compute_args: Some(PreComputeArgs {
317+
chain_task_id: chain_task_id.to_string(),
318+
pre_compute_args: PreComputeArgs {
330319
input_files: urls.into_iter().map(String::from).collect(),
331320
output_dir: output_dir.to_string(),
332321
is_dataset_required: true,
333322
encrypted_dataset_url: HTTP_DATASET_URL.to_string(),
334323
encrypted_dataset_base64_key: ENCRYPTED_DATASET_KEY.to_string(),
335324
encrypted_dataset_checksum: DATASET_CHECKSUM.to_string(),
336325
plain_dataset_filename: PLAIN_DATA_FILE.to_string(),
337-
}),
326+
},
338327
}
339328
}
340329

@@ -378,19 +367,6 @@ mod tests {
378367
);
379368
}
380369

381-
#[test]
382-
fn check_output_folder_returns_err_with_invalid_pre_compute_args() {
383-
let app = PreComputeApp {
384-
chain_task_id: Some(CHAIN_TASK_ID.to_string()),
385-
pre_compute_args: None,
386-
};
387-
388-
let result = app.check_output_folder();
389-
assert_eq!(
390-
result,
391-
Err(ReplicateStatusCause::PreComputeOutputFolderNotFound)
392-
);
393-
}
394370
// endregion
395371

396372
// region download_input_files
@@ -495,9 +471,7 @@ mod tests {
495471
#[test]
496472
fn download_encrypted_dataset_failure_with_invalid_dataset_url() {
497473
let mut app = get_pre_compute_app(CHAIN_TASK_ID, vec![], "");
498-
if let Some(args) = &mut app.pre_compute_args {
499-
args.encrypted_dataset_url = "http://bad-url".to_string();
500-
}
474+
app.pre_compute_args.encrypted_dataset_url = "http://bad-url".to_string();
501475
let actual_content = app.download_encrypted_dataset();
502476
assert_eq!(
503477
actual_content,
@@ -508,11 +482,9 @@ mod tests {
508482
#[test]
509483
fn download_encrypted_dataset_success_with_valid_iexec_gateway() {
510484
let mut app = get_pre_compute_app(CHAIN_TASK_ID, vec![], "");
511-
if let Some(args) = &mut app.pre_compute_args {
512-
args.encrypted_dataset_url = IPFS_DATASET_URL.to_string();
513-
args.encrypted_dataset_checksum =
514-
"0x323b1637c7999942fbebfe5d42fe15dbfe93737577663afa0181938d7ad4a2ac".to_string();
515-
}
485+
app.pre_compute_args.encrypted_dataset_url = IPFS_DATASET_URL.to_string();
486+
app.pre_compute_args.encrypted_dataset_checksum =
487+
"0x323b1637c7999942fbebfe5d42fe15dbfe93737577663afa0181938d7ad4a2ac".to_string();
516488
let actual_content = app.download_encrypted_dataset();
517489
let expected_content = Ok("hello world !\n".as_bytes().to_vec());
518490
assert_eq!(actual_content, expected_content);
@@ -521,9 +493,7 @@ mod tests {
521493
#[test]
522494
fn download_encrypted_dataset_failure_with_invalid_gateway() {
523495
let mut app = get_pre_compute_app(CHAIN_TASK_ID, vec![], "");
524-
if let Some(args) = &mut app.pre_compute_args {
525-
args.encrypted_dataset_url = "/ipfs/INVALID_IPFS_DATASET_URL".to_string();
526-
}
496+
app.pre_compute_args.encrypted_dataset_url = "/ipfs/INVALID_IPFS_DATASET_URL".to_string();
527497
let actual_content = app.download_encrypted_dataset();
528498
let expected_content = Err(ReplicateStatusCause::PreComputeDatasetDownloadFailed);
529499
assert_eq!(actual_content, expected_content);
@@ -532,9 +502,7 @@ mod tests {
532502
#[test]
533503
fn download_encrypted_dataset_failure_with_invalid_dataset_checksum() {
534504
let mut app = get_pre_compute_app(CHAIN_TASK_ID, vec![], "");
535-
if let Some(args) = &mut app.pre_compute_args {
536-
args.encrypted_dataset_checksum = "invalid_dataset_checksum".to_string()
537-
}
505+
app.pre_compute_args.encrypted_dataset_checksum = "invalid_dataset_checksum".to_string();
538506
let actual_content = app.download_encrypted_dataset();
539507
let expected_content = Err(ReplicateStatusCause::PreComputeInvalidDatasetChecksum);
540508
assert_eq!(actual_content, expected_content);
@@ -556,9 +524,7 @@ mod tests {
556524
#[test]
557525
fn decrypt_dataset_failure_with_bad_key() {
558526
let mut app = get_pre_compute_app(CHAIN_TASK_ID, vec![], "");
559-
if let Some(args) = &mut app.pre_compute_args {
560-
args.encrypted_dataset_base64_key = "bad_key".to_string();
561-
}
527+
app.pre_compute_args.encrypted_dataset_base64_key = "bad_key".to_string();
562528
let encrypted_data = app.download_encrypted_dataset().unwrap();
563529
let actual_plain_data = app.decrypt_dataset(&encrypted_data);
564530

@@ -602,9 +568,7 @@ mod tests {
602568
let output_path = temp_dir.path().to_str().unwrap();
603569

604570
let mut app = get_pre_compute_app(CHAIN_TASK_ID, vec![], output_path);
605-
if let Some(args) = &mut app.pre_compute_args {
606-
args.plain_dataset_filename = "/some-folder-123/not-found".to_string();
607-
}
571+
app.pre_compute_args.plain_dataset_filename = "/some-folder-123/not-found".to_string();
608572
let plain_dataset = "Some very useful data.".as_bytes().to_vec();
609573
let saved_dataset = app.save_plain_dataset_file(&plain_dataset);
610574

0 commit comments

Comments
 (0)