1use crate::Result;
2use std::{
3 any::{Any, TypeId},
4 collections::HashMap,
5 net::SocketAddr,
6 sync::{Arc, OnceLock, RwLock},
7};
8
9#[derive(Clone)]
11#[doc(hidden)]
12pub struct CorrelationContext {
13 request_id: String,
14 trace_id: String,
15}
16
17impl CorrelationContext {
18 #[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 pub fn request_id(&self) -> &str {
42 &self.request_id
43 }
44
45 pub fn trace_id(&self) -> &str {
47 &self.trace_id
48 }
49
50 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
61pub 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 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 #[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 #[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 pub fn with_peer_addr(mut self, peer_addr: SocketAddr) -> Self {
122 self.peer_addr = Some(peer_addr);
123 self
124 }
125
126 pub fn peer_addr(&self) -> Option<SocketAddr> {
128 self.peer_addr
129 }
130
131 pub fn method(&self) -> &str {
133 &self.method
134 }
135
136 pub fn path(&self) -> &str {
138 &self.path
139 }
140
141 pub fn request_id(&self) -> &str {
143 self.correlation.request_id()
144 }
145
146 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 pub fn attach_correlation_headers(&self, response: crate::HttpResponse) -> crate::HttpResponse {
157 self.correlation.attach_headers(response)
158 }
159
160 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 pub fn bearer_token(&self) -> Option<&str> {
175 self.header("authorization")?.strip_prefix("Bearer ")
176 }
177
178 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 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 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}