Skip to main content

caelix_core/
context.rs

1use crate::Result;
2use std::{
3    any::{Any, TypeId},
4    collections::HashMap,
5    net::SocketAddr,
6    sync::{Arc, OnceLock, RwLock},
7};
8
9/// Request correlation identifiers shared by lightweight and full request paths.
10#[derive(Clone)]
11#[doc(hidden)]
12pub struct CorrelationContext {
13    request_id: String,
14    trace_id: String,
15}
16
17impl CorrelationContext {
18    /// Builds correlation identifiers from the relevant incoming header values.
19    #[doc(hidden)]
20    pub fn from_header_values(
21        request_id: Option<&str>,
22        traceparent: Option<&str>,
23        trace_id: Option<&str>,
24    ) -> Self {
25        let request_id = request_id
26            .and_then(valid_correlation_id)
27            .map(str::to_owned)
28            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
29        let trace_id = traceparent
30            .and_then(trace_id_from_traceparent)
31            .or_else(|| trace_id.and_then(valid_correlation_id))
32            .unwrap_or(&request_id)
33            .to_owned();
34        Self {
35            request_id,
36            trace_id,
37        }
38    }
39
40    /// Returns the accepted incoming `X-Request-Id`, or a generated UUID.
41    pub fn request_id(&self) -> &str {
42        &self.request_id
43    }
44
45    /// Returns the W3C `traceparent` trace ID, `X-Trace-Id`, or request ID fallback.
46    pub fn trace_id(&self) -> &str {
47        &self.trace_id
48    }
49
50    /// Adds these correlation identifiers to an HTTP response.
51    pub fn attach_headers(&self, mut response: crate::HttpResponse) -> crate::HttpResponse {
52        response.headers.retain(|(name, _)| {
53            !name.eq_ignore_ascii_case("x-request-id") && !name.eq_ignore_ascii_case("x-trace-id")
54        });
55        response.insert_header("X-Request-Id", self.request_id.clone());
56        response.insert_header("X-Trace-Id", self.trace_id.clone());
57        response
58    }
59}
60
61/// Public Caelix type `RequestContext`.
62pub struct RequestContext {
63    method: String,
64    path: String,
65    correlation: CorrelationContext,
66    headers: HashMap<String, String>,
67    cookies: OnceLock<HashMap<String, String>>,
68    peer_addr: Option<SocketAddr>,
69    extensions: RwLock<Option<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>>,
70}
71
72impl RequestContext {
73    /// Runs the `new` public API operation.
74    pub fn new(
75        method: impl Into<String>,
76        path: impl Into<String>,
77        headers: HashMap<String, String>,
78    ) -> Self {
79        let headers = headers
80            .into_iter()
81            .map(|(name, value)| (name.to_ascii_lowercase(), value))
82            .collect::<HashMap<_, _>>();
83        Self::from_normalized_headers(method, path, headers)
84    }
85
86    /// Builds a context from already lowercase header names without copying them again.
87    #[doc(hidden)]
88    pub fn from_normalized_headers(
89        method: impl Into<String>,
90        path: impl Into<String>,
91        headers: HashMap<String, String>,
92    ) -> Self {
93        let correlation = CorrelationContext::from_header_values(
94            headers.get("x-request-id").map(String::as_str),
95            headers.get("traceparent").map(String::as_str),
96            headers.get("x-trace-id").map(String::as_str),
97        );
98        Self::from_normalized_headers_with_correlation(method, path, headers, correlation)
99    }
100
101    /// Builds a context while reusing correlation identifiers already parsed by the adapter.
102    #[doc(hidden)]
103    pub fn from_normalized_headers_with_correlation(
104        method: impl Into<String>,
105        path: impl Into<String>,
106        headers: HashMap<String, String>,
107        correlation: CorrelationContext,
108    ) -> Self {
109        Self {
110            method: method.into(),
111            path: path.into(),
112            correlation,
113            headers,
114            cookies: OnceLock::new(),
115            peer_addr: None,
116            extensions: RwLock::new(None),
117        }
118    }
119
120    /// Sets the immediate socket peer address.
121    pub fn with_peer_addr(mut self, peer_addr: SocketAddr) -> Self {
122        self.peer_addr = Some(peer_addr);
123        self
124    }
125
126    /// Returns the immediate socket peer address, when the runtime supplied one.
127    pub fn peer_addr(&self) -> Option<SocketAddr> {
128        self.peer_addr
129    }
130
131    /// Runs the `method` public API operation.
132    pub fn method(&self) -> &str {
133        &self.method
134    }
135
136    /// Runs the `path` public API operation.
137    pub fn path(&self) -> &str {
138        &self.path
139    }
140
141    /// Returns the accepted incoming `X-Request-Id`, or a generated UUID.
142    pub fn request_id(&self) -> &str {
143        self.correlation.request_id()
144    }
145
146    /// Returns the W3C `traceparent` trace ID, `X-Trace-Id`, or request ID fallback.
147    pub fn trace_id(&self) -> &str {
148        self.correlation.trace_id()
149    }
150
151    pub(crate) fn correlation(&self) -> &CorrelationContext {
152        &self.correlation
153    }
154
155    /// Adds this request's correlation identifiers to an HTTP response.
156    pub fn attach_correlation_headers(&self, response: crate::HttpResponse) -> crate::HttpResponse {
157        self.correlation.attach_headers(response)
158    }
159
160    /// Runs the `header` public API operation.
161    pub fn header(&self, name: &str) -> Option<&str> {
162        self.headers
163            .get(name)
164            .or_else(|| {
165                name.bytes()
166                    .any(|byte| byte.is_ascii_uppercase())
167                    .then(|| self.headers.get(&name.to_ascii_lowercase()))
168                    .flatten()
169            })
170            .map(String::as_str)
171    }
172
173    /// Runs the `bearer_token` public API operation.
174    pub fn bearer_token(&self) -> Option<&str> {
175        self.header("authorization")?.strip_prefix("Bearer ")
176    }
177
178    /// Returns the first cookie with `name`.
179    pub fn cookie(&self, name: &str) -> Option<&str> {
180        self.cookies
181            .get_or_init(|| parse_cookies(self.headers.get("cookie").map(String::as_str)))
182            .get(name)
183            .map(String::as_str)
184    }
185
186    /// Runs the `set` public API operation.
187    pub fn set<T: Send + Sync + 'static>(&self, value: T) -> Result<()> {
188        self.extensions
189            .write()
190            .map_err(|_| {
191                crate::exception::startup_error("request context extensions lock poisoned")
192            })?
193            .get_or_insert_with(HashMap::new)
194            .insert(TypeId::of::<T>(), Arc::new(value));
195        Ok(())
196    }
197
198    /// Runs the `get` public API operation.
199    pub fn get<T: Send + Sync + 'static>(&self) -> Result<Option<Arc<T>>> {
200        let extensions = self.extensions.read().map_err(|_| {
201            crate::exception::startup_error("request context extensions lock poisoned")
202        })?;
203        let value = extensions
204            .as_ref()
205            .and_then(|extensions| extensions.get(&TypeId::of::<T>()))
206            .cloned();
207
208        let Some(value) = value else {
209            return Ok(None);
210        };
211
212        value
213            .clone()
214            .downcast::<T>()
215            .map(Some)
216            .map_err(|_| crate::exception::startup_error("request context extension type mismatch"))
217    }
218}
219
220fn valid_correlation_id(value: &str) -> Option<&str> {
221    let value = value.trim();
222    (!value.is_empty()
223        && value.len() <= 128
224        && value
225            .bytes()
226            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.')))
227    .then_some(value)
228}
229
230fn trace_id_from_traceparent(value: &str) -> Option<&str> {
231    let mut parts = value.trim().split('-');
232    let version = parts.next()?;
233    let trace_id = parts.next()?;
234    let parent_id = parts.next()?;
235    let flags = parts.next()?;
236    if parts.next().is_some()
237        || version.len() != 2
238        || trace_id.len() != 32
239        || parent_id.len() != 16
240        || flags.len() != 2
241        || trace_id.bytes().all(|byte| byte == b'0')
242        || !version
243            .bytes()
244            .chain(trace_id.bytes())
245            .chain(parent_id.bytes())
246            .chain(flags.bytes())
247            .all(|byte| byte.is_ascii_hexdigit())
248    {
249        return None;
250    }
251    Some(trace_id)
252}
253
254fn parse_cookies(header: Option<&str>) -> HashMap<String, String> {
255    let mut cookies = HashMap::new();
256    for pair in header.into_iter().flat_map(|value| value.split(';')) {
257        let pair = pair.trim();
258        let Some((name, value)) = pair.split_once('=') else {
259            continue;
260        };
261        let name = name.trim();
262        if name.is_empty() {
263            continue;
264        }
265        let Ok(name) = cookie::Cookie::parse_encoded(format!("{name}=")) else {
266            continue;
267        };
268        let Ok(value) = cookie::Cookie::parse_encoded(format!("value={}", value.trim())) else {
269            continue;
270        };
271        cookies
272            .entry(name.name().to_string())
273            .or_insert_with(|| value.value().to_string());
274    }
275    cookies
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn parses_cookie_edge_cases_once() {
284        let ctx = RequestContext::new(
285            "GET",
286            "/",
287            HashMap::from([(
288                "Cookie".into(),
289                " theme = dark%20mode ; token=a=b=c; broken; =bad; dup=first; dup=second; na%6De=value"
290                    .into(),
291            )]),
292        );
293        assert_eq!(ctx.cookie("theme"), Some("dark mode"));
294        assert_eq!(ctx.cookie("token"), Some("a=b=c"));
295        assert_eq!(ctx.cookie("dup"), Some("first"));
296        assert_eq!(ctx.cookie("name"), Some("value"));
297        assert_eq!(ctx.cookie("broken"), None);
298    }
299
300    #[test]
301    fn defers_cookie_parsing_and_extension_storage_until_used() {
302        let ctx = RequestContext::new(
303            "GET",
304            "/",
305            HashMap::from([("Cookie".into(), "session=abc".into())]),
306        );
307
308        assert!(ctx.cookies.get().is_none());
309        assert!(ctx.extensions.read().unwrap().is_none());
310        assert_eq!(ctx.get::<String>().unwrap(), None);
311        assert!(ctx.extensions.read().unwrap().is_none());
312
313        assert_eq!(ctx.cookie("session"), Some("abc"));
314        assert!(ctx.cookies.get().is_some());
315
316        ctx.set("authenticated".to_string()).unwrap();
317        assert!(ctx.extensions.read().unwrap().is_some());
318    }
319
320    #[test]
321    fn derives_request_and_trace_ids_from_headers() {
322        let ctx = RequestContext::new(
323            "GET",
324            "/orders",
325            HashMap::from([
326                ("X-Request-Id".into(), "request-123".into()),
327                (
328                    "traceparent".into(),
329                    "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".into(),
330                ),
331            ]),
332        );
333
334        assert_eq!(ctx.request_id(), "request-123");
335        assert_eq!(ctx.trace_id(), "4bf92f3577b34da6a3ce929d0e0e4736");
336    }
337
338    #[test]
339    fn generates_safe_correlation_ids_for_missing_or_invalid_headers() {
340        let ctx = RequestContext::new(
341            "POST",
342            "/orders",
343            HashMap::from([("X-Request-Id".into(), "unsafe\nvalue".into())]),
344        );
345
346        assert!(uuid::Uuid::parse_str(ctx.request_id()).is_ok());
347        assert_eq!(ctx.trace_id(), ctx.request_id());
348    }
349
350    #[test]
351    fn attaches_correlation_ids_to_responses() {
352        let ctx = RequestContext::new(
353            "GET",
354            "/orders",
355            HashMap::from([
356                ("X-Request-Id".into(), "request-123".into()),
357                ("X-Trace-Id".into(), "trace-456".into()),
358            ]),
359        );
360        let response =
361            ctx.attach_correlation_headers(crate::HttpResponse::text(http::StatusCode::OK, "ok"));
362
363        assert!(
364            response
365                .headers
366                .contains(&("X-Request-Id".into(), "request-123".into()))
367        );
368        assert!(
369            response
370                .headers
371                .contains(&("X-Trace-Id".into(), "trace-456".into()))
372        );
373    }
374}