Skip to main content

actix_request_identifier/
lib.rs

1#![deny(unused, missing_docs)]
2//! This is an actix-web middleware to associate every request with a unique ID. This ID can be
3//! used to track errors in an application.
4
5use 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
23/// The default header used for the request ID.
24pub const DEFAULT_HEADER: &str = "x-request-id";
25
26/// Possible error types for the middleware.
27#[derive(Debug, Clone)]
28pub enum Error {
29    /// There is no ID associated with this request.
30    NoAssociatedId,
31}
32
33/// Configuration setting to decide weather the request id from the incoming request header should
34/// be used, if present or if a new one should be generated in any case.
35#[derive(Clone, Copy, PartialEq, Eq, Default)]
36pub enum IdReuse {
37    /// Reuse the incoming request id.
38    UseIncoming,
39    /// Ignore the incoming request id and generate a random one, even if the request supplied an
40    /// id.
41    #[default]
42    IgnoreIncoming,
43}
44
45/// ID wrapper for requests.
46pub 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
55/// A middleware for generating per-request unique IDs
56pub struct RequestIdentifier {
57    header_name: &'static str,
58    id_generator: Generator,
59    use_incoming_id: IdReuse,
60}
61
62/// Request ID that can be extracted in handlers.
63#[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    /// Get the value of the response header for this request id.
85    pub const fn header_value(&self) -> &HeaderValue {
86        &self.0
87    }
88
89    /// Get a string representation of this ID
90    ///
91    /// # Panics
92    ///
93    /// If the header value contains non-ASCII characters
94    pub fn as_str(&self) -> &str {
95        self.0.to_str().expect("Non-ASCII IDs are not supported")
96    }
97}
98
99impl RequestIdentifier {
100    /// Create a default middleware using [`DEFAULT_HEADER`](./constant.DEFAULT_HEADER.html) as the header name and
101    /// UUID v4 for ID generation.
102    #[must_use]
103    #[cfg(feature = "uuid-generator")]
104    pub fn with_uuid() -> Self {
105        Self::default()
106    }
107
108    /// Create a default middleware using [`DEFAULT_HEADER`](./constant.DEFAULT_HEADER.html) as the header name and
109    /// UUID v7 for ID generation.
110    #[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    /// Create a middleware using a custom header name and UUID v4 for ID generation.
117    #[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    /// Change the header name for this middleware.
127    #[must_use]
128    pub const fn header(self, header_name: &'static str) -> Self {
129        Self {
130            header_name,
131            ..self
132        }
133    }
134
135    /// Create a middleware using [`DEFAULT_HEADER`](./constant.DEFAULT_HEADER.html) as the header
136    /// name and IDs as generated by `id_generator`. `id_generator` should return a unique ID
137    /// (ASCII only), each time it is invoked.
138    #[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    /// Change the ID generator for this middleware.
148    #[must_use]
149    pub fn generator(self, id_generator: Generator) -> Self {
150        Self {
151            id_generator,
152            ..self
153        }
154    }
155
156    /// Change the behavior for incoming request id headers. When this is set to
157    /// [`IdReuse::UseIncoming`](./enum.IdReuse.html#variant.UseIncoming) (the default), each request is checked if it
158    /// contains a header by the specified name and if it exists, the id from that header is used, otherwise a random id
159    /// is generated. When this is set to [`IdReuse::IgnoreIncoming`](./enum.IdReuse.html#variant.IgnoreIncoming), the
160    /// id from the request header is ignored.
161    #[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/// Default UUID v4 based ID generator.
182#[cfg(feature = "uuid-generator")]
183fn default_generator() -> HeaderValue {
184    let uuid = Uuid::new_v4();
185    HeaderValue::from_str(&uuid.to_string())
186        // This unwrap can never fail since UUID v4 generated IDs are ASCII-only
187        .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        // This unwrap can never fail since UUID v7 generated IDs are ASCII-only
195        .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        // associate with the request
212        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        // make the id available as an extractor in route handlers
247        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    // Using a macro to reduce code duplication when initializing the test service.
287    // The return type of `test::init_service` is to complicated to create a normal function.
288    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            // use deterministic generator so we can check, if the supplied id is
390            // ignored
391            .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}