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

//! Defines a unique id per `Request` that should be output with all logging.

use hyper::header::HeaderMap;
use log::trace;
use uuid::Uuid;

use crate::state::{FromState, State};

/// A container type for the value returned by `request_id`.
pub(super) struct RequestId {
    val: String,
}

/// Sets a unique identifier for the request if it has not already been stored.
///
/// The unique identifier chosen depends on the the request headers:
///
/// 1. If the header `X-Request-ID` is provided this value is used as-is;
/// 2. Alternatively creates and stores a UUID v4 value.
///
/// This function is invoked by `GothamService` before handing control to its `Router`, to ensure
/// that a value for `RequestId` is always available.
pub(crate) fn set_request_id<'a>(state: &'a mut State) -> &'a str {
    if !state.has::<RequestId>() {
        let request_id = match HeaderMap::borrow_from(state).get("X-Request-ID") {
            Some(ex_req_id) => {
                let id = String::from_utf8(ex_req_id.as_bytes().into()).unwrap();
                trace!(
                    "[{}] RequestId set from external source via X-Request-ID header",
                    id
                );
                RequestId { val: id }
            }
            None => {
                let val = Uuid::new_v4().to_hyphenated().to_string();
                trace!("[{}] RequestId generated internally", val);
                RequestId { val }
            }
        };
        state.put(request_id);
    };

    request_id(state)
}

/// Returns the request ID associated with the current request.
///
/// This is typically used for logging and correlating events that occurred within a request.
///
/// # Panics
///
/// Will panic if `State` does not contain a request ID, which is an invalid state. The request ID
/// should always be populated by Gotham before a `Router` is invoked.
pub fn request_id(state: &State) -> &str {
    match RequestId::try_borrow_from(state) {
        Some(request_id) => &request_id.val,
        None => panic!("RequestId must be populated before application code is invoked"),
    }
}

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

    #[test]
    #[should_panic(expected = "RequestId must be populated before application code is invoked")]
    fn panics_before_request_id_set() {
        let state = State::new();
        request_id(&state);
    }

    #[test]
    fn uses_an_external_request_id() {
        let mut state = State::new();

        let mut headers = HeaderMap::new();
        headers.insert("X-Request-ID", "1-2-3-4".to_owned().parse().unwrap());
        state.put(headers);

        {
            let r = set_request_id(&mut state);
            assert_eq!("1-2-3-4", r);
        };
        assert_eq!("1-2-3-4", request_id(&state));
    }

    #[test]
    fn sets_a_unique_request_id() {
        let mut state = State::new();
        state.put(HeaderMap::new());

        {
            let r = set_request_id(&mut state);
            assert_eq!(4, Uuid::parse_str(r).unwrap().get_version_num());
        };
        assert_eq!(
            4,
            Uuid::parse_str(request_id(&state))
                .unwrap()
                .get_version_num()
        );
    }

    #[test]
    fn does_not_overwrite_existant_request_id() {
        let mut state = State::new();
        state.put(RequestId {
            val: "1-2-3-4".to_string(),
        });

        {
            set_request_id(&mut state);
        }
        assert_eq!("1-2-3-4", request_id(&state));
    }
}