1use std::sync::{
43 Arc,
44 atomic::{AtomicU32, AtomicUsize, Ordering},
45};
46use std::time::{Duration, SystemTime, UNIX_EPOCH};
47
48use serde::Serialize;
49use tokio::io::AsyncWriteExt;
50use tokio::sync::broadcast;
51
52pub const TRACE_CHANNEL_CAPACITY: usize = 1_024;
54pub const MAX_SUBSCRIBERS: usize = 4;
56
57#[derive(Clone, Debug, Serialize)]
61pub struct MatchTraceEvent {
62 pub event_id: u64,
64 pub schema_version: u8,
66 pub received_at_ms: u64,
68 pub duration_ms: u32,
70 pub request: RequestSummary,
72 pub outcome: Outcome,
74 pub dropped_count: u32,
76}
77
78#[derive(Clone, Debug, Serialize)]
80pub struct RequestSummary {
81 pub method: String,
82 pub url_path: String,
83 pub headers: Vec<(String, String)>,
85 #[serde(skip_serializing_if = "Option::is_none")]
88 pub body_json: Option<serde_json::Value>,
89 #[serde(skip_serializing_if = "std::ops::Not::not")]
91 pub body_truncated: bool,
92}
93
94impl RequestSummary {
95 pub fn without_body(method: String, url_path: String, headers: Vec<(String, String)>) -> Self {
97 Self {
98 method,
99 url_path,
100 headers,
101 body_json: None,
102 body_truncated: false,
103 }
104 }
105}
106
107#[derive(Clone, Debug)]
109pub struct TraceConfig {
110 pub capture_body: bool,
112 pub max_body_bytes: usize,
115}
116
117impl Default for TraceConfig {
118 fn default() -> Self {
119 Self {
120 capture_body: false,
121 max_body_bytes: 8_192,
122 }
123 }
124}
125
126#[derive(Clone, Debug, Serialize)]
128#[serde(tag = "type", rename_all = "snake_case")]
129pub enum Outcome {
130 Matched {
131 rule_set_index: usize,
132 rule_index: usize,
133 },
134 Fallback {
135 file_path: String,
136 status: u16,
137 },
138 Miss {
139 status: u16,
140 },
141 Error {
142 kind: String,
143 message: String,
144 },
145}
146
147#[derive(Clone)]
153pub struct TraceEmitter {
154 sender: broadcast::Sender<MatchTraceEvent>,
155 event_counter: Arc<AtomicU32>,
156 dropped_counter: Arc<AtomicU32>,
157 pub config: Arc<TraceConfig>,
159}
160
161impl TraceEmitter {
162 pub fn new() -> Self {
163 Self::with_config(TraceConfig::default())
164 }
165
166 pub fn with_config(config: TraceConfig) -> Self {
167 let (sender, _) = broadcast::channel(TRACE_CHANNEL_CAPACITY);
168 Self {
169 sender,
170 event_counter: Arc::new(AtomicU32::new(0)),
171 dropped_counter: Arc::new(AtomicU32::new(0)),
172 config: Arc::new(config),
173 }
174 }
175
176 pub fn subscribe(&self) -> broadcast::Receiver<MatchTraceEvent> {
178 self.sender.subscribe()
179 }
180
181 pub fn enrich_with_body(
184 &self,
185 summary: &mut RequestSummary,
186 body_json: Option<&serde_json::Value>,
187 ) {
188 if !self.config.capture_body {
189 return;
190 }
191 match body_json {
192 None => {} Some(v) => {
194 match serde_json::to_string(v) {
196 Ok(s) if s.len() <= self.config.max_body_bytes => {
197 summary.body_json = Some(v.clone());
198 }
199 Ok(_) => {
200 summary.body_truncated = true;
201 }
202 Err(_) => {} }
204 }
205 }
206 }
207
208 pub fn emit(
211 &self,
212 received_at_ms: u64,
213 duration_ms: u32,
214 request: RequestSummary,
215 outcome: Outcome,
216 ) {
217 let event_id = self.event_counter.fetch_add(1, Ordering::Relaxed) as u64;
218 let dropped_count = self.dropped_counter.swap(0, Ordering::Relaxed);
219
220 let event = MatchTraceEvent {
221 event_id,
222 schema_version: 1,
223 received_at_ms,
224 duration_ms,
225 request,
226 outcome,
227 dropped_count,
228 };
229
230 if self.sender.send(event).is_err() {
231 self.dropped_counter.fetch_add(1, Ordering::Relaxed);
232 }
233 }
234
235 pub fn has_subscribers(&self) -> bool {
237 self.sender.receiver_count() > 0
238 }
239}
240
241impl Default for TraceEmitter {
242 fn default() -> Self {
243 Self::new()
244 }
245}
246
247#[derive(Clone, Debug, Default)]
251pub enum TraceTransportConfig {
252 #[cfg(unix)]
254 Uds { path: String },
255 Tcp { addr: String },
257 #[default]
259 Disabled,
260}
261
262pub struct TraceTransport;
265
266impl TraceTransport {
267 pub async fn accept_loop(config: TraceTransportConfig, emitter: TraceEmitter) {
279 match config {
280 #[cfg(unix)]
281 TraceTransportConfig::Uds { path } => Self::uds_accept_loop(path, emitter).await,
282 TraceTransportConfig::Tcp { addr } => Self::tcp_accept_loop(addr, emitter).await,
283 TraceTransportConfig::Disabled => {
284 }
286 }
287 }
288
289 async fn tcp_accept_loop(addr: String, emitter: TraceEmitter) {
292 let listener = match tokio::net::TcpListener::bind(&addr).await {
293 Ok(l) => {
294 let bound = l
295 .local_addr()
296 .map(|a| a.to_string())
297 .unwrap_or_else(|_| addr.clone());
298 log::info!("trace transport: TCP listening on {}", bound);
299 l
300 }
301 Err(e) => {
302 log::error!("trace transport: failed to bind TCP {}: {}", addr, e);
303 return;
304 }
305 };
306
307 let active = Arc::new(AtomicUsize::new(0));
308 loop {
309 match listener.accept().await {
310 Ok((stream, peer)) => {
311 let count = active.fetch_add(1, Ordering::Relaxed) + 1;
312 if count > MAX_SUBSCRIBERS {
313 active.fetch_sub(1, Ordering::Relaxed);
314 let active_clone = active.clone();
315 tokio::spawn(async move {
316 let (_, mut writer) = tokio::io::split(stream);
317 let _ = writer
318 .write_all(b"{\"error\":\"max_subscribers_reached\"}\n")
319 .await;
320 drop(active_clone);
321 });
322 continue;
323 }
324 log::debug!("trace: TCP subscriber connected from {}", peer);
325 let rx = emitter.subscribe();
326 let active_clone = active.clone();
327 tokio::spawn(async move {
328 let (_, writer) = tokio::io::split(stream);
329 Self::forward_events(writer, rx).await;
330 active_clone.fetch_sub(1, Ordering::Relaxed);
331 log::debug!("trace: TCP subscriber {} disconnected", peer);
332 });
333 }
334 Err(e) => {
335 log::error!("trace: TCP accept error: {}", e);
336 tokio::time::sleep(Duration::from_millis(100)).await;
337 }
338 }
339 }
340 }
341
342 #[cfg(unix)]
345 async fn uds_accept_loop(path: String, emitter: TraceEmitter) {
346 let _ = std::fs::remove_file(&path);
348
349 let listener = match tokio::net::UnixListener::bind(&path) {
350 Ok(l) => {
351 log::info!("trace transport: UDS listening at {}", path);
352 l
353 }
354 Err(e) => {
355 log::error!("trace transport: failed to bind UDS {}: {}", path, e);
356 return;
357 }
358 };
359
360 let active = Arc::new(AtomicUsize::new(0));
361 loop {
362 match listener.accept().await {
363 Ok((stream, _)) => {
364 let count = active.fetch_add(1, Ordering::Relaxed) + 1;
365 if count > MAX_SUBSCRIBERS {
366 active.fetch_sub(1, Ordering::Relaxed);
367 tokio::spawn(async move {
368 let (_, mut writer) = tokio::io::split(stream);
369 let _ = writer
370 .write_all(b"{\"error\":\"max_subscribers_reached\"}\n")
371 .await;
372 });
373 continue;
374 }
375 log::debug!("trace: UDS subscriber connected");
376 let rx = emitter.subscribe();
377 let active_clone = active.clone();
378 tokio::spawn(async move {
379 let (_, writer) = tokio::io::split(stream);
380 Self::forward_events(writer, rx).await;
381 active_clone.fetch_sub(1, Ordering::Relaxed);
382 log::debug!("trace: UDS subscriber disconnected");
383 });
384 }
385 Err(e) => {
386 log::error!("trace: UDS accept error: {}", e);
387 tokio::time::sleep(Duration::from_millis(100)).await;
388 }
389 }
390 }
391 }
392
393 async fn forward_events<W>(mut writer: W, mut rx: broadcast::Receiver<MatchTraceEvent>)
398 where
399 W: tokio::io::AsyncWrite + Unpin,
400 {
401 loop {
402 let event = match rx.recv().await {
403 Ok(e) => e,
404 Err(broadcast::error::RecvError::Lagged(n)) => {
405 log::debug!("trace: subscriber lagged, {} events dropped", n);
409 continue;
410 }
411 Err(broadcast::error::RecvError::Closed) => break,
412 };
413
414 let mut line = match serde_json::to_string(&event) {
415 Ok(s) => s,
416 Err(e) => {
417 log::error!("trace: serialise error: {}", e);
418 continue;
419 }
420 };
421 line.push('\n');
422
423 if writer.write_all(line.as_bytes()).await.is_err() {
424 break; }
426 }
427 }
428}
429
430pub fn now_ms() -> u64 {
434 SystemTime::now()
435 .duration_since(UNIX_EPOCH)
436 .unwrap_or(Duration::ZERO)
437 .as_millis() as u64
438}
439
440#[cfg(test)]
443mod tests {
444 use super::*;
445
446 #[tokio::test]
447 async fn emit_received_by_subscriber() {
448 let emitter = TraceEmitter::new();
449 let mut rx = emitter.subscribe();
450
451 emitter.emit(
452 1_000_000,
453 5,
454 RequestSummary {
455 method: "GET".into(),
456 url_path: "/api/test".into(),
457 headers: vec![],
458 body_json: None,
459 body_truncated: false,
460 },
461 Outcome::Miss { status: 404 },
462 );
463
464 let event = rx.try_recv().expect("event in channel");
465 assert_eq!(event.event_id, 0);
466 assert_eq!(event.schema_version, 1);
467 assert_eq!(event.request.method, "GET");
468 assert_eq!(event.duration_ms, 5);
469 assert_eq!(event.dropped_count, 0);
470 assert!(matches!(event.outcome, Outcome::Miss { status: 404 }));
471 }
472
473 #[tokio::test]
474 async fn emit_no_subscriber_increments_dropped() {
475 let emitter = TraceEmitter::new();
476 emitter.emit(
477 0,
478 0,
479 RequestSummary {
480 method: "GET".into(),
481 url_path: "/".into(),
482 headers: vec![],
483 body_json: None,
484 body_truncated: false,
485 },
486 Outcome::Miss { status: 404 },
487 );
488 let mut rx = emitter.subscribe();
489 emitter.emit(
490 0,
491 0,
492 RequestSummary {
493 method: "GET".into(),
494 url_path: "/".into(),
495 headers: vec![],
496 body_json: None,
497 body_truncated: false,
498 },
499 Outcome::Miss { status: 200 },
500 );
501 let event = rx.try_recv().expect("second event visible");
502 assert_eq!(
503 event.dropped_count, 1,
504 "first event should be counted dropped"
505 );
506 }
507
508 #[test]
509 fn has_subscribers_reflects_state() {
510 let emitter = TraceEmitter::new();
511 assert!(!emitter.has_subscribers());
512 let _rx = emitter.subscribe();
513 assert!(emitter.has_subscribers());
514 }
515
516 #[tokio::test]
517 async fn outcome_serialises_correctly() {
518 let event = MatchTraceEvent {
519 event_id: 7,
520 schema_version: 1,
521 received_at_ms: 0,
522 duration_ms: 0,
523 request: RequestSummary {
524 method: "POST".into(),
525 url_path: "/x".into(),
526 headers: vec![],
527 body_json: None,
528 body_truncated: false,
529 },
530 outcome: Outcome::Matched {
531 rule_set_index: 0,
532 rule_index: 2,
533 },
534 dropped_count: 0,
535 };
536 let json = serde_json::to_string(&event).unwrap();
537 assert!(json.contains("\"type\":\"matched\""));
538 assert!(json.contains("\"rule_index\":2"));
539 assert!(json.contains("\"schema_version\":1"));
540 }
541
542 #[tokio::test]
543 async fn tcp_transport_delivers_events() {
544 let emitter = TraceEmitter::new();
545 let emitter_clone = emitter.clone();
546
547 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
551 let bound_addr = listener.local_addr().unwrap();
552
553 tokio::spawn(async move {
555 let (stream, _) = listener.accept().await.unwrap();
556 let rx = emitter_clone.subscribe();
557 let (_, writer) = tokio::io::split(stream);
558 TraceTransport::forward_events(writer, rx).await;
559 });
560
561 let mut client = tokio::net::TcpStream::connect(bound_addr).await.unwrap();
563
564 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
566
567 emitter.emit(
568 42,
569 3,
570 RequestSummary {
571 method: "GET".into(),
572 url_path: "/ping".into(),
573 headers: vec![],
574 body_json: None,
575 body_truncated: false,
576 },
577 Outcome::Miss { status: 404 },
578 );
579
580 use tokio::io::AsyncBufReadExt;
582 let mut reader = tokio::io::BufReader::new(&mut client);
583 let mut line = String::new();
584 tokio::time::timeout(
585 std::time::Duration::from_secs(2),
586 reader.read_line(&mut line),
587 )
588 .await
589 .expect("timeout")
590 .expect("read ok");
591
592 let value: serde_json::Value = serde_json::from_str(line.trim()).expect("valid JSON");
593 assert_eq!(value["request"]["url_path"], "/ping");
594 assert_eq!(value["outcome"]["type"], "miss");
595 assert_eq!(value["schema_version"], 1);
596 }
597
598 #[test]
601 fn enrich_with_body_disabled_by_default() {
602 let emitter = TraceEmitter::new(); let mut summary = RequestSummary {
604 method: "POST".into(),
605 url_path: "/".into(),
606 headers: vec![],
607 body_json: None,
608 body_truncated: false,
609 };
610 let body = serde_json::json!({"action": "create"});
611 emitter.enrich_with_body(&mut summary, Some(&body));
612 assert!(
613 summary.body_json.is_none(),
614 "body should not be captured when disabled"
615 );
616 assert!(!summary.body_truncated);
617 }
618
619 #[test]
620 fn enrich_with_body_enabled_captures_small_body() {
621 let emitter = TraceEmitter::with_config(TraceConfig {
622 capture_body: true,
623 max_body_bytes: 8_192,
624 });
625 let mut summary = RequestSummary {
626 method: "POST".into(),
627 url_path: "/".into(),
628 headers: vec![],
629 body_json: None,
630 body_truncated: false,
631 };
632 let body = serde_json::json!({"action": "create", "user_id": 42});
633 emitter.enrich_with_body(&mut summary, Some(&body));
634 assert!(
635 summary.body_json.is_some(),
636 "body should be captured when enabled"
637 );
638 assert_eq!(summary.body_json.unwrap()["action"], "create");
639 assert!(!summary.body_truncated);
640 }
641
642 #[test]
643 fn enrich_with_body_truncates_oversized_body() {
644 let emitter = TraceEmitter::with_config(TraceConfig {
645 capture_body: true,
646 max_body_bytes: 10,
647 });
648 let mut summary = RequestSummary {
649 method: "POST".into(),
650 url_path: "/".into(),
651 headers: vec![],
652 body_json: None,
653 body_truncated: false,
654 };
655 let body = serde_json::json!({"data": "this is longer than 10 bytes"});
656 emitter.enrich_with_body(&mut summary, Some(&body));
657 assert!(
658 summary.body_json.is_none(),
659 "oversized body should be omitted"
660 );
661 assert!(summary.body_truncated, "body_truncated flag should be set");
662 }
663
664 #[test]
665 fn request_summary_body_json_not_in_serialised_output_when_none() {
666 let summary = RequestSummary {
667 method: "GET".into(),
668 url_path: "/api".into(),
669 headers: vec![],
670 body_json: None,
671 body_truncated: false,
672 };
673 let json = serde_json::to_string(&summary).unwrap();
674 assert!(
675 !json.contains("body_json"),
676 "absent body_json must be skipped"
677 );
678 assert!(
679 !json.contains("body_truncated"),
680 "false body_truncated must be skipped"
681 );
682 }
683}