lean_ctx/proxy/forward/
trace_id.rs

1use axum::http::{HeaderMap, HeaderValue};
2use axum::response::Response;
3
4use crate::core::ocla::types::OclaRequestContext;
5
6const TRACE_ID_HEADER: &str = "x-trace-id";
7
8pub fn extract_or_generate_trace_id(headers: &HeaderMap) -> String {
9    headers
10        .get(TRACE_ID_HEADER)
11        .and_then(|value| value.to_str().ok())
12        .filter(|value| !value.is_empty())
13        .map_or_else(
14            || {
15                OclaRequestContext::new(
16                    String::new(),
17                    String::new(),
18                    String::new(),
19                    String::new(),
20                    None,
21                    None,
22                )
23                .trace_id
24            },
25            str::to_owned,
26        )
27}
28
29pub fn inject_trace_id(response: &mut Response, trace_id: &str) {
30    response.headers_mut().insert(
31        TRACE_ID_HEADER,
32        HeaderValue::from_str(trace_id).expect("generated trace ID is a valid header"),
33    );
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39
40    #[test]
41    fn provided_trace_id_is_preserved() {
42        let mut headers = HeaderMap::new();
43        headers.insert(TRACE_ID_HEADER, HeaderValue::from_static("tr-provided"));
44        assert_eq!(extract_or_generate_trace_id(&headers), "tr-provided");
45    }
46
47    #[test]
48    fn missing_trace_id_is_generated() {
49        assert!(extract_or_generate_trace_id(&HeaderMap::new()).starts_with("tr-"));
50    }
51
52    #[test]
53    fn trace_id_is_injected_into_response() {
54        let mut response = Response::new(axum::body::Body::empty());
55        inject_trace_id(&mut response, "tr-test");
56        assert_eq!(response.headers()[TRACE_ID_HEADER], "tr-test");
57    }
58}
lean_ctx/proxy/forward/trace_id.rs

lean_ctx/proxy/forward/
trace_id.rs
