Add evil trace feature for development

docker/Dockerfile

@@ -17,11 +17,12 @@ ENV RUSTFLAGS "-C link-arg=-s"
 
 RUN update-ca-certificates 2> /dev/null || true
 
+# TODO: do not enable otel-evil-trace for production
 RUN apt-get update && apt-get install -qy --no-install-recommends $BUILD_DEPS && \
     curl -sSf https://sh.rustup.rs | bash -s -- -y --default-toolchain stable && \
     export PATH="$HOME/.cargo/bin:$PATH" && \
     echo "Building Mutualized Oblivious DNS relay and target from source" && \
-    cargo build --release --no-default-features --features=otel-full --package modoh-server && \
+    cargo build --release --no-default-features --features=otel-full,otel-evil-trace --package modoh-server && \
     strip --strip-all /tmp/target/release/modoh-server
 
 ########################################

modoh-bin/Cargo.toml

@@ -32,7 +32,7 @@ publish = false
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-default = ["otel-full"]
+default = ["otel-full", "otel-evil-trace"]
 otel-full = ["otel-trace", "otel-metrics", "otel-instance-id"]
 otel-trace = [
   "opentelemetry/trace",
@@ -56,6 +56,9 @@ otel-base = [
 ]
 otel-instance-id = ["dep:uuid"]
 
+# DO NOT USE THIS IN PRODUCTION
+otel-evil-trace = ["modoh-server-lib/evil-trace"]
+
 [dependencies]
 modoh-server-lib = { path = "../modoh-lib", default-features = false, features = [
   "native-tls",

modoh-lib/Cargo.toml

@@ -34,6 +34,7 @@ publish = false
 [features]
 default = ["native-tls"]
 metrics = ["dep:opentelemetry_sdk", "dep:opentelemetry"]
+evil-trace = ["dep:tracing-opentelemetry"]
 native-tls = ["dep:hyper-tls"]
 rustls = []
 
@@ -102,3 +103,7 @@ opentelemetry_sdk = { version = "0.21.1", features = [
   "metrics",
   "rt-tokio",
 ], optional = true }
+
+# tracing requests traveled among relays and targets
+# NOTE: DO NOT USE THIS IN PRODUCTION
+tracing-opentelemetry = { version = "0.22.0", optional = true }

modoh-lib/src/constants.rs

@@ -53,3 +53,10 @@ pub const JWKS_REFETCH_DELAY_SEC: u64 = 300;
 pub const JWKS_REFETCH_TIMEOUT_SEC: u64 = 3;
 /// Expected maximum size of JWKS in bytes
 pub const EXPECTED_MAX_JWKS_SIZE: u64 = 1024 * 64;
+
+#[cfg(feature = "evil-trace")]
+pub const EVIL_TRACE_HEADER_NAME: &str = "traceparent";
+#[cfg(feature = "evil-trace")]
+pub const EVIL_TRACE_FLAGS: &str = "01";
+#[cfg(feature = "evil-trace")]
+pub const EVIL_TRACE_VERSION: &str = "00";

modoh-lib/src/evil_trace.rs

@@ -0,0 +1,92 @@
+use crate::constants::EVIL_TRACE_HEADER_NAME;
+use anyhow::anyhow;
+use http::Request;
+use hyper::header::HeaderName;
+use opentelemetry::trace::{SpanContext, SpanId, TraceFlags, TraceId};
+
+#[allow(unused)]
+/// Trace context
+pub(crate) struct TraceCX {
+  pub(crate) version: String, // 00
+  pub(crate) trace_id: String,
+  pub(crate) span_id: String,
+  pub(crate) trace_flags: String,
+}
+
+impl TryFrom<String> for TraceCX {
+  type Error = anyhow::Error;
+
+  fn try_from(value: String) -> Result<Self, Self::Error> {
+    let mut iter = value.split('-');
+    let version = iter.next().ok_or_else(|| anyhow!("Invalid trace context"))?;
+    let trace_id = iter.next().ok_or_else(|| anyhow!("Invalid trace context"))?;
+    let span_id = iter.next().ok_or_else(|| anyhow!("Invalid trace context"))?;
+    let trace_flags = iter.next().ok_or_else(|| anyhow!("Invalid trace context"))?;
+    Ok(TraceCX {
+      version: version.to_string(),
+      trace_id: trace_id.to_string(),
+      span_id: span_id.to_string(),
+      trace_flags: trace_flags.to_string(),
+    })
+  }
+}
+
+impl From<TraceCX> for SpanContext {
+  fn from(val: TraceCX) -> Self {
+    let trace_flags = match val.trace_flags.as_str() {
+      "00" => TraceFlags::NOT_SAMPLED,
+      "01" => TraceFlags::SAMPLED,
+      _ => TraceFlags::default(),
+    };
+    SpanContext::new(
+      TraceId::from_hex(&val.trace_id).unwrap(),
+      SpanId::from_hex(&val.span_id).unwrap(),
+      trace_flags,
+      false,
+      opentelemetry::trace::TraceState::default(),
+    )
+  }
+}
+
+/// Get span context built from trace header from http request
+/// Definition of http header for trace
+/// https://uptrace.dev/opentelemetry/opentelemetry-traceparent.html
+/// ```test:
+/// # {version}-{trace_id}-{span_id}-{trace_flags}
+/// traceparent: 00-80e1afed08e019fc1110464cfa66635c-7a085853722dc6d2-01
+/// ```
+pub(crate) fn get_span_cx_from_request<B>(req: &Request<B>) -> Option<SpanContext> {
+  let header_name = HeaderName::from_static(EVIL_TRACE_HEADER_NAME);
+  let header_value = req.headers().get(&header_name)?;
+  let header_value = header_value.to_str().ok()?;
+  TraceCX::try_from(header_value.to_string()).ok().map(SpanContext::from)
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use http::header::HeaderValue;
+  use opentelemetry::trace::TraceState;
+
+  #[test]
+  fn test_get_trace_cx_from_request() {
+    let req = Request::builder()
+      .header(
+        HeaderName::from_static(EVIL_TRACE_HEADER_NAME),
+        HeaderValue::from_static("00-80e1afed08e019fc1110464cfa66635c-7a085853722dc6d2-01"),
+      )
+      .body(http_body_util::Empty::<hyper::body::Bytes>::new())
+      .unwrap();
+    let span_context = get_span_cx_from_request(&req).unwrap();
+    assert_eq!(
+      span_context,
+      SpanContext::new(
+        TraceId::from_hex("80e1afed08e019fc1110464cfa66635c").unwrap(),
+        SpanId::from_hex("7a085853722dc6d2").unwrap(),
+        TraceFlags::SAMPLED,
+        false,
+        TraceState::default(),
+      )
+    );
+  }
+}

modoh-lib/src/hyper_client.rs

@@ -37,6 +37,30 @@ where
     &self,
     req: Request<B>,
   ) -> std::result::Result<Response<Incoming>, hyper_util::client::legacy::Error> {
+    #[cfg(feature = "evil-trace")]
+    {
+      use crate::constants::{EVIL_TRACE_FLAGS, EVIL_TRACE_HEADER_NAME, EVIL_TRACE_VERSION};
+      use opentelemetry::trace::TraceContextExt;
+      use tracing_opentelemetry::OpenTelemetrySpanExt;
+      let current_span_context = tracing::Span::current().context().span().span_context().clone();
+      let header_value = format!(
+        "{}-{}-{}-{}",
+        EVIL_TRACE_VERSION,
+        current_span_context.trace_id(),
+        current_span_context.span_id(),
+        EVIL_TRACE_FLAGS
+      );
+      let mut req = req;
+      let headers = req.headers_mut();
+      headers.insert(
+        http::HeaderName::from_static(EVIL_TRACE_HEADER_NAME),
+        http::HeaderValue::from_str(&header_value).unwrap_or(http::HeaderValue::from_static("")),
+      );
+      println!("evil-trace enabled. header: {:?}", req.headers());
+      self.inner.request(req).await
+    }
+
+    #[cfg(not(feature = "evil-trace"))]
     self.inner.request(req).await
   }
 }

modoh-lib/src/lib.rs

@@ -16,6 +16,9 @@ mod validator;
 #[cfg(feature = "metrics")]
 mod metrics;
 
+#[cfg(feature = "evil-trace")]
+mod evil_trace;
+
 use crate::{count::RequestCount, error::*, globals::Globals, router::Router, trace::*};
 use hyper_client::HttpClient;
 use hyper_executor::LocalExecutor;

modoh-lib/src/router/router_main.rs

@@ -12,6 +12,7 @@ use hyper::{
 use hyper_util::{client::legacy::connect::Connect, rt::TokioIo, server::conn::auto::Builder as ConnectionBuilder};
 use std::{net::SocketAddr, sync::Arc, time::Duration};
 use tokio::time::timeout;
+use tracing::Instrument as _;
 
 #[derive(Clone)]
 /// (M)ODoH Router main object
@@ -61,7 +62,23 @@ where
         timeout_sec + Duration::from_secs(1),
         server_clone.serve_connection(
           stream,
-          service_fn(move |req: Request<Incoming>| serve_request_with_validation(req, peer_addr, self_clone.clone())),
+          service_fn(move |req: Request<Incoming>| {
+            let current_span = tracing::info_span!("router_serve", method = ?req.method(), uri = ?req.uri(), peer_addr = ?peer_addr, xff = ?req.headers().get("x-forwarded-for"), forwarded = ?req.headers().get("forwarded"));
+
+            #[cfg(feature = "evil-trace")]
+            {
+              use opentelemetry::trace::TraceContextExt;
+              use tracing_opentelemetry::OpenTelemetrySpanExt;
+              let span_cx = crate::evil_trace::get_span_cx_from_request(&req);
+              if span_cx.is_some() {
+                debug!("evil-trace enabled. parente span context: {:?}", span_cx);
+                let span_cx = span_cx.unwrap();
+                let context = current_span.context().with_remote_span_context(span_cx);
+                current_span.set_parent(context);
+              }
+            }
+            serve_request_with_validation(req, peer_addr, self_clone.clone()).instrument(current_span)
+    }),
         ),
       )
       .await

modoh-lib/src/router/router_serve_req.rs

@@ -7,9 +7,7 @@ use crate::{
 use hyper::{body::Incoming, header, Request, StatusCode};
 use hyper_util::client::legacy::connect::Connect;
 use std::net::SocketAddr;
-use tracing::instrument;
 
-#[instrument(name = "router_serve", skip_all, fields(method = ?req.method(), uri = ?req.uri(), peer_addr = ?peer_addr, xff = ?req.headers().get("x-forwarded-for"), forwarded = ?req.headers().get("forwarded")))]
 /// Service wrapper with validation
 pub async fn serve_request_with_validation<C>(
   req: Request<Incoming>,