actix_request_identifier/
lib.rs1#![deny(unused, missing_docs)]
2use std::{
6 fmt::{self, Display},
7 pin::Pin,
8 task::{Context, Poll},
9};
10
11use actix_web::{
12 dev::{Payload, Service, ServiceRequest, ServiceResponse, Transform},
13 error::ResponseError,
14 http::header::{HeaderName, HeaderValue},
15 Error as ActixError, FromRequest, HttpMessage, HttpRequest,
16};
17use futures::{
18 future::{ok, ready, Ready},
19 Future,
20};
21use uuid::Uuid;
22
23pub const DEFAULT_HEADER: &str = "x-request-id";
25
26#[derive(Debug, Clone)]
28pub enum Error {
29 NoAssociatedId,
31}
32
33#[derive(Clone, Copy, PartialEq, Eq, Default)]
36pub enum IdReuse {
37 UseIncoming,
39 #[default]
42 IgnoreIncoming,
43}
44
45pub struct RequestIdMiddleware<S> {
47 service: S,
48 header_name: HeaderName,
49 id_generator: Generator,
50 use_incoming_id: IdReuse,
51}
52
53type Generator = fn() -> HeaderValue;
54
55pub struct RequestIdentifier {
57 header_name: &'static str,
58 id_generator: Generator,
59 use_incoming_id: IdReuse,
60}
61
62#[derive(Clone)]
64pub struct RequestId(HeaderValue);
65
66impl Display for RequestId {
67 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
68 write!(f, "{}", self.as_str())
69 }
70}
71
72impl ResponseError for Error {}
73
74impl fmt::Display for Error {
75 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
76 use Error::NoAssociatedId;
77 match self {
78 NoAssociatedId => write!(fmt, "NoAssociatedId"),
79 }
80 }
81}
82
83impl RequestId {
84 pub const fn header_value(&self) -> &HeaderValue {
86 &self.0
87 }
88
89 pub fn as_str(&self) -> &str {
95 self.0.to_str().expect("Non-ASCII IDs are not supported")
96 }
97}
98
99impl RequestIdentifier {
100 #[must_use]
103 #[cfg(feature = "uuid-generator")]
104 pub fn with_uuid() -> Self {
105 Self::default()
106 }
107
108 #[must_use]
111 #[cfg(feature = "uuid-v7-generator")]
112 pub fn with_uuid_v7() -> Self {
113 Self::with_generator(uuid_v7_generator)
114 }
115
116 #[must_use]
118 #[cfg(feature = "uuid-generator")]
119 pub fn with_header(header_name: &'static str) -> Self {
120 Self {
121 header_name,
122 ..Default::default()
123 }
124 }
125
126 #[must_use]
128 pub const fn header(self, header_name: &'static str) -> Self {
129 Self {
130 header_name,
131 ..self
132 }
133 }
134
135 #[must_use]
139 pub fn with_generator(id_generator: Generator) -> Self {
140 Self {
141 id_generator,
142 header_name: DEFAULT_HEADER,
143 use_incoming_id: IdReuse::default(),
144 }
145 }
146
147 #[must_use]
149 pub fn generator(self, id_generator: Generator) -> Self {
150 Self {
151 id_generator,
152 ..self
153 }
154 }
155
156 #[must_use]
162 pub fn use_incoming_id(self, use_incoming_id: IdReuse) -> Self {
163 Self {
164 use_incoming_id,
165 ..self
166 }
167 }
168}
169
170#[cfg(feature = "uuid-generator")]
171impl Default for RequestIdentifier {
172 fn default() -> Self {
173 Self {
174 header_name: DEFAULT_HEADER,
175 id_generator: default_generator,
176 use_incoming_id: IdReuse::default(),
177 }
178 }
179}
180
181#[cfg(feature = "uuid-generator")]
183fn default_generator() -> HeaderValue {
184 let uuid = Uuid::new_v4();
185 HeaderValue::from_str(&uuid.to_string())
186 .unwrap()
188}
189
190#[cfg(feature = "uuid-v7-generator")]
191fn uuid_v7_generator() -> HeaderValue {
192 let uuid = Uuid::now_v7();
193 HeaderValue::from_str(&uuid.to_string())
194 .unwrap()
196}
197
198impl<S, B> Transform<S, ServiceRequest> for RequestIdentifier
199where
200 S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = ActixError>,
201 S::Future: 'static,
202 B: 'static,
203{
204 type Response = S::Response;
205 type Error = S::Error;
206 type Transform = RequestIdMiddleware<S>;
207 type InitError = ();
208 type Future = Ready<Result<Self::Transform, Self::InitError>>;
209
210 fn new_transform(&self, service: S) -> Self::Future {
211 ok(RequestIdMiddleware {
213 service,
214 header_name: HeaderName::from_static(self.header_name),
215 id_generator: self.id_generator,
216 use_incoming_id: self.use_incoming_id,
217 })
218 }
219}
220
221#[allow(clippy::type_complexity)]
222impl<S, B> Service<ServiceRequest> for RequestIdMiddleware<S>
223where
224 S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = ActixError>,
225 S::Future: 'static,
226 B: 'static,
227{
228 type Response = S::Response;
229 type Error = S::Error;
230 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
231
232 fn poll_ready(&self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
233 self.service.poll_ready(ctx)
234 }
235
236 fn call(&self, request: ServiceRequest) -> Self::Future {
237 let header_name = self.header_name.clone();
238 let header_value = match self.use_incoming_id {
239 IdReuse::UseIncoming => request
240 .headers()
241 .get(&header_name)
242 .map_or_else(self.id_generator, Clone::clone),
243 IdReuse::IgnoreIncoming => (self.id_generator)(),
244 };
245
246 let request_id = RequestId(header_value.clone());
248 request.extensions_mut().insert(request_id);
249
250 let fut = self.service.call(request);
251 Box::pin(async move {
252 let mut response = fut.await?;
253
254 response.headers_mut().insert(header_name, header_value);
255
256 Ok(response)
257 })
258 }
259}
260
261impl FromRequest for RequestId {
262 type Error = Error;
263 type Future = Ready<Result<Self, Self::Error>>;
264
265 fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
266 ready(
267 req.extensions()
268 .get::<RequestId>()
269 .cloned()
270 .ok_or(Error::NoAssociatedId),
271 )
272 }
273}
274
275#[cfg(test)]
276mod tests {
277 use super::*;
278 use actix_web::{test, web, App};
279 use bytes::Bytes;
280
281 #[allow(clippy::unused_async)]
282 async fn handler(id: RequestId) -> String {
283 id.as_str().to_string()
284 }
285
286 macro_rules! service {
289 ($middleware:expr) => {
290 test::init_service(
291 App::new()
292 .wrap($middleware)
293 .route("/", web::get().to(handler)),
294 )
295 .await
296 };
297 }
298
299 async fn test_get(middleware: RequestIdentifier) -> ServiceResponse {
300 let service = service!(middleware);
301 test::call_service(&service, test::TestRequest::get().uri("/").to_request()).await
302 }
303
304 #[actix_web::test]
305 async fn default_identifier() {
306 let resp = test_get(RequestIdentifier::with_uuid()).await;
307 let uid = resp
308 .headers()
309 .get(HeaderName::from_static(DEFAULT_HEADER))
310 .map(|v| v.to_str().unwrap().to_string())
311 .unwrap();
312 let body: Bytes = test::read_body(resp).await;
313 let body = String::from_utf8_lossy(&body);
314 assert_eq!(uid, body);
315 }
316
317 #[cfg(feature = "uuid-v7-generator")]
318 #[actix_web::test]
319 async fn uuid_v7() {
320 let resp = test_get(RequestIdentifier::with_uuid_v7()).await;
321 let uid = resp
322 .headers()
323 .get(HeaderName::from_static(DEFAULT_HEADER))
324 .map(|v| v.to_str().unwrap().to_string())
325 .unwrap();
326 let body: Bytes = test::read_body(resp).await;
327 let body = String::from_utf8_lossy(&body);
328 assert_eq!(uid, body);
329 }
330
331 #[actix_web::test]
332 async fn deterministic_identifier() {
333 let resp = test_get(RequestIdentifier::with_generator(|| {
334 HeaderValue::from_static("look ma, i'm an id")
335 }))
336 .await;
337 let uid = resp
338 .headers()
339 .get(HeaderName::from_static(DEFAULT_HEADER))
340 .map(|v| v.to_str().unwrap().to_string())
341 .unwrap();
342 let body: Bytes = test::read_body(resp).await;
343 let body = String::from_utf8_lossy(&body);
344 assert_eq!(uid, body);
345 }
346
347 #[actix_web::test]
348 async fn custom_header() {
349 let resp = test_get(RequestIdentifier::with_header("custom-header")).await;
350 assert!(resp
351 .headers()
352 .get(HeaderName::from_static(DEFAULT_HEADER))
353 .is_none());
354 let uid = resp
355 .headers()
356 .get(HeaderName::from_static("custom-header"))
357 .map(|v| v.to_str().unwrap().to_string())
358 .unwrap();
359 let body: Bytes = test::read_body(resp).await;
360 let body = String::from_utf8_lossy(&body);
361 assert_eq!(uid, body);
362 }
363
364 #[actix_web::test]
365 async fn existing_request_id() {
366 let uuid4 = Uuid::new_v4().to_string();
367 let service =
368 service!(RequestIdentifier::with_uuid().use_incoming_id(IdReuse::UseIncoming));
369 let req = test::TestRequest::get()
370 .insert_header((DEFAULT_HEADER, uuid4.as_str()))
371 .uri("/")
372 .to_request();
373 let resp = test::call_service(&service, req).await;
374 let uid = resp
375 .headers()
376 .get(HeaderName::from_static(DEFAULT_HEADER))
377 .map(|v| v.to_str().unwrap().to_string())
378 .unwrap();
379 assert_eq!(uid, uuid4);
380 let body: Bytes = test::read_body(resp).await;
381 let body = String::from_utf8_lossy(&body);
382 assert_eq!(body, uuid4);
383 }
384
385 #[actix_web::test]
386 async fn ignore_existing_request_id() {
387 let uuid4 = Uuid::new_v4().to_string();
388 let service = service!(RequestIdentifier::with_uuid()
389 .generator(|| HeaderValue::from_static("0")));
392 let req = test::TestRequest::get()
393 .insert_header((DEFAULT_HEADER, uuid4.as_str()))
394 .uri("/")
395 .to_request();
396 let resp = test::call_service(&service, req).await;
397 let uid = resp
398 .headers()
399 .get(HeaderName::from_static(DEFAULT_HEADER))
400 .map(|v| v.to_str().unwrap().to_string())
401 .unwrap();
402 assert_eq!(uid, "0");
403 let body: Bytes = test::read_body(resp).await;
404 let body = String::from_utf8_lossy(&body);
405 assert_eq!(body, "0");
406 }
407}