Skip to content

Commit

Permalink
stackdriver: parse trace context without regex (#47)
Browse files Browse the repository at this point in the history
  • Loading branch information
@djc
djc committed Mar 12, 2024
1 parent 57d9bb9 commit 93dc0d9
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 77 deletions.
4 changes: 1 addition & 3 deletions opentelemetry-stackdriver/Cargo.toml
View file
Expand Up @@ -25,8 +25,6 @@ thiserror = "1.0.30"
tonic = { version = "0.11", features = ["gzip", "tls", "transport"] }
yup-oauth2 = { version = "8.1.0", optional = true }
once_cell = { version = "1.19", optional = true }
# we don't need unicode support
regex = { version = "1.10", default-features = false, features = ["std", "perf"], optional = true }

# Futures
futures-core = "0.3"
Expand All @@ -38,7 +36,7 @@ default = ["yup-authorizer", "tls-native-roots"]
yup-authorizer = ["hyper-rustls", "yup-oauth2"]
tls-native-roots = ["tonic/tls-roots"]
tls-webpki-roots = ["tonic/tls-webpki-roots"]
propagator = ["once_cell", "regex"]
propagator = ["once_cell"]

[dev-dependencies]
reqwest = "0.11.9"
Expand Down
100 changes: 26 additions & 74 deletions opentelemetry-stackdriver/src/google_trace_context_propagator.rs
View file
@@ -1,9 +1,10 @@
use std::str::FromStr;

use once_cell::sync::Lazy;
use opentelemetry::propagation::text_map_propagator::FieldIter;
use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState};
use opentelemetry::Context;
use regex::Regex;

