From f3711a8f5d36a9f290678220f8d7502e3aa10669 Mon Sep 17 00:00:00 2001 From: Yoshua Wuyts Date: Wed, 3 Jun 2020 19:27:44 +0200 Subject: [PATCH] Add trace submodule --- Cargo.toml | 1 + src/headers/headers.rs | 12 ++ src/lib.rs | 5 +- src/trace/mod.rs | 268 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 283 insertions(+), 3 deletions(-) create mode 100644 src/trace/mod.rs diff --git a/Cargo.toml b/Cargo.toml index 3c61714c..33740c11 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ url = "2.1.1" serde_json = "1.0.51" serde = { version = "1.0.106", features = ["derive"] } serde_urlencoded = "0.6.1" +rand = "0.7.3" [dev-dependencies] http = "0.2.0" diff --git a/src/headers/headers.rs b/src/headers/headers.rs index 774f15ac..8555a010 100644 --- a/src/headers/headers.rs +++ b/src/headers/headers.rs @@ -183,6 +183,18 @@ impl Debug for Headers { } } +impl AsRef for Headers { + fn as_ref(&self) -> &Headers { + self + } +} + +impl AsMut for Headers { + fn as_mut(&mut self) -> &mut Headers { + self + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/lib.rs b/src/lib.rs index df0d7619..fb4bbc46 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -132,6 +132,7 @@ mod version; cfg_unstable! { pub mod upgrade; + pub mod trace; mod client; mod server; @@ -165,10 +166,8 @@ pub use crate::url::Url; #[doc(inline)] pub use crate::cookies::Cookie; -#[doc(inline)] -pub mod trailers; - pub mod security; +pub mod trailers; #[cfg(feature = "hyperium_http")] mod hyperium_http; diff --git a/src/trace/mod.rs b/src/trace/mod.rs new file mode 100644 index 00000000..ab6a273a --- /dev/null +++ b/src/trace/mod.rs @@ -0,0 +1,268 @@ +//! Extract and inject [trace context](https://w3c.github.io/trace-context/) headers. +//! +//! ## Examples +//! +//! ``` +//! use http_types::trace::TraceContext; +//! +//! let mut res = http_types::Response::new(200); +//! +//! res.insert_header( +//! "traceparent", +//! "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01" +//! ); +//! +//! let context = TraceContext::extract(&res).unwrap(); +//! +//! let trace_id = u128::from_str_radix("0af7651916cd43dd8448eb211c80319c", 16); +//! let parent_id = u64::from_str_radix("00f067aa0ba902b7", 16); +//! +//! assert_eq!(context.trace_id(), trace_id.unwrap()); +//! assert_eq!(context.parent_id(), parent_id.ok()); +//! assert_eq!(context.sampled(), true); +//! ``` + +use rand::Rng; +use std::fmt; + +use crate::Headers; + +/// A TraceContext object +#[derive(Debug)] +pub struct TraceContext { + id: u64, + version: u8, + trace_id: u128, + parent_id: Option, + flags: u8, +} + +impl TraceContext { + /// Create and return TraceContext object based on `traceparent` HTTP header. + /// + /// ## Examples + /// ``` + /// use http_types::trace::TraceContext; + /// + /// let mut res = http_types::Response::new(200); + /// res.insert_header( + /// "traceparent", + /// "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01" + /// ); + /// + /// let context = TraceContext::extract(&res).unwrap(); + /// + /// let trace_id = u128::from_str_radix("0af7651916cd43dd8448eb211c80319c", 16); + /// let parent_id = u64::from_str_radix("00f067aa0ba902b7", 16); + /// + /// assert_eq!(context.trace_id(), trace_id.unwrap()); + /// assert_eq!(context.parent_id(), parent_id.ok()); + /// assert_eq!(context.sampled(), true); + /// ``` + pub fn extract(headers: impl AsRef) -> crate::Result { + let headers = headers.as_ref(); + let mut rng = rand::thread_rng(); + + let traceparent = match headers.get("traceparent") { + Some(header) => header.as_str(), + None => return Ok(Self::new_root()), + }; + + let parts: Vec<&str> = traceparent.split('-').collect(); + + Ok(Self { + id: rng.gen(), + version: u8::from_str_radix(parts[0], 16)?, + trace_id: u128::from_str_radix(parts[1], 16)?, + parent_id: Some(u64::from_str_radix(parts[2], 16)?), + flags: u8::from_str_radix(parts[3], 16)?, + }) + } + + /// Generate a new TraceContect object without a parent. + /// + /// By default root TraceContext objects are sampled. + /// To mark it unsampled, call `context.set_sampled(false)`. + /// + /// ## Examples + /// ``` + /// use http_types::trace::TraceContext; + /// + /// let context = TraceContext::new_root(); + /// + /// assert_eq!(context.parent_id(), None); + /// assert_eq!(context.sampled(), true); + /// ``` + pub fn new_root() -> Self { + let mut rng = rand::thread_rng(); + + Self { + id: rng.gen(), + version: 0, + trace_id: rng.gen(), + parent_id: None, + flags: 1, + } + } + + /// Add the traceparent header to the http headers + /// + /// ## Examples + /// ``` + /// use http_types::trace::TraceContext; + /// use http_types::{Request, Response, Url, Method}; + /// + /// let mut req = Request::new(Method::Get, Url::parse("https://example.com").unwrap()); + /// req.insert_header( + /// "traceparent", + /// "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01" + /// ); + /// + /// let parent = TraceContext::extract(&req).unwrap(); + /// + /// let mut res = Response::new(200); + /// parent.inject(&mut res); + /// + /// let child = TraceContext::extract(&res).unwrap(); + /// + /// assert_eq!(child.version(), parent.version()); + /// assert_eq!(child.trace_id(), parent.trace_id()); + /// assert_eq!(child.parent_id(), Some(parent.id())); + /// ``` + pub fn inject(&self, mut headers: impl AsMut) { + let headers = headers.as_mut(); + headers.insert("traceparent", format!("{}", self)); + } + + /// Generate a child of the current TraceContext and return it. + /// + /// The child will have a new randomly genrated `id` and its `parent_id` will be set to the + /// `id` of this TraceContext. + pub fn child(&self) -> Self { + let mut rng = rand::thread_rng(); + + Self { + id: rng.gen(), + version: self.version, + trace_id: self.trace_id, + parent_id: Some(self.id), + flags: self.flags, + } + } + + /// Return the id of the TraceContext. + pub fn id(&self) -> u64 { + self.id + } + + /// Return the version of the TraceContext spec used. + /// + /// You probably don't need this. + pub fn version(&self) -> u8 { + self.version + } + + /// Return the trace id of the TraceContext. + /// + /// All children will have the same `trace_id`. + pub fn trace_id(&self) -> u128 { + self.trace_id + } + + /// Return the id of the parent TraceContext. + #[inline] + pub fn parent_id(&self) -> Option { + self.parent_id + } + + /// Returns true if the trace is sampled + /// + /// ## Examples + /// + /// ``` + /// use http_types::trace::TraceContext; + /// use http_types::Response; + /// + /// let mut res = Response::new(200); + /// res.insert_header("traceparent", "00-00000000000000000000000000000001-0000000000000002-01"); + /// let context = TraceContext::extract(&res).unwrap(); + /// assert_eq!(context.sampled(), true); + /// ``` + pub fn sampled(&self) -> bool { + (self.flags & 0b00000001) == 1 + } + + /// Change sampled flag + /// + /// ## Examples + /// + /// ``` + /// use http_types::trace::TraceContext; + /// + /// let mut context = TraceContext::new_root(); + /// assert_eq!(context.sampled(), true); + /// context.set_sampled(false); + /// assert_eq!(context.sampled(), false); + /// ``` + pub fn set_sampled(&mut self, sampled: bool) { + let x = sampled as u8; + self.flags ^= (x ^ self.flags) & (1 << 0); + } +} + +impl fmt::Display for TraceContext { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "{:02x}-{:032x}-{:016x}-{:02x}", + self.version, self.trace_id, self.id, self.flags + ) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn default() -> crate::Result<()> { + let mut headers = crate::Headers::new(); + headers.insert("traceparent", "00-01-deadbeef-00"); + let context = TraceContext::extract(&mut headers)?; + assert_eq!(context.version(), 0); + assert_eq!(context.trace_id(), 1); + assert_eq!(context.parent_id().unwrap(), 3735928559); + assert_eq!(context.flags, 0); + assert_eq!(context.sampled(), false); + Ok(()) + } + + #[test] + fn no_header() -> crate::Result<()> { + let mut headers = crate::Headers::new(); + let context = TraceContext::extract(&mut headers)?; + assert_eq!(context.version(), 0); + assert_eq!(context.parent_id(), None); + assert_eq!(context.flags, 1); + assert_eq!(context.sampled(), true); + Ok(()) + } + + #[test] + fn not_sampled() -> crate::Result<()> { + let mut headers = crate::Headers::new(); + headers.insert("traceparent", "00-01-02-00"); + let context = TraceContext::extract(&mut headers)?; + assert_eq!(context.sampled(), false); + Ok(()) + } + + #[test] + fn sampled() -> crate::Result<()> { + let mut headers = crate::Headers::new(); + headers.insert("traceparent", "00-01-02-01"); + let context = TraceContext::extract(&mut headers)?; + assert_eq!(context.sampled(), true); + Ok(()) + } +}