|
| 1 | +use crate::compute::{ |
| 2 | + errors::ReplicateStatusCause, |
| 3 | + utils::env_utils::{TeeSessionEnvironmentVariable, get_env_var_or_error}, |
| 4 | +}; |
| 5 | +use reqwest::{Error, blocking::Client, header::AUTHORIZATION}; |
| 6 | +use serde::Serialize; |
| 7 | + |
| 8 | +/// Represents payload that can be sent to the worker API to report the outcome of the |
| 9 | +/// post‑compute stage. |
| 10 | +/// |
| 11 | +/// The JSON structure expected by the REST endpoint is: |
| 12 | +/// ```json |
| 13 | +/// { |
| 14 | +/// "cause": "<ReplicateStatusCause as string>" |
| 15 | +/// } |
| 16 | +/// ``` |
| 17 | +/// |
| 18 | +/// # Arguments |
| 19 | +/// |
| 20 | +/// * `cause` - A reference to the ReplicateStatusCause indicating why the post-compute operation exited |
| 21 | +/// |
| 22 | +/// # Example |
| 23 | +/// |
| 24 | +/// ``` |
| 25 | +/// use crate::api::worker_api::ExitMessage; |
| 26 | +/// use crate::compute::errors::ReplicateStatusCause; |
| 27 | +/// |
| 28 | +/// let exit_message = ExitMessage::from(&ReplicateStatusCause::PostComputeInvalidTeeSignature); |
| 29 | +/// ``` |
| 30 | +#[derive(Serialize, Debug)] |
| 31 | +pub struct ExitMessage<'a> { |
| 32 | + #[serde(rename = "cause")] |
| 33 | + pub cause: &'a ReplicateStatusCause, |
| 34 | +} |
| 35 | + |
| 36 | +impl<'a> From<&'a ReplicateStatusCause> for ExitMessage<'a> { |
| 37 | + fn from(cause: &'a ReplicateStatusCause) -> Self { |
| 38 | + Self { cause } |
| 39 | + } |
| 40 | +} |
| 41 | + |
| 42 | +/// Thin wrapper around a [`Client`] that knows how to reach the iExec worker API. |
| 43 | +/// |
| 44 | +/// This client can be created directly with a base URL using [`new()`], or |
| 45 | +/// configured from environment variables using [`from_env()`]. |
| 46 | +/// |
| 47 | +/// # Example |
| 48 | +/// |
| 49 | +/// ``` |
| 50 | +/// use crate::api::worker_api::WorkerApiClient; |
| 51 | +/// |
| 52 | +/// let client = WorkerApiClient::new("http://worker:13100"); |
| 53 | +/// ``` |
| 54 | +pub struct WorkerApiClient { |
| 55 | + base_url: String, |
| 56 | + client: Client, |
| 57 | +} |
| 58 | + |
| 59 | +const DEFAULT_WORKER_HOST: &str = "worker:13100"; |
| 60 | + |
| 61 | +impl WorkerApiClient { |
| 62 | + pub fn new(base_url: &str) -> Self { |
| 63 | + WorkerApiClient { |
| 64 | + base_url: base_url.to_string(), |
| 65 | + client: Client::builder().build().unwrap(), |
| 66 | + } |
| 67 | + } |
| 68 | + |
| 69 | + /// Creates a new WorkerApiClient instance with configuration from environment variables. |
| 70 | + /// |
| 71 | + /// This method retrieves the worker host from the [`WORKER_HOST_ENV_VAR`] environment variable. |
| 72 | + /// If the variable is not set or empty, it defaults to `"worker:13100"`. |
| 73 | + /// |
| 74 | + /// # Returns |
| 75 | + /// |
| 76 | + /// * `WorkerApiClient` - A new client configured with the appropriate base URL |
| 77 | + /// |
| 78 | + /// # Example |
| 79 | + /// |
| 80 | + /// ``` |
| 81 | + /// use crate::api::worker_api::WorkerApiClient; |
| 82 | + /// |
| 83 | + /// let client = WorkerApiClient::from_env(); |
| 84 | + /// ``` |
| 85 | + pub fn from_env() -> Self { |
| 86 | + let worker_host = get_env_var_or_error( |
| 87 | + TeeSessionEnvironmentVariable::WORKER_HOST_ENV_VAR, |
| 88 | + ReplicateStatusCause::PostComputeWorkerAddressMissing, |
| 89 | + ) |
| 90 | + .unwrap_or_else(|_| DEFAULT_WORKER_HOST.to_string()); |
| 91 | + |
| 92 | + let base_url = format!("http://{}", &worker_host); |
| 93 | + Self::new(&base_url) |
| 94 | + } |
| 95 | + |
| 96 | + /// Sends an exit cause for a post-compute operation to the Worker API. |
| 97 | + /// |
| 98 | + /// This method reports the exit cause of a post-compute operation to the Worker API, |
| 99 | + /// which can be used for tracking and debugging purposes. |
| 100 | + /// |
| 101 | + /// # Arguments |
| 102 | + /// |
| 103 | + /// * `authorization` - The authorization token to use for the API request |
| 104 | + /// * `chain_task_id` - The chain task ID for which to report the exit cause |
| 105 | + /// * `exit_cause` - The exit cause to report |
| 106 | + /// |
| 107 | + /// # Returns |
| 108 | + /// |
| 109 | + /// * `Ok(())` - If the exit cause was successfully reported |
| 110 | + /// * `Err(Error)` - If the exit cause could not be reported due to an HTTP error |
| 111 | + /// |
| 112 | + /// # Errors |
| 113 | + /// |
| 114 | + /// This function will return an [`Error`] if the request could not be sent or |
| 115 | + /// the server responded with a non‑success status. |
| 116 | + /// |
| 117 | + /// # Example |
| 118 | + /// |
| 119 | + /// ``` |
| 120 | + /// use crate::api::worker_api::{ExitMessage, WorkerApiClient}; |
| 121 | + /// use crate::compute::errors::ReplicateStatusCause; |
| 122 | + /// |
| 123 | + /// let client = WorkerApiClient::new("http://worker:13100"); |
| 124 | + /// let exit_message = ExitMessage::from(&ReplicateStatusCause::PostComputeInvalidTeeSignature); |
| 125 | + /// |
| 126 | + /// match client.send_exit_cause_for_post_compute_stage( |
| 127 | + /// "authorization_token", |
| 128 | + /// "0x123456789abcdef", |
| 129 | + /// &exit_message, |
| 130 | + /// ) { |
| 131 | + /// Ok(()) => println!("Exit cause reported successfully"), |
| 132 | + /// Err(error) => eprintln!("Failed to report exit cause: {}", error), |
| 133 | + /// } |
| 134 | + /// ``` |
| 135 | + pub fn send_exit_cause_for_post_compute_stage( |
| 136 | + &self, |
| 137 | + authorization: &str, |
| 138 | + chain_task_id: &str, |
| 139 | + exit_cause: &ExitMessage, |
| 140 | + ) -> Result<(), Error> { |
| 141 | + let url = format!("{}/compute/post/{}/exit", self.base_url, chain_task_id); |
| 142 | + let response = self |
| 143 | + .client |
| 144 | + .post(&url) |
| 145 | + .header(AUTHORIZATION, authorization) |
| 146 | + .json(exit_cause) |
| 147 | + .send()?; |
| 148 | + |
| 149 | + if response.status().is_success() { |
| 150 | + Ok(()) |
| 151 | + } else { |
| 152 | + Err(response.error_for_status().unwrap_err()) |
| 153 | + } |
| 154 | + } |
| 155 | +} |
| 156 | + |
| 157 | +#[cfg(test)] |
| 158 | +mod tests { |
| 159 | + use super::*; |
| 160 | + use crate::compute::utils::env_utils::TeeSessionEnvironmentVariable::*; |
| 161 | + use serde_json::{json, to_string}; |
| 162 | + use temp_env::with_vars; |
| 163 | + use wiremock::{ |
| 164 | + Mock, MockServer, ResponseTemplate, |
| 165 | + matchers::{body_json, header, method, path}, |
| 166 | + }; |
| 167 | + |
| 168 | + // region ExitMessage() |
| 169 | + #[test] |
| 170 | + fn should_serialize_exit_message() { |
| 171 | + let causes = [ |
| 172 | + ( |
| 173 | + ReplicateStatusCause::PostComputeInvalidTeeSignature, |
| 174 | + "PostComputeInvalidTeeSignature", |
| 175 | + ), |
| 176 | + ( |
| 177 | + ReplicateStatusCause::PostComputeWorkerAddressMissing, |
| 178 | + "PostComputeWorkerAddressMissing", |
| 179 | + ), |
| 180 | + ( |
| 181 | + ReplicateStatusCause::PostComputeFailedUnknownIssue, |
| 182 | + "PostComputeFailedUnknownIssue", |
| 183 | + ), |
| 184 | + ]; |
| 185 | + |
| 186 | + for (cause, message) in causes { |
| 187 | + let exit_message = ExitMessage::from(&cause); |
| 188 | + let serialized = to_string(&exit_message).expect("Failed to serialize"); |
| 189 | + let expected = format!("{{\"cause\":\"{message}\"}}"); |
| 190 | + assert_eq!(serialized, expected); |
| 191 | + } |
| 192 | + } |
| 193 | + // endregion |
| 194 | + |
| 195 | + // region get_worker_api_client |
| 196 | + #[test] |
| 197 | + fn should_get_worker_api_client_with_env_var() { |
| 198 | + with_vars( |
| 199 | + vec![(WORKER_HOST_ENV_VAR.name(), Some("custom-worker-host:9999"))], |
| 200 | + || { |
| 201 | + let client = WorkerApiClient::from_env(); |
| 202 | + assert_eq!(client.base_url, "http://custom-worker-host:9999"); |
| 203 | + }, |
| 204 | + ); |
| 205 | + } |
| 206 | + |
| 207 | + #[test] |
| 208 | + fn should_get_worker_api_client_without_env_var() { |
| 209 | + with_vars(vec![(WORKER_HOST_ENV_VAR.name(), None::<&str>)], || { |
| 210 | + let client = WorkerApiClient::from_env(); |
| 211 | + assert_eq!(client.base_url, format!("http://{}", DEFAULT_WORKER_HOST)); |
| 212 | + }); |
| 213 | + } |
| 214 | + // endregion |
| 215 | + |
| 216 | + // region send_exit_cause_for_post_compute_stage() |
| 217 | + const CHALLENGE: &str = "challenge"; |
| 218 | + const CHAIN_TASK_ID: &str = "0x123456789abcdef"; |
| 219 | + |
| 220 | + #[tokio::test] |
| 221 | + async fn should_send_exit_cause() { |
| 222 | + let mock_server = MockServer::start().await; |
| 223 | + let server_url = mock_server.uri(); |
| 224 | + |
| 225 | + let expected_body = json!({ |
| 226 | + "cause": ReplicateStatusCause::PostComputeInvalidTeeSignature, |
| 227 | + }); |
| 228 | + |
| 229 | + Mock::given(method("POST")) |
| 230 | + .and(path(format!("/compute/post/{}/exit", CHAIN_TASK_ID))) |
| 231 | + .and(header("Authorization", CHALLENGE)) |
| 232 | + .and(body_json(&expected_body)) |
| 233 | + .respond_with(ResponseTemplate::new(200)) |
| 234 | + .expect(1) |
| 235 | + .mount(&mock_server) |
| 236 | + .await; |
| 237 | + |
| 238 | + let result = tokio::task::spawn_blocking(move || { |
| 239 | + let exit_message = |
| 240 | + ExitMessage::from(&ReplicateStatusCause::PostComputeInvalidTeeSignature); |
| 241 | + let worker_api_client = WorkerApiClient::new(&server_url); |
| 242 | + worker_api_client.send_exit_cause_for_post_compute_stage( |
| 243 | + CHALLENGE, |
| 244 | + CHAIN_TASK_ID, |
| 245 | + &exit_message, |
| 246 | + ) |
| 247 | + }) |
| 248 | + .await |
| 249 | + .expect("Task panicked"); |
| 250 | + |
| 251 | + assert!(result.is_ok()); |
| 252 | + } |
| 253 | + |
| 254 | + #[tokio::test] |
| 255 | + async fn should_not_send_exit_cause() { |
| 256 | + let mock_server = MockServer::start().await; |
| 257 | + let server_url = mock_server.uri(); |
| 258 | + |
| 259 | + Mock::given(method("POST")) |
| 260 | + .and(path(format!("/compute/post/{}/exit", CHAIN_TASK_ID))) |
| 261 | + .respond_with(ResponseTemplate::new(404)) |
| 262 | + .expect(1) |
| 263 | + .mount(&mock_server) |
| 264 | + .await; |
| 265 | + |
| 266 | + let result = tokio::task::spawn_blocking(move || { |
| 267 | + let exit_message = |
| 268 | + ExitMessage::from(&ReplicateStatusCause::PostComputeFailedUnknownIssue); |
| 269 | + let worker_api_client = WorkerApiClient::new(&server_url); |
| 270 | + worker_api_client.send_exit_cause_for_post_compute_stage( |
| 271 | + CHALLENGE, |
| 272 | + CHAIN_TASK_ID, |
| 273 | + &exit_message, |
| 274 | + ) |
| 275 | + }) |
| 276 | + .await |
| 277 | + .expect("Task panicked"); |
| 278 | + |
| 279 | + assert!(result.is_err()); |
| 280 | + |
| 281 | + if let Err(error) = result { |
| 282 | + assert_eq!(error.status().unwrap(), 404); |
| 283 | + } |
| 284 | + } |
| 285 | + // endregion |
| 286 | +} |
0 commit comments