Skip to content
Merged
Show file tree
Hide file tree
Changes from 10 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions pre-compute/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ reqwest = { version = "0.12.15", features = ["blocking", "json"] }
serde = "1.0.219"
sha256 = "1.6.0"
sha3 = "0.10.8"
strum = "0.27.2"
strum_macros = "0.27.2"
thiserror = "2.0.12"

[dev-dependencies]
Expand Down
157 changes: 69 additions & 88 deletions pre-compute/src/api/worker_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,40 +4,6 @@ use crate::compute::{
};
use log::error;
use reqwest::{blocking::Client, header::AUTHORIZATION};
use serde::Serialize;

/// Represents payload that can be sent to the worker API to report the outcome of the
/// pre‑compute stage.
///
/// The JSON structure expected by the REST endpoint is:
/// ```json
/// {
/// "cause": "<ReplicateStatusCause as string>"
/// }
/// ```
///
/// # Arguments
///
/// * `cause` - A reference to the ReplicateStatusCause indicating why the pre-compute operation exited
///
/// # Example
///
/// ```rust
/// use tee_worker_pre_compute::api::worker_api::ExitMessage;
/// use tee_worker_pre_compute::compute::errors::ReplicateStatusCause;
///
/// let exit_message = ExitMessage::from(&ReplicateStatusCause::PreComputeInvalidTeeSignature);
/// ```
#[derive(Serialize, Debug)]
pub struct ExitMessage<'a> {
pub cause: &'a ReplicateStatusCause,
}

impl<'a> From<&'a ReplicateStatusCause> for ExitMessage<'a> {
fn from(cause: &'a ReplicateStatusCause) -> Self {
Self { cause }
}
}

/// Thin wrapper around a [`Client`] that knows how to reach the iExec worker API.
///
Expand Down Expand Up @@ -93,21 +59,21 @@ impl WorkerApiClient {
Self::new(&base_url)
}

