1use std::{
4 io::Read,
5 net::TcpListener,
6 path::{Path, PathBuf},
7 process::{Child, Command, Stdio},
8 time::{Duration, Instant},
9};
10
11use anyhow::{bail, Context, Result};
12
13use crate::scripts_lane_a::AuthSmokeResults;
14
15const STARTUP_TIMEOUT: Duration = Duration::from_secs(15);
16const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
17const TRACEPARENT: &str = "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01";
18
19pub fn test_trace_headers(_args: &[String]) -> Result<()> {
20 let binary = build_soma_binary()?;
21 let mut results = AuthSmokeResults::default();
22
23 run_scenario(&binary, &mut results, "off", "off", |port| {
24 off_mode_checks(port)
25 })?;
26 run_scenario(&binary, &mut results, "trusted", "trusted", |port| {
27 trusted_mode_checks(port, false)
28 })?;
29 run_scenario(
30 &binary,
31 &mut results,
32 "trusted-with-baggage",
33 "trusted-with-baggage",
34 |port| trusted_mode_checks(port, true),
35 )?;
36
37 println!("\n{} passed, {} failed", results.pass, results.fail);
38 if results.fail == 0 {
39 Ok(())
40 } else {
41 bail!("{} trace-header smoke check(s) failed", results.fail)
42 }
43}
44
45fn build_soma_binary() -> Result<PathBuf> {
46 println!("==> Building soma binary once...");
47 let output = Command::new("cargo")
48 .args([
49 "build",
50 "--bin",
51 "soma",
52 "--features",
53 "full",
54 "--message-format=json",
55 ])
56 .output()
57 .context("run cargo build")?;
58 if !output.status.success() {
59 bail!(
60 "cargo build failed:\n{}",
61 String::from_utf8_lossy(&output.stderr)
62 );
63 }
64 for line in String::from_utf8_lossy(&output.stdout).lines() {
65 let Ok(message) = serde_json::from_str::<serde_json::Value>(line) else {
66 continue;
67 };
68 if message["reason"] == "compiler-artifact"
69 && message["target"]["name"] == "soma"
70 && message["executable"].is_string()
71 {
72 return Ok(PathBuf::from(
73 message["executable"]
74 .as_str()
75 .context("executable path should be a string")?,
76 ));
77 }
78 }
79 bail!("cargo build did not report a soma binary artifact")
80}
81
82fn run_scenario(
83 binary: &Path,
84 results: &mut AuthSmokeResults,
85 label: &str,
86 trace_headers_env: &str,
87 checks: impl FnOnce(u16) -> Result<Vec<(String, bool)>>,
88) -> Result<()> {
89 println!("==> Scenario: SOMA_MCP_TRACE_HEADERS={trace_headers_env}");
90 let home = tempfile::tempdir().context("create isolated SOMA_HOME")?;
91 let port = free_port()?;
92 let mut server = ServerGuard::spawn(binary, home.path(), port, trace_headers_env)?;
93 server.wait_for_health(port)?;
94
95 for (check_label, passed) in checks(port)? {
96 record(results, label, &check_label, passed);
97 }
98 let log = server.captured_log();
99 assert_safe_log_fields(&log, results, label, trace_headers_env);
100 Ok(())
101}
102
103fn record(results: &mut AuthSmokeResults, scenario: &str, label: &str, passed: bool) {
104 let label = format!("{scenario}: {label}");
105 if passed {
106 results.pass(&label);
107 } else {
108 results.fail(&label);
109 }
110}
111
112fn off_mode_checks(port: u16) -> Result<Vec<(String, bool)>> {
113 let status = curl_status(port, &[("traceparent", TRACEPARENT)])?;
114 let preflight = curl_preflight(port, "TraceParent")?;
115 Ok(vec![
116 (
117 "status tool call succeeds with traceparent present".to_owned(),
118 status == 200,
119 ),
120 (
121 "CORS preflight denies traceparent".to_owned(),
122 !preflight.to_ascii_lowercase().contains("traceparent"),
123 ),
124 ])
125}
126
127fn trusted_mode_checks(port: u16, baggage_enabled: bool) -> Result<Vec<(String, bool)>> {
128 let mut checks = Vec::new();
129 let status = curl_status(
130 port,
131 &[
132 ("traceparent", TRACEPARENT),
133 ("tracestate", "vendor=value"),
134 ("baggage", "region=us-east-1"),
135 ],
136 )?;
137 checks.push((
138 "valid trace context tool call succeeds".to_owned(),
139 status == 200,
140 ));
141
142 let duplicate = curl_status(
143 port,
144 &[
145 ("traceparent", TRACEPARENT),
146 (
147 "traceparent",
148 "00-11112222333344445555666677778888-1111222233334444-01",
149 ),
150 ],
151 )?;
152 checks.push((
153 "duplicate traceparent does not crash the server".to_owned(),
154 duplicate == 200,
155 ));
156
157 let non_ascii = curl_status(port, &[("traceparent", "00-\u{00e9}-invalid-01")])?;
158 checks.push((
159 "non-ASCII traceparent does not crash the server".to_owned(),
160 non_ascii == 200,
161 ));
162
163 let preflight = curl_preflight(port, "TraceParent, TraceState, Baggage")?;
164 let preflight = preflight.to_ascii_lowercase();
165 checks.push((
166 "CORS preflight allows traceparent".to_owned(),
167 preflight.contains("traceparent"),
168 ));
169 checks.push((
170 "CORS preflight allows tracestate".to_owned(),
171 preflight.contains("tracestate"),
172 ));
173 checks.push((
174 "baggage CORS allowance matches mode".to_owned(),
175 preflight.contains("baggage") == baggage_enabled,
176 ));
177 Ok(checks)
178}
179
180fn assert_safe_log_fields(log: &str, results: &mut AuthSmokeResults, scenario: &str, mode: &str) {
181 match mode {
182 "off" => record(
183 results,
184 scenario,
185 "logs show HTTP trace extraction disabled",
186 log.contains("http_trace_headers_present=false"),
187 ),
188 "trusted" => {
189 record(
190 results,
191 scenario,
192 "logs contain the safe trace ID prefix",
193 log.contains("trace_id_prefix=Some(\"0af76519\")"),
194 );
195 record(
196 results,
197 scenario,
198 "logs prove baggage was stripped",
199 log.contains("baggage_member_count=0"),
200 );
201 }
202 "trusted-with-baggage" => {
203 record(
204 results,
205 scenario,
206 "logs contain the safe trace ID prefix",
207 log.contains("trace_id_prefix=Some(\"0af76519\")"),
208 );
209 record(
210 results,
211 scenario,
212 "logs contain only the safe baggage count",
213 log.contains("baggage_member_count=1"),
214 );
215 }
216 _ => unreachable!("known trace-header mode"),
217 }
218 record(
219 results,
220 scenario,
221 "raw baggage value is absent from logs",
222 !log.contains("region=us-east-1") && !log.contains("accessToken"),
223 );
224}
225
226fn curl_status(port: u16, headers: &[(&str, &str)]) -> Result<u16> {
227 let timeout = REQUEST_TIMEOUT.as_secs().to_string();
228 let mut cmd = Command::new("curl");
229 cmd.args([
230 "-s",
231 "-o",
232 "/dev/null",
233 "-w",
234 "%{http_code}",
235 "-X",
236 "POST",
237 "-H",
238 "Content-Type: application/json",
239 "-H",
240 "Accept: application/json, text/event-stream",
241 "--max-time",
242 &timeout,
243 ]);
244 for (name, value) in headers {
245 cmd.args(["-H", &format!("{name}: {value}")]);
246 }
247 cmd.args([
248 "-d",
249 r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"soma","arguments":{"action":"status"}}}"#,
250 &format!("http://127.0.0.1:{port}/mcp"),
251 ]);
252 let output = cmd.output().context("run curl tool call")?;
253 String::from_utf8_lossy(&output.stdout)
254 .trim()
255 .parse::<u16>()
256 .context("parse curl status code")
257}
258
259fn curl_preflight(port: u16, requested_headers: &str) -> Result<String> {
260 let output = curl_preflight_command(port, requested_headers)
261 .output()
262 .context("run curl preflight")?;
263 checked_preflight_output(output)
264}
265
266fn checked_preflight_output(output: std::process::Output) -> Result<String> {
267 if !output.status.success() {
268 bail!(
269 "curl preflight failed with {}: {}",
270 output.status,
271 String::from_utf8_lossy(&output.stderr).trim()
272 );
273 }
274 Ok(String::from_utf8_lossy(&output.stdout).into_owned())
275}
276
277fn curl_preflight_command(port: u16, requested_headers: &str) -> Command {
278 let timeout = REQUEST_TIMEOUT.as_secs().to_string();
279 let mut command = Command::new("curl");
280 command.args([
281 "-s",
282 "-i",
283 "--max-time",
284 &timeout,
285 "-X",
286 "OPTIONS",
287 "-H",
288 &format!("Origin: http://127.0.0.1:{port}"),
289 "-H",
290 "Access-Control-Request-Method: POST",
291 "-H",
292 &format!("Access-Control-Request-Headers: {requested_headers}"),
293 &format!("http://127.0.0.1:{port}/mcp"),
294 ]);
295 command
296}
297
298fn free_port() -> Result<u16> {
299 let listener = TcpListener::bind("127.0.0.1:0").context("bind ephemeral port")?;
300 Ok(listener.local_addr()?.port())
301}
302
303struct ServerGuard {
304 child: Child,
305 log_path: PathBuf,
306}
307
308impl ServerGuard {
309 fn spawn(binary: &Path, home: &Path, port: u16, trace_headers: &str) -> Result<Self> {
310 let log_path = home.join("server.log");
311 let log_file = std::fs::File::create(&log_path).context("create server log")?;
312 let mut command = Command::new(binary);
313 for (key, _) in
314 std::env::vars_os().filter(|(key, _)| key.to_string_lossy().starts_with("SOMA_"))
315 {
316 command.env_remove(key);
317 }
318 let child = command
319 .arg("serve")
320 .current_dir(home)
321 .env("HOME", home)
322 .env("SOMA_HOME", home)
323 .env("SOMA_MCP_HOST", "127.0.0.1")
324 .env("SOMA_MCP_PORT", port.to_string())
325 .env("SOMA_MCP_NO_AUTH", "true")
326 .env("SOMA_MCP_TRACE_HEADERS", trace_headers)
327 .env("SOMA_RUNTIME_MODE", "local")
328 .env("RUST_LOG", "soma_mcp=info,soma=info")
329 .stdout(Stdio::from(
330 log_file.try_clone().context("clone log handle")?,
331 ))
332 .stderr(Stdio::from(log_file))
333 .spawn()
334 .context("spawn soma serve")?;
335 Ok(Self { child, log_path })
336 }
337
338 fn wait_for_health(&mut self, port: u16) -> Result<()> {
339 let deadline = Instant::now() + STARTUP_TIMEOUT;
340 loop {
341 if let Some(status) = self.child.try_wait().context("poll soma process")? {
342 bail!("soma exited before becoming healthy: {status}");
343 }
344 if Instant::now() > deadline {
345 bail!("server did not become healthy within {STARTUP_TIMEOUT:?}");
346 }
347 let output = Command::new("curl")
348 .args([
349 "-s",
350 "-o",
351 "/dev/null",
352 "-w",
353 "%{http_code}",
354 "--max-time",
355 "1",
356 &format!("http://127.0.0.1:{port}/health"),
357 ])
358 .output();
359 if output.is_ok_and(|output| String::from_utf8_lossy(&output.stdout).trim() == "200") {
360 return Ok(());
361 }
362 std::thread::sleep(Duration::from_millis(200));
363 }
364 }
365
366 fn captured_log(&self) -> String {
367 let mut contents = String::new();
368 if let Ok(mut file) = std::fs::File::open(&self.log_path) {
369 let _ = file.read_to_string(&mut contents);
370 }
371 strip_ansi(&contents)
372 }
373}
374
375fn strip_ansi(input: &str) -> String {
392 let mut output = String::with_capacity(input.len());
393 let mut chars = input.chars().peekable();
394 while let Some(ch) = chars.next() {
395 if ch == '\u{1b}' && chars.peek() == Some(&'[') {
396 chars.next(); for c in chars.by_ref() {
398 if c.is_ascii_alphabetic() {
399 break;
400 }
401 }
402 continue;
403 }
404 output.push(ch);
405 }
406 output
407}
408
409#[cfg(test)]
410mod strip_ansi_tests {
411 use super::strip_ansi;
412
413 #[test]
414 fn removes_sgr_escape_sequences_between_field_name_and_value() {
415 let colored = "\u{1b}[3mtrace_id_prefix\u{1b}[0m\u{1b}[2m=\u{1b}[0mSome(\"0af76519\")";
416 assert_eq!(strip_ansi(colored), "trace_id_prefix=Some(\"0af76519\")");
417 }
418
419 #[test]
420 fn leaves_plain_text_untouched() {
421 assert_eq!(
422 strip_ansi("http_trace_headers_present=false"),
423 "http_trace_headers_present=false"
424 );
425 }
426}
427
428impl Drop for ServerGuard {
429 fn drop(&mut self) {
430 let _ = self.child.kill();
431 let _ = self.child.wait();
432 }
433}
434
435#[cfg(test)]
436mod tests {
437 use super::*;
438
439 #[test]
440 fn preflight_curl_uses_the_request_timeout() {
441 let command = curl_preflight_command(40060, "TraceParent");
442 let args: Vec<_> = command
443 .get_args()
444 .map(|arg| arg.to_string_lossy())
445 .collect();
446 let timeout_index = args
447 .iter()
448 .position(|arg| arg == "--max-time")
449 .expect("preflight curl should have --max-time");
450 assert_eq!(
451 args.get(timeout_index + 1).map(|arg| arg.as_ref()),
452 Some("5")
453 );
454 }
455
456 #[cfg(unix)]
457 #[test]
458 fn preflight_rejects_a_nonzero_curl_exit_status() {
459 use std::os::unix::process::ExitStatusExt;
460
461 let output = std::process::Output {
462 status: std::process::ExitStatus::from_raw(7 << 8),
463 stdout: Vec::new(),
464 stderr: b"timed out".to_vec(),
465 };
466
467 let error = checked_preflight_output(output).expect_err("nonzero curl must fail smoke");
468 assert!(error.to_string().contains("curl preflight failed"));
469 assert!(error.to_string().contains("timed out"));
470 }
471}