/// Propagates span context in the Google Cloud Trace format,
/// using the __X-Cloud-Trace-Context__ header.
Expand All @@ -25,63 +26,30 @@ pub struct GoogleTraceContextPropagator {
// - trace flags is optional, 0 to 9 (0 - not sampled, missing or any other number - sampled)

const CLOUD_TRACE_CONTEXT_HEADER: &str = "X-Cloud-Trace-Context";
const GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR: &str =
r"^(?P<trace_id>[0-9a-f]{32})/(?P<span_id>[0-9]{1,20})(;o=(?P<trace_flags>[0-9]))?$";

static TRACE_CONTEXT_HEADER_FIELDS: Lazy<[String; 1]> =
Lazy::new(|| [CLOUD_TRACE_CONTEXT_HEADER.to_owned()]);

static GOOGLE_PROPAGATION_HEADER_VALUE_REGEX: Lazy<Option<Regex>> =
Lazy::new(|| Regex::new(GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR).ok());

impl GoogleTraceContextPropagator {
/// Create a new `GoogleTraceContextPropagator`.
pub fn new() -> Self {
GoogleTraceContextPropagator { _private: () }
}

fn extract_span_context(&self, extractor: &dyn Extractor) -> Result<SpanContext, ()> {
let header_value = extractor
.get(CLOUD_TRACE_CONTEXT_HEADER)
.map(|v| v.trim())
.ok_or(())?;

let regex = GOOGLE_PROPAGATION_HEADER_VALUE_REGEX.as_ref().ok_or(())?;

// we could do a quick check if the header value matches the regex here to avoid more expensive regex capture in the next step if it doesn't match
// but the assumption is that the header value will almost always be valid, so that would add an unnecessary check in majority of cases
let caps = regex.captures(header_value).ok_or(())?;

// trace id is mandatory, if it's missing, the header value is invalid
let trace_id_hex = caps.name("trace_id").map(|m| m.as_str()).ok_or(())?;

// span id is mandatory, if it's missing, the header value is invalid
let span_id_dec = caps.name("span_id").map(|m| m.as_str()).ok_or(())?;

// the request is sampled by default, it's not sampled only if explicitly set to 0
let trace_flags = caps.name("trace_flags").map_or(TraceFlags::SAMPLED, |m| {
if m.as_str() == "0" {
TraceFlags::NOT_SAMPLED
} else {
TraceFlags::SAMPLED
}
});

Self::construct_span_context(trace_flags, trace_id_hex, span_id_dec)
}

fn construct_span_context(
trace_flags: TraceFlags,
trace_id_hex: &str,
span_id_dec: &str,
) -> Result<SpanContext, ()> {
let trace_id = TraceId::from_hex(trace_id_hex).map_err(|_| ())?;
let (trace_id, rest) = match header_value.split_once('/') {
Some((trace_id, rest)) if trace_id.len() == 32 => (trace_id, rest),
_ => return Err(()),
};

let span_id = span_id_dec
.parse::<u64>()
.map(|v| SpanId::from_bytes(v.to_be_bytes())) // we can create SPAN ID only from bytes or hex string
.map_err(|_| ())?;
let (span_id, trace_flags) = match rest.split_once(";o=") {
Some((span_id, trace_flags)) => (span_id, trace_flags),
None => (rest, "1"),
};

let trace_id = TraceId::from_hex(trace_id).map_err(|_| ())?;
let span_id = SpanId::from(u64::from_str(span_id).map_err(|_| ())?);
let trace_flags = TraceFlags::new(u8::from_str(trace_flags).map_err(|_| ())?);
let span_context = SpanContext::new(trace_id, span_id, trace_flags, true, TraceState::NONE);

// Ensure span is valid
Expand Down Expand Up @@ -128,29 +96,9 @@ mod tests {
use opentelemetry::trace::TraceState;
use std::collections::HashMap;

#[test]
fn test_google_propagation_header_value_regex_str_valid() {
// Try to create a Regex from the string
let regex_result = Regex::new(GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR);

// Assert that the Regex was created successfully
assert!(
regex_result.is_ok(),
"Failed to create Regex from GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR"
);

// If the Regex was created successfully, validate it against a known valid string
let regex = regex_result.unwrap();
let valid_string = "105445aa7843bc8bf206b12000100000/1;o=1";
assert!(
regex.is_match(valid_string),
"Regex does not match known valid string"
);
}

#[test]
fn test_extract_span_context_valid() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let mut headers = HashMap::new();
headers.insert(
// hashmap implementation of Extractor trait uses lowercase keys
Expand All @@ -169,7 +117,7 @@ mod tests {

#[test]
fn test_extract_span_context_valid_without_options() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let mut headers = HashMap::new();
headers.insert(
// hashmap implementation of Extractor trait uses lowercase keys
Expand All @@ -188,7 +136,7 @@ mod tests {

#[test]
fn test_extract_span_context_valid_not_sampled() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let mut headers = HashMap::new();
headers.insert(
// hashmap implementation of Extractor trait uses lowercase keys
Expand All @@ -207,15 +155,15 @@ mod tests {

#[test]
fn test_extract_span_context_invalid() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let headers = HashMap::new();

assert!(propagator.extract_span_context(&headers).is_err());
}

#[test]
fn test_inject_context_valid() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let mut headers = HashMap::new();
let span = TestSpan(SpanContext::new(
TraceId::from_hex("105445aa7843bc8bf206b12000100000").unwrap(),
Expand All @@ -236,21 +184,25 @@ mod tests {

#[test]
fn test_extract_with_context_valid() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let mut headers = HashMap::new();
headers.insert(
CLOUD_TRACE_CONTEXT_HEADER.to_string().to_lowercase(),
"105445aa7843bc8bf206b12000100000/1;o=1".to_string(),
"105445aa7843bc8bf206b12000100000/10;o=1".to_string(),
);
let cx = Context::current();

let new_cx = propagator.extract_with_context(&cx, &headers);
assert!(new_cx.span().span_context().is_valid());
assert_eq!(
new_cx.span().span_context().span_id().to_string(),
"000000000000000a"
);
}

#[test]
fn test_extract_with_context_invalid_trace_id() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let mut headers = HashMap::new();
// Insert a trace ID with less than 32 characters
headers.insert(
Expand All @@ -266,7 +218,7 @@ mod tests {

#[test]
fn test_extract_with_context_invalid_span_id() {
let propagator = GoogleTraceContextPropagator::new();
let propagator = GoogleTraceContextPropagator::default();
let mut headers = HashMap::new();
// Insert a trace ID with less than 32 characters
headers.insert(
Expand Down

0 comments on commit 93dc0d9

Please sign in to comment.