/// Sends an exit cause for a pre-compute operation to the Worker API.
/// Sends exit causes for a pre-compute operation to the Worker API.
///
/// This method reports the exit cause of a pre-compute operation to the Worker API,
/// This method reports the exit causes of a pre-compute operation to the Worker API,
/// which can be used for tracking and debugging purposes.
///
/// # Arguments
///
/// * `authorization` - The authorization token to use for the API request
/// * `chain_task_id` - The chain task ID for which to report the exit cause
/// * `exit_cause` - The exit cause to report
/// * `chain_task_id` - The chain task ID for which to report the exit causes
/// * `exit_causes` - The list of exit causes to report
///
/// # Returns
///
/// * `Ok(())` - If the exit cause was successfully reported
/// * `Err(Error)` - If the exit cause could not be reported due to an HTTP error
/// * `Ok(())` - If the exit causes were successfully reported
/// * `Err(Error)` - If the exit causes could not be reported due to an HTTP error
///
/// # Errors
///
Expand All @@ -117,33 +83,33 @@ impl WorkerApiClient {
/// # Example
///
/// ```rust
/// use tee_worker_pre_compute::api::worker_api::{ExitMessage, WorkerApiClient};
/// use tee_worker_pre_compute::api::worker_api::WorkerApiClient;
/// use tee_worker_pre_compute::compute::errors::ReplicateStatusCause;
///
/// let client = WorkerApiClient::new("http://worker:13100");
/// let exit_message = ExitMessage::from(&ReplicateStatusCause::PreComputeInvalidTeeSignature);
/// let exit_causes = vec![ReplicateStatusCause::PreComputeInvalidTeeSignature];
///
/// match client.send_exit_cause_for_pre_compute_stage(
/// match client.send_exit_causes_for_pre_compute_stage(
/// "authorization_token",
/// "0x123456789abcdef",
/// &exit_message,
/// &exit_causes,
/// ) {
/// Ok(()) => println!("Exit cause reported successfully"),
/// Err(error) => eprintln!("Failed to report exit cause: {error}"),
/// Ok(()) => println!("Exit causes reported successfully"),
/// Err(error) => eprintln!("Failed to report exit causes: {error}"),
/// }
/// ```
pub fn send_exit_cause_for_pre_compute_stage(
pub fn send_exit_causes_for_pre_compute_stage(
&self,
authorization: &str,
chain_task_id: &str,
exit_cause: &ExitMessage,
exit_causes: &[ReplicateStatusCause],
) -> Result<(), ReplicateStatusCause> {
let url = format!("{}/compute/pre/{chain_task_id}/exit", self.base_url);
let url = format!("{}/compute/pre/{chain_task_id}/exit-causes", self.base_url);
match self
.client
.post(&url)
.header(AUTHORIZATION, authorization)
.json(exit_cause)
.json(exit_causes)
.send()
{
Ok(resp) => {
Expand All @@ -152,12 +118,12 @@ impl WorkerApiClient {
Ok(())
} else {
let body = resp.text().unwrap_or_default();
error!("Failed to send exit cause: [status:{status}, body:{body}]");
error!("Failed to send exit causes: [status:{status}, body:{body}]");
Err(ReplicateStatusCause::PreComputeFailedUnknownIssue)
}
}
Err(err) => {
error!("HTTP request failed when sending exit cause to {url}: {err:?}");
error!("HTTP request failed when sending exit causes to {url}: {err:?}");
Err(ReplicateStatusCause::PreComputeFailedUnknownIssue)
}
}
Expand All @@ -175,36 +141,50 @@ mod tests {
matchers::{body_json, header, method, path},
};

// region ExitMessage()
// region Serialization tests
#[test]
fn should_serialize_exit_message() {
let causes = [
fn serialize_replicate_status_cause_succeeds_when_single_cause() {
let causes = vec![
(
ReplicateStatusCause::PreComputeInvalidTeeSignature,
"PRE_COMPUTE_INVALID_TEE_SIGNATURE",
r#"{"cause":"PRE_COMPUTE_INVALID_TEE_SIGNATURE","message":"Invalid TEE signature"}"#,
),
(
ReplicateStatusCause::PreComputeWorkerAddressMissing,
"PRE_COMPUTE_WORKER_ADDRESS_MISSING",
r#"{"cause":"PRE_COMPUTE_WORKER_ADDRESS_MISSING","message":"Worker address related environment variable is missing"}"#,
),
(
ReplicateStatusCause::PreComputeFailedUnknownIssue,
"PRE_COMPUTE_FAILED_UNKNOWN_ISSUE",
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("0xDataset1".to_string()),
r#"{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 0xDataset1"}"#,
),
];

for (cause, message) in causes {
let exit_message = ExitMessage::from(&cause);
let serialized = to_string(&exit_message).expect("Failed to serialize");
let expected = format!("{{\"cause\":\"{message}\"}}");
assert_eq!(serialized, expected);
for (cause, expected_json) in causes {
let serialized = to_string(&cause).expect("Failed to serialize");
assert_eq!(serialized, expected_json);
}
}

#[test]
fn serialize_vec_of_causes_succeeds_when_multiple_causes() {
let causes = vec![
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 0xDatasetA"},{"cause":"PRE_COMPUTE_INVALID_DATASET_CHECKSUM","message":"Invalid dataset checksum for dataset 0xDatasetB"}]"#;
assert_eq!(serialized, expected);
}
// endregion

// region get_worker_api_client
#[test]
fn should_get_worker_api_client_with_env_var() {
fn from_env_creates_client_with_custom_host_when_env_var_set() {
with_vars(
vec![(WorkerHostEnvVar.name(), Some("custom-worker-host:9999"))],
|| {
Expand All @@ -215,29 +195,32 @@ mod tests {
}

#[test]
fn should_get_worker_api_client_without_env_var() {
fn from_env_creates_client_with_default_host_when_env_var_unset() {
temp_env::with_vars_unset(vec![WorkerHostEnvVar.name()], || {
let client = WorkerApiClient::from_env();
assert_eq!(client.base_url, format!("http://{DEFAULT_WORKER_HOST}"));
});
}
// endregion

// region send_exit_cause_for_pre_compute_stage()
// region send_exit_causes_for_pre_compute_stage()
const CHALLENGE: &str = "challenge";
const CHAIN_TASK_ID: &str = "0x123456789abcdef";

#[tokio::test]
async fn should_send_exit_cause() {
async fn send_exit_causes_succeeds_when_api_returns_success() {
let mock_server = MockServer::start().await;
let server_url = mock_server.uri();

let expected_body = json!({
"cause": ReplicateStatusCause::PreComputeInvalidTeeSignature,
});
let expected_body = json!([
{
"cause": "PRE_COMPUTE_INVALID_TEE_SIGNATURE",
"message": "Invalid TEE signature"
}
]);

Mock::given(method("POST"))
.and(path(format!("/compute/pre/{CHAIN_TASK_ID}/exit")))
.and(path(format!("/compute/pre/{CHAIN_TASK_ID}/exit-causes")))
.and(header("Authorization", CHALLENGE))
.and(body_json(&expected_body))
.respond_with(ResponseTemplate::new(200))
Expand All @@ -246,13 +229,12 @@ mod tests {
.await;

let result = tokio::task::spawn_blocking(move || {
let exit_message =
ExitMessage::from(&ReplicateStatusCause::PreComputeInvalidTeeSignature);
let exit_causes = vec![ReplicateStatusCause::PreComputeInvalidTeeSignature];
let worker_api_client = WorkerApiClient::new(&server_url);
worker_api_client.send_exit_cause_for_pre_compute_stage(
worker_api_client.send_exit_causes_for_pre_compute_stage(
CHALLENGE,
CHAIN_TASK_ID,
&exit_message,
&exit_causes,
)
})
.await
Expand All @@ -262,26 +244,25 @@ mod tests {
}

#[tokio::test]
async fn should_not_send_exit_cause() {
async fn send_exit_causes_fails_when_api_returns_error() {
testing_logger::setup();
let mock_server = MockServer::start().await;
let server_url = mock_server.uri();

Mock::given(method("POST"))
.and(path(format!("/compute/pre/{CHAIN_TASK_ID}/exit")))
.and(path(format!("/compute/pre/{CHAIN_TASK_ID}/exit-causes")))
.respond_with(ResponseTemplate::new(503).set_body_string("Service Unavailable"))
.expect(1)
.mount(&mock_server)
.await;

let result = tokio::task::spawn_blocking(move || {
let exit_message =
ExitMessage::from(&ReplicateStatusCause::PreComputeFailedUnknownIssue);
let exit_causes = vec![ReplicateStatusCause::PreComputeFailedUnknownIssue];
let worker_api_client = WorkerApiClient::new(&server_url);
let response = worker_api_client.send_exit_cause_for_pre_compute_stage(
let response = worker_api_client.send_exit_causes_for_pre_compute_stage(
CHALLENGE,
CHAIN_TASK_ID,
&exit_message,
&exit_causes,
);
testing_logger::validate(|captured_logs| {
let logs = captured_logs
Expand All @@ -292,7 +273,7 @@ mod tests {
assert_eq!(logs.len(), 1);
assert_eq!(
logs[0].body,
"Failed to send exit cause: [status:503 Service Unavailable, body:Service Unavailable]"
"Failed to send exit causes: [status:503 Service Unavailable, body:Service Unavailable]"
);
});
response
Expand All @@ -308,14 +289,14 @@ mod tests {
}

#[test]
fn test_send_exit_cause_http_request_failure() {
fn send_exit_causes_fails_when_http_request_invalid() {
testing_logger::setup();
let exit_message = ExitMessage::from(&ReplicateStatusCause::PreComputeFailedUnknownIssue);
let exit_causes = vec![ReplicateStatusCause::PreComputeFailedUnknownIssue];
let worker_api_client = WorkerApiClient::new("wrong_url");
let result = worker_api_client.send_exit_cause_for_pre_compute_stage(
let result = worker_api_client.send_exit_causes_for_pre_compute_stage(
CHALLENGE,
CHAIN_TASK_ID,
&exit_message,
&exit_causes,
);
testing_logger::validate(|captured_logs| {
let logs = captured_logs
Expand All @@ -326,7 +307,7 @@ mod tests {
assert_eq!(logs.len(), 1);
assert_eq!(
logs[0].body,
"HTTP request failed when sending exit cause to wrong_url/compute/pre/0x123456789abcdef/exit: reqwest::Error { kind: Builder, source: RelativeUrlWithoutBase }"
"HTTP request failed when sending exit causes to wrong_url/compute/pre/0x123456789abcdef/exit-causes: reqwest::Error { kind: Builder, source: RelativeUrlWithoutBase }"
);
});
assert!(result.is_err());
Expand Down
Loading