1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

extern crate actix_web;
extern crate rand;

use actix_web::http::header::HeaderValue;
use actix_web::middleware::{Middleware, Response};
use actix_web::{FromRequest, HttpRequest, HttpResponse, Result};
use rand::distributions::Alphanumeric;
use rand::Rng;

/// The header set by the middleware
pub const REQUEST_ID_HEADER : &str = "request-id";

/// The HTTP Request ID
/// 
/// **note:** must contain as String that is valid to put in HTTP Header values
/// using base62 / base64 is a great way to sanitize the string
/// 
/// It can also be extracted from a request and Helper converter to be able to extract the RequestID easily in an handler
#[derive(Debug, Clone, PartialEq)]
pub struct RequestID(String);

/// Permits retrieving the HttpRequest associated RequestID
pub trait RequestIDGetter {
    /// Returns the HttpRequest RequestID, if the HttpRequest currently has none
    /// it creates one and associates it to the HttpRequest.
    fn request_id(&self) -> RequestID;
}

impl<S> RequestIDGetter for HttpRequest<S> {
    fn request_id(&self) -> RequestID {
        if let Some(req_id) = self.extensions().get::<RequestID>() {
            return req_id.clone();
        }

        let id: String = rand::thread_rng()
            .sample_iter(&Alphanumeric)
            .take(10)
            .collect::<String>();
        self.extensions_mut().insert(RequestID(id.clone()));
        RequestID(id)
    }
}


impl<S> FromRequest<S> for RequestID {
    type Config = ();
    type Result = RequestID;

    #[inline]
    fn from_request(req: &HttpRequest<S>, _: &Self::Config) -> Self::Result {
        req.request_id()
    }
}

/// The RequestID Middleware. It sets a `request-id` HTTP header to the HttpResponse
pub struct RequestIDHeader;
impl<S> Middleware<S> for RequestIDHeader {
    fn response(&self, req: &HttpRequest<S>, mut resp: HttpResponse) -> Result<Response> {
        if let Ok(v) = HeaderValue::from_str(&(req.request_id().0)) {
            resp.headers_mut().append(REQUEST_ID_HEADER, v);
        }

        Ok(Response::Done(resp))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use actix_web::test::TestRequest;

    trait ResponseGetterHelper {
        fn response(self) -> HttpResponse;
    }
    impl ResponseGetterHelper for Response {
        fn response(self) -> HttpResponse {
            match self {
                Response::Done(resp) => resp,
                _ => panic!(),
            }
        }
    }
    
    #[test]
    fn request_id_is_consistent_for_same_request() {
        let req = TestRequest::default().finish();

        assert_eq!(req.request_id(), req.request_id());
        assert_eq!(req.request_id(), RequestID::extract(&req));
    }

    #[test]
    fn request_id_is_new_between_different_requests() {
        let req1 = TestRequest::default().finish();
        let req2 = TestRequest::default().finish();

        assert!(req1.request_id() != req2.request_id());
        assert_eq!(req1.request_id(), req1.request_id());
        assert_eq!(req2.request_id(), req2.request_id());
    }

    #[test]
    fn middleware_adds_request_id_in_headers() {
        let req = TestRequest::default().finish();

        let resp: HttpResponse = HttpResponse::Ok().into();
        let resp = RequestIDHeader.response(&req, resp).unwrap().response();

        let req_id = req.request_id();

        assert_eq!(
            resp.headers().get(REQUEST_ID_HEADER).unwrap().as_bytes(),
            req_id.0.as_bytes()
        );
    }
}