1use std::time::{SystemTime, UNIX_EPOCH};
2
3use cc_lb_plugin_wire::limits::{
4 IMPLEMENTED_FUNCTIONS_MAX, SELF_CHECK_FUEL, SELF_CHECK_OUTPUT_MAX_BYTES, SELF_CHECK_WALL_MS,
5};
6use cc_lb_plugin_wire::self_check::{
7 SelfCheckError, SelfCheckRequest, SelfCheckResponse, SelfCheckStatus,
8};
9use cc_lb_plugin_wire::wire_function::all_wire_functions;
10use thiserror::Error;
11
12use crate::handshake::{BuildPluginError, build_plugin};
13
14const SELF_CHECK_EXPORT: &str = "cc_lb_self_check";
15
16pub fn execute_self_check(
17 plugin_bytes: &[u8],
18) -> Result<SelfCheckResponse, SelfCheckExecutionError> {
19 let mut plugin =
20 build_plugin(plugin_bytes, SELF_CHECK_WALL_MS, SELF_CHECK_FUEL).map_err(|source| {
21 match source {
22 BuildPluginError::Instantiate { reason } => {
23 SelfCheckExecutionError::Instantiate { reason }
24 }
25 }
26 })?;
27
28 if !plugin.function_exists(SELF_CHECK_EXPORT) {
29 return Err(SelfCheckExecutionError::MissingSelfCheckExport);
30 }
31
32 let request = build_request()?;
33 let request = serde_json::to_string(&request).map_err(|source| {
34 SelfCheckExecutionError::SerializeRequest {
35 reason: source.to_string(),
36 }
37 })?;
38 let response = plugin
39 .call::<&str, String>(SELF_CHECK_EXPORT, request.as_str())
40 .map_err(|source| SelfCheckExecutionError::Call {
41 reason: source.to_string(),
42 })?;
43
44 if response.len() > SELF_CHECK_OUTPUT_MAX_BYTES {
45 return Err(SelfCheckExecutionError::OutputTooLarge {
46 bytes: response.len(),
47 max: SELF_CHECK_OUTPUT_MAX_BYTES,
48 });
49 }
50
51 let response: SelfCheckResponse = serde_json::from_str(&response).map_err(|source| {
52 SelfCheckExecutionError::DecodeResponse {
53 reason: source.to_string(),
54 }
55 })?;
56 response.validate()?;
57 validate_status_failures(&response)?;
58
59 Ok(response)
60}
61
62fn build_request() -> Result<SelfCheckRequest, SelfCheckExecutionError> {
63 let functions = all_wire_functions();
64 if functions.len() > IMPLEMENTED_FUNCTIONS_MAX {
65 return Err(SelfCheckExecutionError::Validation(
66 SelfCheckError::TooManyFunctions,
67 ));
68 }
69
70 let initiated_at = SystemTime::now()
71 .duration_since(UNIX_EPOCH)
72 .map_err(|source| SelfCheckExecutionError::Clock {
73 reason: source.to_string(),
74 })?
75 .as_secs();
76 let initiated_at =
77 i64::try_from(initiated_at).map_err(|source| SelfCheckExecutionError::Clock {
78 reason: source.to_string(),
79 })?;
80 let request = SelfCheckRequest {
81 functions_to_test: functions
82 .iter()
83 .map(|function| (*function).to_owned())
84 .collect(),
85 initiated_at,
86 };
87 request.validate()?;
88 Ok(request)
89}
90
91fn validate_status_failures(response: &SelfCheckResponse) -> Result<(), SelfCheckExecutionError> {
92 match response.status {
93 SelfCheckStatus::Success if !response.failures.is_empty() => {
94 Err(SelfCheckExecutionError::SuccessWithFailures {
95 count: response.failures.len(),
96 })
97 }
98 SelfCheckStatus::Failure if response.failures.is_empty() => {
99 Err(SelfCheckExecutionError::FailureWithoutFailures)
100 }
101 SelfCheckStatus::Failure => Err(SelfCheckExecutionError::FailureStatus {
102 failures: response.failures.len(),
103 }),
104 SelfCheckStatus::Success => Ok(()),
105 }
106}
107
108#[non_exhaustive]
109#[derive(Debug, Error)]
110pub enum SelfCheckExecutionError {
111 #[error("self-check validation failed: {0}")]
112 Validation(#[from] SelfCheckError),
113 #[error("self-check plugin instantiation failed: {reason}")]
114 Instantiate { reason: String },
115 #[error("plugin does not export cc_lb_self_check")]
116 MissingSelfCheckExport,
117 #[error("self-check request serialization failed: {reason}")]
118 SerializeRequest { reason: String },
119 #[error("self-check call failed: {reason}")]
120 Call { reason: String },
121 #[error("self-check output size {bytes} exceeds maximum {max}")]
122 OutputTooLarge { bytes: usize, max: usize },
123 #[error("self-check response decode failed: {reason}")]
124 DecodeResponse { reason: String },
125 #[error("self-check success response included {count} failure(s)")]
126 SuccessWithFailures { count: usize },
127 #[error("self-check failure response did not include failures")]
128 FailureWithoutFailures,
129 #[error("self-check reported failure status with {failures} failure(s)")]
130 FailureStatus { failures: usize },
131 #[error("self-check timestamp generation failed: {reason}")]
132 Clock { reason: String },
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138
139 #[test]
140 fn execute_self_check_accepts_success_response() {
141 let wasm = self_check_module(
142 r#"{"status":"success","failures":[],"completed_at":1}"#,
143 false,
144 );
145
146 let response = execute_self_check(&wasm).expect("self-check succeeds");
147
148 assert_eq!(response.status, SelfCheckStatus::Success);
149 assert!(response.failures.is_empty());
150 }
151
152 #[test]
153 fn execute_self_check_rejects_failure_status() {
154 let wasm = self_check_module(
155 r#"{"status":"failure","failures":[{"stage":"wire_function_test","message":"bad wire shape"}],"completed_at":1}"#,
156 false,
157 );
158
159 let err = execute_self_check(&wasm).expect_err("failure status rejected at executor level");
160
161 match err {
162 SelfCheckExecutionError::FailureStatus { failures } => assert_eq!(failures, 1),
163 other => panic!("expected failure status rejection, got {other:?}"),
164 }
165 }
166
167 #[test]
168 fn execute_self_check_rejects_missing_export() {
169 let wasm = wat::parse_str(r#"(module (func (export "shape") (result i32) (i32.const 0)))"#)
170 .expect("wat parses");
171
172 let err = execute_self_check(&wasm).expect_err("missing export rejected");
173
174 match err {
175 SelfCheckExecutionError::MissingSelfCheckExport => {}
176 other => panic!("expected missing export, got {other:?}"),
177 }
178 }
179
180 #[test]
181 fn execute_self_check_rejects_user_host_imports() {
182 let wasm = self_check_module(
183 r#"{"status":"success","failures":[],"completed_at":1}"#,
184 true,
185 );
186
187 let err = execute_self_check(&wasm).expect_err("host import rejected");
188
189 match err {
190 SelfCheckExecutionError::Instantiate { .. } | SelfCheckExecutionError::Call { .. } => {}
191 other => panic!("expected purity failure, got {other:?}"),
192 }
193 }
194
195 #[test]
196 fn execute_self_check_rejects_oversized_output() {
197 let output = "x".repeat(SELF_CHECK_OUTPUT_MAX_BYTES + 1);
198 let wasm = self_check_module(&output, false);
199
200 let err = execute_self_check(&wasm).expect_err("oversized response rejected");
201
202 match err {
203 SelfCheckExecutionError::OutputTooLarge { bytes, max } => {
204 assert_eq!(bytes, SELF_CHECK_OUTPUT_MAX_BYTES + 1);
205 assert_eq!(max, SELF_CHECK_OUTPUT_MAX_BYTES);
206 }
207 other => panic!("expected oversized output, got {other:?}"),
208 }
209 }
210
211 #[test]
212 fn execute_self_check_rejects_success_with_failures() {
213 let wasm = self_check_module(
214 r#"{"status":"success","failures":[{"stage":"wire_function_test","message":"bad"}],"completed_at":1}"#,
215 false,
216 );
217
218 let err = execute_self_check(&wasm).expect_err("status/failures mismatch rejected");
219
220 match err {
221 SelfCheckExecutionError::SuccessWithFailures { count } => assert_eq!(count, 1),
222 other => panic!("expected status/failures mismatch, got {other:?}"),
223 }
224 }
225
226 #[test]
227 fn execute_self_check_rejects_failure_without_failures() {
228 let wasm = self_check_module(
229 r#"{"status":"failure","failures":[],"completed_at":1}"#,
230 false,
231 );
232
233 let err = execute_self_check(&wasm).expect_err("status/failures mismatch rejected");
234
235 match err {
236 SelfCheckExecutionError::FailureWithoutFailures => {}
237 other => panic!("expected status/failures mismatch, got {other:?}"),
238 }
239 }
240
241 #[test]
242 fn execute_self_check_rejects_response_validation_errors() {
243 let wasm = self_check_module(
244 r#"{"status":"success","failures":[],"completed_at":0}"#,
245 false,
246 );
247
248 let err = execute_self_check(&wasm).expect_err("invalid response rejected");
249
250 match err {
251 SelfCheckExecutionError::Validation(SelfCheckError::InvalidTimestamp(_)) => {}
252 other => panic!("expected response validation error, got {other:?}"),
253 }
254 }
255
256 fn self_check_module(output: &str, import_user_host: bool) -> Vec<u8> {
257 let output_helper = bytes_helper("self_check_out", output.as_bytes());
258 let user_import = if import_user_host {
259 r#"(import "extism:host/user" "cc_lb_log" (func $cc_lb_log (param i64 i64)))"#
260 } else {
261 ""
262 };
263 let user_call = if import_user_host {
264 " (call $cc_lb_log (call $self_check_out) (call $self_check_out))"
265 } else {
266 ""
267 };
268
269 let wat = format!(
270 r#"
271(module
272 (import "extism:host/env" "alloc" (func $alloc (param i64) (result i64)))
273 (import "extism:host/env" "store_u8" (func $store_u8 (param i64 i32)))
274 (import "extism:host/env" "output_set" (func $output_set (param i64 i64)))
275 {user_import}
276 {output_helper}
277 (func (export "cc_lb_self_check") (result i32)
278{user_call}
279 (call $output_set (call $self_check_out) (i64.const {len}))
280 (i32.const 0))
281)
282"#,
283 len = output.len()
284 );
285 wat::parse_str(&wat).expect("self-check wat parses")
286 }
287
288 fn bytes_helper(name: &str, bytes: &[u8]) -> String {
289 let mut stores = String::new();
290 for (index, byte) in bytes.iter().enumerate() {
291 stores.push_str(&format!(
292 " (call $store_u8 (i64.add (local.get $ptr) (i64.const {index})) (i32.const {byte}))\n"
293 ));
294 }
295 format!(
296 r#"
297(func ${name} (result i64)
298 (local $ptr i64)
299 (local.set $ptr (call $alloc (i64.const {len})))
300{stores} (local.get $ptr))
301"#,
302 len = bytes.len()
303 )
304 }
305}