From 0060ed374a83d0d686d8ce97b46f409524376cef Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 12:05:41 -0400 Subject: [PATCH 1/6] Add datadog-mock-intake crate for APM payload-level integration tests In-process mock Datadog intake that accepts the APM endpoints serverless components flush to (/api/v0.2/stats, /api/v0.2/traces, /api/v0.1/pipeline_stats), decodes msgpack/protobuf payloads on arrival, and exposes typed query methods plus raw request capture for test assertions. Ported from the bottlecap fake intake and served with hyper; preserves the 2 MiB wire-body limit, gzip/zstd decoding, stats failure injection, request summaries, and JSON dumps. --- Cargo.lock | 27 +- crates/datadog-mock-intake/Cargo.toml | 30 + crates/datadog-mock-intake/src/lib.rs | 2067 +++++++++++++++++ .../tests/common/mock_server.rs | 157 -- 4 files changed, 2121 insertions(+), 160 deletions(-) create mode 100644 crates/datadog-mock-intake/Cargo.toml create mode 100644 crates/datadog-mock-intake/src/lib.rs delete mode 100644 crates/datadog-trace-agent/tests/common/mock_server.rs diff --git a/Cargo.lock b/Cargo.lock index a8d1b467..f42fb6ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -585,6 +585,27 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "datadog-mock-intake" +version = "0.1.0" +dependencies = [ + "bytes", + "flate2", + "http-body-util", + "hyper", + "hyper-util", + "libdd-trace-protobuf", + "prost 0.14.3", + "reqwest", + "rmp-serde", + "serde", + "serde_json", + "tempfile", + "thiserror 1.0.69", + "tokio", + "zstd", +] + [[package]] name = "datadog-opentelemetry" version = "0.5.2" @@ -664,8 +685,8 @@ dependencies = [ "async-trait", "bytes", "datadog-fips", + "datadog-mock-intake", "duplicate", - "flate2", "http-body-util", "hyper", "hyper-http-proxy", @@ -3015,9 +3036,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.45" +version = "0.23.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +checksum = "6725596c3f2c3a0aef021139e145d4eafe314a6623e4680ca83852b2c67ab2ba" dependencies = [ "aws-lc-rs", "once_cell", diff --git a/crates/datadog-mock-intake/Cargo.toml b/crates/datadog-mock-intake/Cargo.toml new file mode 100644 index 00000000..db1cf336 --- /dev/null +++ b/crates/datadog-mock-intake/Cargo.toml @@ -0,0 +1,30 @@ +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "datadog-mock-intake" +version = "0.1.0" +license.workspace = true +edition.workspace = true + +[lib] +bench = false + +[dependencies] +bytes = "1.10.1" +flate2 = { version = "1", default-features = false, features = ["rust_backend"] } +http-body-util = "0.1" +hyper = { version = "1.6", features = ["http1", "server"] } +hyper-util = { version = "0.1", features = ["tokio"] } +libdd-trace-protobuf = "4.0.1" +prost = "0.14.1" +rmp-serde = { version = "1.3.1", default-features = false } +serde = { version = "1.0.145", features = ["derive"] } +serde_json = "1.0" +thiserror = { version = "1.0.58", default-features = false } +tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "sync", "time"] } +zstd = { version = "0.13.3", default-features = false } + +[dev-dependencies] +reqwest = { version = "0.12.23", default-features = false } +tempfile = "3.3.0" diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs new file mode 100644 index 00000000..8f477e6e --- /dev/null +++ b/crates/datadog-mock-intake/src/lib.rs @@ -0,0 +1,2067 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! In-process mock Datadog intake for APM payload-level integration tests. +//! +//! Spawns a hyper server on a local port that accepts the same APM endpoints +//! serverless components flush to, decodes msgpack / protobuf payloads on +//! arrival, and stores the decoded structs. Callers then use typed query +//! methods to assert on payload contents. +//! +//! Endpoints supported: +//! +//! - `POST /api/v0.2/stats`: msgpack, gzip-compressed, `pb::StatsPayload` +//! - `POST /api/v0.2/traces`: protobuf (optionally zstd-compressed), `pb::AgentPayload` +//! - `POST /api/v0.1/pipeline_stats`: msgpack (struct-as-map), gzip-compressed, DSM pipeline stats +//! +//! Every completed POST attempt to one of these endpoints, including decode +//! failures and failure-injection rejections, is also recorded as a raw +//! [`CapturedRequest`], queryable via [`MockIntake::requests_for_path`]. +//! Typed query methods return only accepted, successfully decoded payloads. +//! +//! `MockIntake::start()` binds `127.0.0.1:0` with defaults and is used by the +//! trace-agent integration tests in this workspace and by the APM / DSM +//! integration tests in `bottlecap`. The standalone `fake-intake` binary in +//! the `datadog-lambda-extension` repository wraps this crate with request +//! summaries, optional stats failure injection +//! (`FAKE_INTAKE_FAIL_STATS_FIRST_N`), and optional JSON dumps +//! (`FAKE_INTAKE_DUMP_DIR`) for local debugging against a live tracer. +//! +//! This crate is self-contained (no host config, logging framework, or trace +//! processing dependencies) so it stays usable from multiple repositories. +//! +//! DSM JSON dumps contain only the fields this fixture decodes +//! ([`PipelineStatsPayload`] below); serde ignores the rest of the wire +//! payload, including the `serde_bytes` latency sketches. + +use std::collections::BTreeMap; +use std::convert::Infallible; +use std::fmt::Write as _; +use std::io::Read; +use std::path::PathBuf; +use std::process; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use bytes::Bytes; +use http_body_util::{BodyExt, Full, Limited}; +use hyper::body::Incoming; +use hyper::header; +use hyper::{Method, Request, Response, StatusCode, service::service_fn}; +use hyper_util::rt::TokioIo; +use libdd_trace_protobuf::pb; +use prost::Message; +use serde::{Deserialize, Serialize}; +use tokio::net::TcpListener; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; + +/// Maximum accepted request body size, matching the 2 MiB default extractor +/// limit the intake was originally served with. Enforced on the wire body, +/// before decompression. +const MAX_BODY_SIZE: usize = 2_097_152; + +/// A DSM pipeline-stats payload as it lands on `/api/v0.1/pipeline_stats`. +/// Only the fields tests assert on are decoded; serde ignores the rest +/// (including the `serde_bytes` latency sketches). JSON dumps therefore +/// contain only these fields, not the full wire payload. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsPayload { + #[serde(rename = "Env")] + pub env: String, + #[serde(rename = "Service")] + pub service: String, + #[serde(rename = "TracerVersion")] + pub tracer_version: String, + #[serde(rename = "Version")] + pub version: String, + #[serde(rename = "Tags")] + pub tags: Vec, + #[serde(rename = "Stats")] + pub stats: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsBucket { + #[serde(rename = "Stats")] + pub stats: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsPoint { + #[serde(rename = "Hash")] + pub hash: u64, + #[serde(rename = "ParentHash")] + pub parent_hash: u64, + #[serde(rename = "EdgeTags")] + pub edge_tags: Vec, +} + +/// Options controlling a mock-intake server. The default +/// (`MockIntakeOptions::default()`) disables all diagnostics. +#[derive(Clone, Debug, Default)] +pub struct MockIntakeOptions { + /// Port to bind on `127.0.0.1`. `0` lets the OS assign a free port; the + /// actual address is available via `MockIntake::base_url()`. + pub port: u16, + /// Emit one summary line per handled request to stderr. + pub request_summaries: bool, + /// Return HTTP 500 for the first N stats request attempts. Attempts are + /// still decoded so they can be summarized and dumped, but rejected + /// payloads are not captured by `stats_payloads()`. + pub fail_stats_first_n: usize, + /// Write one JSON envelope per successfully decoded request attempt + /// (including rejected stats attempts) into this directory. + pub dump_dir: Option, +} + +/// Errors from `MockIntake::start_with_options`. +#[derive(Debug, thiserror::Error)] +pub enum MockIntakeError { + #[error("mock_intake: failed to bind listener on 127.0.0.1:{port}: {source}")] + Bind { + port: u16, + #[source] + source: std::io::Error, + }, + #[error( + "mock_intake: failed to initialize dump directory {}: {source}", + .dir.display() + )] + DumpDir { + dir: PathBuf, + #[source] + source: std::io::Error, + }, +} + +/// A raw request captured by the intake: method, path, headers in arrival +/// order, and the original wire body (pre-decompression). Recorded for every +/// completed POST attempt to a supported endpoint, including attempts that +/// were rejected by failure injection or failed to decode. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CapturedRequest { + pub method: String, + pub path: String, + pub headers: Vec<(String, String)>, + pub body: Vec, +} + +/// Captured, decoded APM payloads for a single test run. +#[derive(Debug, Default)] +struct Captured { + requests: Vec, + stats: Vec, + traces: Vec, + pipeline_stats: Vec, +} + +/// Shared server state. The request handler writes to the mutex; callers read +/// via `MockIntake::stats_payloads()` / `trace_payloads()` and friends. +#[derive(Debug)] +struct SharedState { + captured: Mutex, + options: MockIntakeOptions, + /// Monotonic request identity for summaries and dumps. + request_counter: AtomicU64, + /// Stats request attempts, used for `fail_stats_first_n` injection. + stats_attempts: AtomicU64, +} + +/// The APM endpoints the intake decodes. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Endpoint { + Stats, + Traces, + PipelineStats, +} + +impl Endpoint { + fn from_path(path: &str) -> Option { + match path { + "/api/v0.2/stats" => Some(Self::Stats), + "/api/v0.2/traces" => Some(Self::Traces), + "/api/v0.1/pipeline_stats" => Some(Self::PipelineStats), + _ => None, + } + } +} + +/// A running mock-intake server. Drop shuts it down. +#[derive(Debug)] +pub struct MockIntake { + base_url: String, + state: std::sync::Arc, + shutdown_tx: Option>, + task: Option>, +} + +impl MockIntake { + /// Bind to `127.0.0.1` on an OS-assigned port and start serving. + pub async fn start() -> Self { + Self::start_with_options(MockIntakeOptions::default()) + .await + .expect("mock_intake: failed to start server on 127.0.0.1:0") + } + + /// Start serving with explicit options. Fails on bind errors or when the + /// dump directory cannot be created; nothing is served otherwise. + pub async fn start_with_options(options: MockIntakeOptions) -> Result { + if let Some(dir) = &options.dump_dir { + std::fs::create_dir_all(dir).map_err(|source| MockIntakeError::DumpDir { + dir: dir.clone(), + source, + })?; + } + + let listener = TcpListener::bind(("127.0.0.1", options.port)) + .await + .map_err(|source| MockIntakeError::Bind { + port: options.port, + source, + })?; + let addr = listener + .local_addr() + .map_err(|source| MockIntakeError::Bind { + port: options.port, + source, + })?; + let base_url = format!("http://{addr}"); + + let state = std::sync::Arc::new(SharedState { + captured: Mutex::new(Captured::default()), + options, + request_counter: AtomicU64::new(0), + stats_attempts: AtomicU64::new(0), + }); + + let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>(); + let task_state = std::sync::Arc::clone(&state); + let task = tokio::spawn(async move { + let state = task_state; + loop { + tokio::select! { + result = listener.accept() => { + let (stream, _) = match result { + Ok(conn) => conn, + Err(e) => { + eprintln!("mock_intake: accept error: {e}"); + break; + } + }; + + let io = TokioIo::new(stream); + let state = std::sync::Arc::clone(&state); + tokio::spawn(async move { + let service = service_fn(move |req: Request| { + let state = std::sync::Arc::clone(&state); + async move { handle_request(state, req).await } + }); + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection(io, service) + .await; + }); + } + _ = &mut shutdown_rx => { + break; + } + } + } + }); + + Ok(Self { + base_url, + state, + shutdown_tx: Some(shutdown_tx), + task: Some(task), + }) + } + + /// Full URL for the stats endpoint. + #[must_use] + pub fn stats_url(&self) -> String { + format!("{}/api/v0.2/stats", self.base_url) + } + + /// Full URL for the traces endpoint. + #[must_use] + pub fn traces_url(&self) -> String { + format!("{}/api/v0.2/traces", self.base_url) + } + + /// Full URL for the DSM pipeline-stats endpoint. + #[must_use] + pub fn pipeline_stats_url(&self) -> String { + format!("{}/api/v0.1/pipeline_stats", self.base_url) + } + + /// Base URL (scheme + host + port, no path). Use as the `apm_dd_url` for + /// components that build their own endpoint path (e.g. `DsmProcessor`). + #[must_use] + pub fn base_url(&self) -> String { + self.base_url.clone() + } + + /// Raw requests captured for `path`, in arrival order. Includes attempts + /// that were rejected by failure injection or failed to decode; excludes + /// requests rejected at the transport level (unknown path, unsupported + /// method, or oversized body). + #[must_use] + pub fn requests_for_path(&self, path: &str) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: requests mutex poisoned") + .requests + .iter() + .filter(|req| req.path == path) + .cloned() + .collect() + } + + /// All DSM pipeline-stats payloads captured so far, in arrival order. + #[must_use] + pub fn pipeline_stats_payloads(&self) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: pipeline_stats mutex poisoned") + .pipeline_stats + .clone() + } + + /// All `StatsPayload`s captured so far, in arrival order. Rejected stats + /// attempts (failure injection) are excluded. + #[must_use] + pub fn stats_payloads(&self) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: stats mutex poisoned") + .stats + .clone() + } + + /// All `AgentPayload`s captured so far, in arrival order. + #[must_use] + pub fn trace_payloads(&self) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: traces mutex poisoned") + .traces + .clone() + } +} + +impl Drop for MockIntake { + fn drop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.task.take() { + task.abort(); + } + } +} + +/// Build an empty-bodied response with the given status. +fn response(status: StatusCode) -> Response> { + Response::builder() + .status(status) + .body(Full::new(Bytes::new())) + .expect("mock_intake: static response construction cannot fail") +} + +/// Handle one incoming request: route it, enforce the body size limit, then +/// capture and dispatch it to the endpoint-specific handler. +async fn handle_request( + state: std::sync::Arc, + req: Request, +) -> Result>, Infallible> { + let path = req.uri().path().to_string(); + let Some(endpoint) = Endpoint::from_path(&path) else { + return Ok(response(StatusCode::NOT_FOUND)); + }; + if req.method() != Method::POST { + let mut resp = response(StatusCode::METHOD_NOT_ALLOWED); + resp.headers_mut() + .insert(header::ALLOW, header::HeaderValue::from_static("POST")); + return Ok(resp); + } + + let headers = req.headers().clone(); + let body = match Limited::new(req.into_body(), MAX_BODY_SIZE).collect().await { + Ok(collected) => collected.to_bytes(), + Err(err) => { + if err + .downcast_ref::() + .is_some() + { + return Ok(response(StatusCode::PAYLOAD_TOO_LARGE)); + } + eprintln!("mock_intake: failed to read request body: {err}"); + return Ok(response(StatusCode::BAD_REQUEST)); + } + }; + + let captured = CapturedRequest { + method: Method::POST.as_str().to_string(), + path: path.clone(), + headers: headers + .iter() + .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string())) + .collect(), + body: body.to_vec(), + }; + state + .captured + .lock() + .expect("mock_intake: requests mutex poisoned") + .requests + .push(captured); + + let status = match endpoint { + Endpoint::Stats => handle_stats(&state, &headers, &body), + Endpoint::Traces => handle_traces(&state, &headers, &body), + Endpoint::PipelineStats => handle_pipeline_stats(&state, &headers, &body), + }; + Ok(response(status)) +} + +/// Result of handling one intake request, before summary and dump emission. +struct HandledRequest { + request_id: u64, + status: StatusCode, + decoded: Option, +} + +fn handle_stats( + state: &std::sync::Arc, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> StatusCode { + let request_id = next_request_id(state); + let attempt = state.stats_attempts.fetch_add(1, Ordering::SeqCst) + 1; + let inject_failure = attempt <= state.options.fail_stats_first_n as u64; + + let handled: HandledRequest = match decompress(headers, body) { + Ok(d) => match rmp_serde::from_slice::(&d) { + Ok(payload) => { + let status = if inject_failure { + StatusCode::INTERNAL_SERVER_ERROR + } else { + state + .captured + .lock() + .expect("mock_intake: stats mutex poisoned") + .stats + .push(payload.clone()); + StatusCode::ACCEPTED + }; + HandledRequest { + request_id, + status, + decoded: Some(payload), + } + } + Err(err) => { + eprintln!("mock_intake: failed to decode StatsPayload msgpack: {err}"); + HandledRequest { + request_id, + status: failure_status(inject_failure), + decoded: None, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: failure_status(inject_failure), + decoded: None, + } + } + }; + + let payload_count = handled.decoded.as_ref().map_or(0, |p| p.stats.len()); + + if state.options.request_summaries { + let groups = handled.decoded.as_ref().map(stats_hits_by_key); + log_summary( + handled.request_id, + "/api/v0.2/stats", + headers, + handled.status, + payload_count, + groups.as_ref(), + ); + } + + if let Some(payload) = &handled.decoded { + dump_request( + state, + handled.request_id, + "/api/v0.2/stats", + headers, + handled.status, + serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + ); + } + + handled.status +} + +fn handle_traces( + state: &std::sync::Arc, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> StatusCode { + let request_id = next_request_id(state); + let handled: HandledRequest = match decompress(headers, body) { + Ok(d) => match pb::AgentPayload::decode(d.as_slice()) { + Ok(payload) => { + state + .captured + .lock() + .expect("mock_intake: traces mutex poisoned") + .traces + .push(payload.clone()); + HandledRequest { + request_id, + status: StatusCode::ACCEPTED, + decoded: Some(payload), + } + } + Err(err) => { + eprintln!("mock_intake: failed to decode AgentPayload protobuf: {err}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }; + + // Tracer payload count across the ordinary and indexed collections. + let payload_count = handled + .decoded + .as_ref() + .map_or(0, |p| p.tracer_payloads.len() + p.idx_tracer_payloads.len()); + + if state.options.request_summaries { + log_summary( + handled.request_id, + "/api/v0.2/traces", + headers, + handled.status, + payload_count, + None, + ); + } + + if let Some(payload) = &handled.decoded { + dump_request( + state, + handled.request_id, + "/api/v0.2/traces", + headers, + handled.status, + agent_payload_to_json(payload), + ); + } + + handled.status +} + +fn handle_pipeline_stats( + state: &std::sync::Arc, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> StatusCode { + let request_id = next_request_id(state); + let handled: HandledRequest = match decompress(headers, body) { + Ok(d) => match rmp_serde::from_slice::(&d) { + Ok(payload) => { + state + .captured + .lock() + .expect("mock_intake: pipeline_stats mutex poisoned") + .pipeline_stats + .push(payload.clone()); + HandledRequest { + request_id, + status: StatusCode::ACCEPTED, + decoded: Some(payload), + } + } + Err(err) => { + eprintln!("mock_intake: failed to decode pipeline stats msgpack: {err}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }; + + if state.options.request_summaries { + log_summary( + handled.request_id, + "/api/v0.1/pipeline_stats", + headers, + handled.status, + handled.decoded.as_ref().map_or(0, |_| 1), + None, + ); + } + + if let Some(payload) = &handled.decoded { + dump_request( + state, + handled.request_id, + "/api/v0.1/pipeline_stats", + headers, + handled.status, + serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + ); + } + + handled.status +} + +/// Status used for a stats request that failed to decode or was rejected by +/// failure injection. Decode failures outside the injection window keep the +/// historical `400 Bad Request` behavior. +fn failure_status(inject_failure: bool) -> StatusCode { + if inject_failure { + StatusCode::INTERNAL_SERVER_ERROR + } else { + StatusCode::BAD_REQUEST + } +} + +fn next_request_id(state: &SharedState) -> u64 { + state.request_counter.fetch_add(1, Ordering::SeqCst) + 1 +} + +fn content_encoding(headers: &hyper::HeaderMap) -> String { + headers + .get("content-encoding") + .and_then(|v| v.to_str().ok()) + .map_or_else(|| "identity".to_string(), ToString::to_string) +} + +/// Emit one summary line per handled request. Payload contents stay out of +/// logs; for stats, only aggregated hits per full grouping key are shown. +fn log_summary( + request_id: u64, + endpoint: &str, + headers: &hyper::HeaderMap, + status: StatusCode, + payload_count: usize, + groups: Option<&BTreeMap>, +) { + let mut line = format!( + "mock-intake: request {request_id} | POST {endpoint} | encoding={} | status={} | payloads={payload_count}", + content_encoding(headers), + status.as_u16(), + ); + if let Some(groups) = groups { + let rendered: Vec = groups + .iter() + .map(|(key, hits)| format!("hits={hits} {}", key.render())) + .collect(); + let _ = write!(line, " | groups={} | {}", groups.len(), rendered.join("; ")); + } + eprintln!("{line}"); +} + +/// Serialize a JSON envelope for one decoded request attempt and write it to +/// the dump directory. Diagnostic only: failures are reported but never +/// change the intake response. No capture lock is held during serialization +/// or filesystem I/O. +fn dump_request( + state: &SharedState, + request_id: u64, + endpoint: &str, + headers: &hyper::HeaderMap, + status: StatusCode, + payload: serde_json::Value, +) { + let Some(dump_dir) = &state.options.dump_dir else { + return; + }; + + let envelope = serde_json::json!({ + "request_id": request_id, + "endpoint": endpoint, + "encoding": content_encoding(headers), + "status": status.as_u16(), + "payload": payload, + }); + let Ok(mut json) = serde_json::to_vec_pretty(&envelope) else { + eprintln!("mock_intake: failed to serialize dump envelope for request {request_id}"); + return; + }; + json.push(b'\n'); + + // Collision-resistant filename: slug + request id + pid + nanos. A final + // uniqueness check with create_new guarantees no earlier dump is ever + // overwritten, even across process restarts. + let slug = endpoint.replace('/', "_"); + let base = format!( + "{slug}-{request_id:04}-{}-{}", + process::id(), + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_nanos()) + ); + let mut suffix = 0; + loop { + let path = if suffix == 0 { + dump_dir.join(format!("{base}.json")) + } else { + dump_dir.join(format!("{base}-{suffix}.json")) + }; + match std::fs::OpenOptions::new() + .create_new(true) + .write(true) + .open(&path) + { + Ok(mut file) => { + if let Err(e) = std::io::Write::write_all(&mut file, &json) { + eprintln!("mock_intake: failed to write dump {}: {e}", path.display()); + } + return; + } + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { + suffix += 1; + if suffix > 1000 { + eprintln!( + "mock_intake: giving up on dump filename for request {request_id} after 1000 collisions" + ); + return; + } + } + Err(e) => { + eprintln!("mock_intake: failed to create dump {}: {e}", path.display()); + return; + } + } + } +} + +/// Manual JSON serialization for `pb::AgentPayload`, which does not implement +/// `Serialize` in `libdd-trace-protobuf` 4.0.1. Its tracer-payload collections +/// do, so the fields are embedded directly. Fields mirror the protobuf tags. +fn agent_payload_to_json(payload: &pb::AgentPayload) -> serde_json::Value { + serde_json::json!({ + "host_name": payload.host_name, + "env": payload.env, + "tracer_payloads": payload.tracer_payloads, + "tags": payload.tags, + "agent_version": payload.agent_version, + "target_tps": payload.target_tps, + "error_tps": payload.error_tps, + "rare_sampler_enabled": payload.rare_sampler_enabled, + "idx_tracer_payloads": payload.idx_tracer_payloads, + }) +} + +// --------------------------------------------------------------------------- +// Stats grouping +// --------------------------------------------------------------------------- + +/// Full aggregation-dimension key for one `ClientGroupedStats` entry. +/// +/// Includes every dimension stats aggregation distinguishes: client identity +/// and origin context plus the grouped-stat dimensions. Excludes +/// measurements and delivery metadata (hits, errors, duration, sketches, +/// top-level hits, runtime ID, sequence, tracer version/language) and time +/// bucket fields, so identical dimensions combine across time windows. +/// +/// Tag lists are normalized (sorted) here for deterministic summaries; +/// captured payloads are never modified. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +struct StatsGroupKey { + // Client dimensions (from `ClientStatsPayload`). + client_service: String, + client_env: String, + client_version: String, + client_hostname: String, + client_container_id: String, + client_tags: Vec, + client_git_commit_sha: String, + client_image_tag: String, + client_process_tags: String, + client_process_tags_hash: u64, + client_agent_aggregation: String, + + // Grouped dimensions (from `ClientGroupedStats`). + stats_service: String, + stats_name: String, + stats_resource: String, + stats_type: String, + stats_db_type: String, + stats_http_status_code: u32, + stats_grpc_status_code: String, + stats_http_method: String, + stats_http_endpoint: String, + stats_span_kind: String, + stats_synthetics: bool, + stats_is_trace_root: i32, + stats_service_source: String, + stats_peer_tags: Vec, + stats_span_derived_primary_tags: Vec, + stats_additional_metric_tags: Vec, +} + +impl StatsGroupKey { + fn new(client: &pb::ClientStatsPayload, grouped: &pb::ClientGroupedStats) -> Self { + let sorted = |mut tags: Vec| { + tags.sort(); + tags + }; + Self { + client_service: client.service.clone(), + client_env: client.env.clone(), + client_version: client.version.clone(), + client_hostname: client.hostname.clone(), + client_container_id: client.container_id.clone(), + client_tags: sorted(client.tags.clone()), + client_git_commit_sha: client.git_commit_sha.clone(), + client_image_tag: client.image_tag.clone(), + client_process_tags: client.process_tags.clone(), + client_process_tags_hash: client.process_tags_hash, + client_agent_aggregation: client.agent_aggregation.clone(), + + stats_service: grouped.service.clone(), + stats_name: grouped.name.clone(), + stats_resource: grouped.resource.clone(), + stats_type: grouped.r#type.clone(), + stats_db_type: grouped.db_type.clone(), + stats_http_status_code: grouped.http_status_code, + stats_grpc_status_code: grouped.grpc_status_code.clone(), + stats_http_method: grouped.http_method.clone(), + stats_http_endpoint: grouped.http_endpoint.clone(), + stats_span_kind: grouped.span_kind.clone(), + stats_synthetics: grouped.synthetics, + stats_is_trace_root: grouped.is_trace_root, + stats_service_source: grouped.service_source.clone(), + stats_peer_tags: sorted(grouped.peer_tags.clone()), + stats_span_derived_primary_tags: sorted(grouped.span_derived_primary_tags.clone()), + stats_additional_metric_tags: sorted(grouped.additional_metric_tags.clone()), + } + } + + /// Stable, escaped one-line representation in a fixed field order. + fn render(&self) -> String { + let mut fields: Vec = Vec::new(); + fields.push(format!("client.service={}", escape(&self.client_service))); + fields.push(format!("client.env={}", escape(&self.client_env))); + fields.push(format!("client.version={}", escape(&self.client_version))); + fields.push(format!("client.hostname={}", escape(&self.client_hostname))); + fields.push(format!( + "client.container_id={}", + escape(&self.client_container_id) + )); + fields.push(format!("client.tags={}", render_tags(&self.client_tags))); + fields.push(format!( + "client.git_commit_sha={}", + escape(&self.client_git_commit_sha) + )); + fields.push(format!( + "client.image_tag={}", + escape(&self.client_image_tag) + )); + fields.push(format!( + "client.process_tags={}", + escape(&self.client_process_tags) + )); + fields.push(format!( + "client.process_tags_hash={}", + self.client_process_tags_hash + )); + fields.push(format!( + "client.agent_aggregation={}", + escape(&self.client_agent_aggregation) + )); + + fields.push(format!("stats.service={}", escape(&self.stats_service))); + fields.push(format!("stats.name={}", escape(&self.stats_name))); + fields.push(format!("stats.resource={}", escape(&self.stats_resource))); + fields.push(format!("stats.type={}", escape(&self.stats_type))); + fields.push(format!("stats.db_type={}", escape(&self.stats_db_type))); + fields.push(format!( + "stats.http_status_code={}", + self.stats_http_status_code + )); + fields.push(format!( + "stats.grpc_status_code={}", + escape(&self.stats_grpc_status_code) + )); + fields.push(format!( + "stats.http_method={}", + escape(&self.stats_http_method) + )); + fields.push(format!( + "stats.http_endpoint={}", + escape(&self.stats_http_endpoint) + )); + fields.push(format!("stats.span_kind={}", escape(&self.stats_span_kind))); + fields.push(format!("stats.synthetics={}", self.stats_synthetics)); + fields.push(format!( + "stats.is_trace_root={}", + escape(&trilean_name(self.stats_is_trace_root)) + )); + fields.push(format!( + "stats.service_source={}", + escape(&self.stats_service_source) + )); + fields.push(format!( + "stats.peer_tags={}", + render_tags(&self.stats_peer_tags) + )); + fields.push(format!( + "stats.span_derived_primary_tags={}", + render_tags(&self.stats_span_derived_primary_tags) + )); + fields.push(format!( + "stats.additional_metric_tags={}", + render_tags(&self.stats_additional_metric_tags) + )); + fields.join(" ") + } +} + +/// JSON-escape a value for stable output in summaries. +fn escape(value: &str) -> String { + serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string()) +} + +/// Render a sorted tag list as `["a:1","b:2"]` for summaries. +fn render_tags(tags: &[String]) -> String { + let rendered: Vec = tags.iter().map(|t| escape(t)).collect(); + format!("[{}]", rendered.join(",")) +} + +fn trilean_name(value: i32) -> String { + match pb::Trilean::try_from(value) { + Ok(t) => t.as_str_name().to_string(), + Err(_) => value.to_string(), + } +} + +/// Sum `hits` across client payloads and time buckets, grouped by the full +/// aggregation key. Returned as a `BTreeMap` for deterministic ordering. +fn stats_hits_by_key(payload: &pb::StatsPayload) -> BTreeMap { + let mut groups: BTreeMap = BTreeMap::new(); + for client in &payload.stats { + for bucket in &client.stats { + for grouped in &bucket.stats { + *groups + .entry(StatsGroupKey::new(client, grouped)) + .or_insert(0) += grouped.hits; + } + } + } + groups +} + +// --------------------------------------------------------------------------- +// Body decoding +// --------------------------------------------------------------------------- + +/// Decompress a request body based on its `Content-Encoding` header. +/// Supports `gzip` and `zstd`. An unknown or absent encoding is treated as +/// identity: the body is returned unchanged. +fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, String> { + let encoding = headers + .get("content-encoding") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_ascii_lowercase(); + + match encoding.as_str() { + "gzip" => { + let mut decoder = flate2::read::GzDecoder::new(body.as_ref()); + let mut out = Vec::new(); + decoder + .read_to_end(&mut out) + .map_err(|e| format!("mock_intake: gzip decode failed: {e}"))?; + Ok(out) + } + "zstd" => zstd::stream::decode_all(body.as_ref()) + .map_err(|e| format!("mock_intake: zstd decode failed: {e}")), + _ => { + if !encoding.is_empty() { + eprintln!( + "mock_intake: unrecognized Content-Encoding '{encoding}', treating as identity" + ); + } + Ok(body.to_vec()) + } + } +} + +#[cfg(test)] +mod tests { + + use super::*; + + /// POST a raw body to a path on the intake with an optional + /// Content-Encoding header. Returns the response status code. + async fn post(base_url: &str, path: &str, encoding: Option<&str>, body: Vec) -> StatusCode { + let client = reqwest::Client::new(); + let url = format!("{base_url}{path}"); + let mut request = client.post(&url); + if let Some(enc) = encoding { + request = request.header("content-encoding", enc); + } + request + .header("content-type", "application/msgpack") + .body(body) + .send() + .await + .expect("test: request to mock intake failed") + .status() + } + + fn gzip(data: Vec) -> Vec { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + std::io::Write::write_all(&mut encoder, &data).expect("test: gzip compression failed"); + encoder.finish().expect("test: gzip compression failed") + } + + /// Build a `pb::StatsPayload` with one client payload, one bucket, and the + /// given grouped stats entries. + fn stats_payload( + client: pb::ClientStatsPayload, + start: u64, + grouped: Vec, + ) -> pb::StatsPayload { + pb::StatsPayload { + stats: vec![pb::ClientStatsPayload { + stats: vec![pb::ClientStatsBucket { + start, + duration: 10_000_000_000, + stats: grouped, + agent_time_shift: 0, + }], + ..client + }], + ..pb::StatsPayload::default() + } + } + + fn grouped_stats( + service: &str, + resource: &str, + hits: u64, + overrides: impl FnOnce(&mut pb::ClientGroupedStats), + ) -> pb::ClientGroupedStats { + let mut g = pb::ClientGroupedStats { + service: service.to_string(), + name: "smoke.request".to_string(), + resource: resource.to_string(), + r#type: "web".to_string(), + span_kind: "server".to_string(), + http_status_code: 200, + is_trace_root: pb::Trilean::True as i32, + hits, + ..pb::ClientGroupedStats::default() + }; + overrides(&mut g); + g + } + + fn client_payload(env: &str) -> pb::ClientStatsPayload { + pb::ClientStatsPayload { + service: "mock-intake-smoke".to_string(), + env: env.to_string(), + version: "smoke".to_string(), + ..pb::ClientStatsPayload::default() + } + } + + /// Serialize a stats payload whose wire size is exactly `target` bytes by + /// padding the client hostname. Above the msgpack str16 boundary the + /// encoding grows linearly with the string, so a fixed-point adjustment + /// converges; fails the test if it cannot hit the target exactly. + fn stats_payload_wire_of_size(target: usize) -> Vec { + let base = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let base_len = rmp_serde::to_vec_named(&base) + .expect("test: msgpack encode failed") + .len(); + assert!( + target > base_len, + "test: target {} must exceed base payload size {base_len}", + target + ); + + let mut pad = target - base_len; + for _ in 0..8 { + let mut client = client_payload("local"); + client.hostname = "x".repeat(pad); + let body = rmp_serde::to_vec_named(&stats_payload( + client, + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + match body.len().cmp(&target) { + std::cmp::Ordering::Equal => return body, + std::cmp::Ordering::Greater => { + pad -= body.len() - target; + } + std::cmp::Ordering::Less => { + pad += target - body.len(); + } + } + } + panic!("test: could not build a payload of exactly {target} bytes"); + } + + async fn start_default() -> MockIntake { + MockIntake::start().await + } + + #[tokio::test] + async fn default_startup_binds_port_zero_and_exposes_url_helpers() { + let intake = start_default().await; + assert!(intake.base_url().starts_with("http://127.0.0.1:")); + assert!(intake.stats_url().ends_with("/api/v0.2/stats")); + assert!(intake.traces_url().ends_with("/api/v0.2/traces")); + assert!( + intake + .pipeline_stats_url() + .ends_with("/api/v0.1/pipeline_stats") + ); + } + + #[tokio::test] + async fn instances_are_independent() { + let a = start_default().await; + let b = start_default().await; + assert_ne!(a.base_url(), b.base_url()); + + let payload = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let body = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + assert_eq!( + post(&a.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::ACCEPTED + ); + assert_eq!(a.stats_payloads().len(), 1); + assert!(b.stats_payloads().is_empty()); + + // Dropping `a` must not shut down `b`. + drop(a); + let client = reqwest::Client::new(); + // Any HTTP response proves `b` is still listening; an unmatched route + // answers 404. + let response = client + .get(format!("{}/", b.base_url())) + .send() + .await + .expect("test: intake b should still be listening"); + assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND); + } + + #[tokio::test] + async fn drop_shuts_down_server() { + let intake = start_default().await; + let url = intake.base_url(); + drop(intake); + // The listener socket is released on shutdown; a fresh bind on the + // same port must succeed shortly afterwards. + let port = url.rsplit(':').next().and_then(|p| p.parse::().ok()); + let port = port.expect("test: no port in base url"); + let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(5); + let mut bound = false; + while tokio::time::Instant::now() < deadline { + if TcpListener::bind(("127.0.0.1", port)).await.is_ok() { + bound = true; + break; + } + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + } + assert!(bound, "listener port was not released after Drop"); + } + + #[tokio::test] + async fn stats_endpoint_decodes_msgpack_and_gzip() { + let intake = start_default().await; + let payload = stats_payload( + client_payload("local"), + 42, + vec![grouped_stats("svc", "GET /a", 3, |_| {})], + ); + let raw = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + Some("gzip"), + gzip(raw.clone()) + ) + .await, + StatusCode::ACCEPTED + ); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, raw).await, + StatusCode::ACCEPTED + ); + + let captured = intake.stats_payloads(); + assert_eq!(captured.len(), 2); + assert_eq!(captured[0].stats[0].stats[0].stats[0].hits, 3); + assert_eq!(captured[0].stats[0].stats[0].start, 42); + } + + #[tokio::test] + async fn traces_endpoint_decodes_protobuf() { + let intake = start_default().await; + let payload = pb::AgentPayload { + env: "local".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + ..pb::AgentPayload::default() + }; + let body = payload.encode_to_vec(); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/traces", None, body).await, + StatusCode::ACCEPTED + ); + let captured = intake.trace_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].env, "local"); + assert_eq!(captured[0].tracer_payloads.len(), 1); + } + + #[tokio::test] + async fn traces_endpoint_decodes_zstd_compressed_protobuf() { + let intake = start_default().await; + let payload = pb::AgentPayload { + env: "local".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + ..pb::AgentPayload::default() + }; + let body = zstd::stream::encode_all(payload.encode_to_vec().as_slice(), 0) + .expect("test: zstd compression failed"); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/traces", Some("zstd"), body).await, + StatusCode::ACCEPTED + ); + let captured = intake.trace_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].env, "local"); + } + + #[tokio::test] + async fn dsm_endpoint_decodes_msgpack_gzip() { + let intake = start_default().await; + let payload = PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: vec!["a:b".to_string()], + stats: vec![PipelineStatsBucket { + stats: vec![PipelineStatsPoint { + hash: 7, + parent_hash: 0, + edge_tags: vec!["direction:out".to_string()], + }], + }], + }; + let body = gzip(rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed")); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.1/pipeline_stats", + Some("gzip"), + body + ) + .await, + StatusCode::ACCEPTED + ); + let captured = intake.pipeline_stats_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].stats[0].stats[0].hash, 7); + } + + #[tokio::test] + async fn malformed_bodies_return_400_and_are_not_typed_captured() { + let intake = start_default().await; + for path in [ + "/api/v0.2/stats", + "/api/v0.2/traces", + "/api/v0.1/pipeline_stats", + ] { + assert_eq!( + post(&intake.base_url(), path, None, vec![0xFF; 16]).await, + StatusCode::BAD_REQUEST, + "malformed body on {path} should return 400" + ); + } + assert!(intake.stats_payloads().is_empty()); + assert!(intake.trace_payloads().is_empty()); + assert!(intake.pipeline_stats_payloads().is_empty()); + // Raw capture records completed attempts even when decoding failed. + for path in [ + "/api/v0.2/stats", + "/api/v0.2/traces", + "/api/v0.1/pipeline_stats", + ] { + assert_eq!( + intake.requests_for_path(path).len(), + 1, + "malformed POST on {path} should still be raw-captured" + ); + } + } + + #[tokio::test] + async fn fail_stats_first_n_rejects_exactly_n_then_accepts() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 2, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let payload = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let body = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body.clone()).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body.clone()).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::ACCEPTED + ); + + // Rejected attempts are not typed-captured; only the successful retry is. + let captured = intake.stats_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].stats[0].stats[0].stats[0].hits, 1); + // Raw capture records every completed attempt, including rejections. + assert_eq!(intake.requests_for_path("/api/v0.2/stats").len(), 3); + } + + #[tokio::test] + async fn concurrent_stats_requests_respect_rejection_limit() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 3, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let payload = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let body = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + + let client = reqwest::Client::new(); + let url = format!("{}/api/v0.2/stats", intake.base_url()); + let mut handles = Vec::new(); + for _ in 0..10 { + let client = client.clone(); + let url = url.clone(); + let body = body.clone(); + handles.push(tokio::spawn(async move { + client + .post(&url) + .body(body) + .send() + .await + .expect("test: request failed") + .status() + .as_u16() + })); + } + let mut rejected = 0; + let mut accepted = 0; + for handle in handles { + match handle.await.expect("test: task panicked") { + 500 => rejected += 1, + 202 => accepted += 1, + other => panic!("unexpected status {other}"), + } + } + assert_eq!(rejected, 3, "exactly the first 3 attempts must be rejected"); + assert_eq!(accepted, 7); + assert_eq!(intake.stats_payloads().len(), 7); + } + + #[tokio::test] + async fn traces_and_dsm_do_not_consume_stats_rejections() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 1, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let trace = pb::AgentPayload::default().encode_to_vec(); + let dsm = gzip( + rmp_serde::to_vec_named(&PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: Vec::new(), + stats: Vec::new(), + }) + .expect("test: msgpack encode failed"), + ); + + assert_eq!( + post(&intake.base_url(), "/api/v0.2/traces", None, trace).await, + StatusCode::ACCEPTED + ); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.1/pipeline_stats", + Some("gzip"), + dsm + ) + .await, + StatusCode::ACCEPTED + ); + // The stats rejection budget is untouched by the two requests above. + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + None, + rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})] + )) + .expect("test: msgpack encode failed") + ) + .await, + StatusCode::INTERNAL_SERVER_ERROR + ); + } + + #[tokio::test] + async fn unknown_path_returns_404_and_is_not_captured() { + let intake = start_default().await; + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + assert_eq!( + post(&intake.base_url(), "/api/v2/profile", None, body).await, + StatusCode::NOT_FOUND + ); + assert!(intake.stats_payloads().is_empty()); + assert!(intake.trace_payloads().is_empty()); + assert!(intake.pipeline_stats_payloads().is_empty()); + assert!(intake.requests_for_path("/api/v2/profile").is_empty()); + } + + #[tokio::test] + async fn non_post_methods_on_known_paths_return_405_and_are_not_captured() { + let intake = start_default().await; + let client = reqwest::Client::new(); + for path in [ + "/api/v0.2/stats", + "/api/v0.2/traces", + "/api/v0.1/pipeline_stats", + ] { + let response = client + .get(format!("{}{path}", intake.base_url())) + .send() + .await + .expect("test: request failed"); + assert_eq!(response.status(), reqwest::StatusCode::METHOD_NOT_ALLOWED); + assert!( + response.headers().get("allow").is_some_and(|v| v == "POST"), + "405 response must advertise Allow: POST" + ); + assert!(intake.requests_for_path(path).is_empty()); + } + assert!(intake.stats_payloads().is_empty()); + assert!(intake.trace_payloads().is_empty()); + assert!(intake.pipeline_stats_payloads().is_empty()); + } + + #[tokio::test] + async fn oversized_bodies_return_413_and_are_not_captured_or_counted() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 1, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + // One byte over the wire-body limit. + let oversized = vec![0u8; MAX_BODY_SIZE + 1]; + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, oversized).await, + StatusCode::PAYLOAD_TOO_LARGE + ); + assert!(intake.requests_for_path("/api/v0.2/stats").is_empty()); + assert!(intake.stats_payloads().is_empty()); + + // The rejected oversized attempt must not consume the failure-injection + // budget: the next valid stats request is rejected as attempt 1 of 1. + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert!(intake.stats_payloads().is_empty()); + } + + #[tokio::test] + async fn body_at_exact_size_limit_is_accepted() { + let intake = start_default().await; + let body = stats_payload_wire_of_size(MAX_BODY_SIZE); + assert_eq!(body.len(), MAX_BODY_SIZE); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::ACCEPTED + ); + assert_eq!(intake.stats_payloads().len(), 1); + } + + #[tokio::test] + async fn grouping_combines_across_time_buckets_and_clients() { + let payload = pb::StatsPayload { + stats: vec![ + pb::ClientStatsPayload { + stats: vec![ + pb::ClientStatsBucket { + start: 100, + duration: 10, + stats: vec![ + grouped_stats("svc", "GET /a", 2, |_| {}), + grouped_stats("svc", "GET /b", 1, |_| {}), + ], + agent_time_shift: 0, + }, + pb::ClientStatsBucket { + start: 200, + duration: 10, + stats: vec![grouped_stats("svc", "GET /a", 3, |_| {})], + agent_time_shift: 0, + }, + ], + ..client_payload("prod") + }, + // Same dimensions in a second client payload combine too. + pb::ClientStatsPayload { + stats: vec![pb::ClientStatsBucket { + start: 300, + duration: 10, + stats: vec![grouped_stats("svc", "GET /a", 4, |_| {})], + agent_time_shift: 0, + }], + ..client_payload("prod") + }, + ], + ..pb::StatsPayload::default() + }; + + let groups = stats_hits_by_key(&payload); + assert_eq!(groups.len(), 2, "matching keys must combine across buckets"); + let hits: Vec = groups.values().copied().collect(); + assert_eq!(hits, vec![2 + 3 + 4, 1]); + } + + type PayloadMutation = Box; + + #[test] + #[allow(clippy::too_many_lines)] + fn grouping_separates_changed_dimensions() { + let base = || pb::StatsPayload { + stats: vec![pb::ClientStatsPayload { + stats: vec![pb::ClientStatsBucket { + start: 1, + duration: 1, + stats: vec![grouped_stats("svc", "GET /a", 1, |_| {})], + agent_time_shift: 0, + }], + ..client_payload("prod") + }], + ..pb::StatsPayload::default() + }; + + let mut dimensions: Vec<(&str, PayloadMutation)> = Vec::new(); + dimensions.push(( + "resource", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].resource = "GET /b".to_string(); + }), + )); + dimensions.push(( + "peer_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].peer_tags = vec!["grpc.target:other".to_string()]; + }), + )); + dimensions.push(( + "additional_metric_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].additional_metric_tags = vec!["x:y".to_string()]; + }), + )); + dimensions.push(( + "is_trace_root", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].is_trace_root = pb::Trilean::False as i32; + }), + )); + dimensions.push(( + "http_status_code", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].http_status_code = 500; + }), + )); + dimensions.push(( + "grpc_status_code", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].grpc_status_code = "14".to_string(); + }), + )); + dimensions.push(( + "client_env", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].env = "staging".to_string(); + }), + )); + dimensions.push(( + "client_version", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].version = "2.0".to_string(); + }), + )); + dimensions.push(( + "span_kind", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].span_kind = "client".to_string(); + }), + )); + dimensions.push(( + "http_method", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].http_method = "POST".to_string(); + }), + )); + dimensions.push(( + "http_endpoint", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].http_endpoint = "/api".to_string(); + }), + )); + dimensions.push(( + "service_source", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].service_source = "dd.trace".to_string(); + }), + )); + dimensions.push(( + "span_derived_primary_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].span_derived_primary_tags = vec!["t:1".to_string()]; + }), + )); + dimensions.push(( + "db_type", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].db_type = "postgres".to_string(); + }), + )); + dimensions.push(( + "synthetics", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].synthetics = true; + }), + )); + dimensions.push(( + "name", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].name = "other.request".to_string(); + }), + )); + dimensions.push(( + "type", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].r#type = "db".to_string(); + }), + )); + dimensions.push(( + "client_hostname", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].hostname = "host-1".to_string(); + }), + )); + dimensions.push(( + "client_container_id", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].container_id = "abc".to_string(); + }), + )); + dimensions.push(( + "client_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].tags = vec!["k:v".to_string()]; + }), + )); + dimensions.push(( + "client_git_commit_sha", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].git_commit_sha = "deadbeef".to_string(); + }), + )); + dimensions.push(( + "client_image_tag", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].image_tag = "v1".to_string(); + }), + )); + dimensions.push(( + "client_process_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].process_tags = "proc:a".to_string(); + }), + )); + dimensions.push(( + "client_process_tags_hash", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].process_tags_hash = 99; + }), + )); + dimensions.push(( + "client_agent_aggregation", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].agent_aggregation = "counts".to_string(); + }), + )); + dimensions.push(( + "stats_service", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].service = "other-svc".to_string(); + }), + )); + + let baseline = stats_hits_by_key(&base()); + assert_eq!(baseline.len(), 1); + for (name, mutate) in dimensions { + let mut payload = base(); + mutate(&mut payload); + let groups = stats_hits_by_key(&payload); + assert_eq!( + groups.len(), + 1, + "{name} change must not merge with baseline" + ); + assert_ne!( + groups.keys().next(), + baseline.keys().next(), + "{name} must change the grouping key" + ); + } + } + + #[test] + fn grouping_ignores_measurements_and_delivery_metadata() { + let base = grouped_stats("svc", "GET /a", 1, |_| {}); + let variant = grouped_stats("svc", "GET /a", 7, |g| { + g.errors = 2; + g.duration = 500; + g.top_level_hits = 3; + g.ok_summary = vec![1, 2, 3]; + g.error_summary = vec![4, 5]; + }); + let client_a = client_payload("prod"); + let mut client_b = client_payload("prod"); + client_b.runtime_id = "other".to_string(); + client_b.sequence = 12; + client_b.lang = "rust".to_string(); + client_b.tracer_version = "9.9".to_string(); + + let key_base = StatsGroupKey::new(&client_a, &base); + let key_variant = StatsGroupKey::new(&client_b, &variant); + assert_eq!( + key_base, key_variant, + "measurement or delivery changes must not split groups" + ); + } + + #[test] + fn summary_output_is_deterministic_and_normalizes_tag_order() { + let mut client = client_payload("prod"); + client.tags = vec!["z:1".to_string(), "a:2".to_string()]; + let mut grouped = grouped_stats("svc", "GET /a", 5, |_| {}); + grouped.peer_tags = vec!["b:x".to_string(), "a:y".to_string()]; + + let key = StatsGroupKey::new(&client, &grouped); + let first = key.render(); + let second = StatsGroupKey::new(&client, &grouped).render(); + assert_eq!(first, second, "rendering must be deterministic"); + assert!( + first.contains("client.tags=[\"a:2\",\"z:1\"]"), + "tag lists must be sorted: {first}" + ); + assert!( + first.contains("stats.peer_tags=[\"a:y\",\"b:x\"]"), + "tag lists must be sorted: {first}" + ); + assert!(!first.contains("hits="), "summary key must exclude hits"); + assert!(first.starts_with("client.service=")); + } + + #[tokio::test] + async fn json_dumps_parse_for_all_endpoints_and_rejections() { + let tmp = tempfile::tempdir().expect("test: tempdir failed"); + let intake = MockIntake::start_with_options(MockIntakeOptions { + dump_dir: Some(tmp.path().to_path_buf()), + fail_stats_first_n: 1, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let stats = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 2, |_| {})], + ); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + None, + rmp_serde::to_vec_named(&stats).expect("test: msgpack encode failed") + ) + .await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + None, + rmp_serde::to_vec_named(&stats).expect("test: msgpack encode failed") + ) + .await, + StatusCode::ACCEPTED + ); + + let trace = pb::AgentPayload { + env: "local".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + ..pb::AgentPayload::default() + }; + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/traces", + None, + trace.encode_to_vec() + ) + .await, + StatusCode::ACCEPTED + ); + + let dsm = PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: Vec::new(), + stats: Vec::new(), + }; + assert_eq!( + post( + &intake.base_url(), + "/api/v0.1/pipeline_stats", + None, + rmp_serde::to_vec_named(&dsm).expect("test: msgpack encode failed") + ) + .await, + StatusCode::ACCEPTED + ); + + let mut files: Vec<_> = std::fs::read_dir(tmp.path()) + .expect("test: read dump dir failed") + .filter_map(std::result::Result::ok) + .map(|e| e.path()) + .collect(); + files.sort(); + assert_eq!(files.len(), 4, "one dump per decoded attempt: {files:?}"); + + let mut seen = std::collections::HashSet::new(); + for path in &files { + let raw = std::fs::read(path).expect("test: read dump failed"); + let envelope: serde_json::Value = + serde_json::from_slice(&raw).expect("dump envelope must parse as JSON"); + let dumped_endpoint = envelope["endpoint"].as_str().expect("endpoint key missing"); + let dumped_status = envelope["status"].as_u64().expect("status key missing"); + assert!( + envelope["payload"].is_object(), + "decoded payload must be present" + ); + assert!( + envelope["request_id"].is_u64(), + "request id must be present" + ); + assert!( + envelope["encoding"].as_str() == Some("identity"), + "absent content-encoding must report identity" + ); + seen.insert(format!("{dumped_endpoint}/{dumped_status}")); + } + // Rejected stats attempt (500) and its successful retry (202) are + // both visible with their own identities and statuses. + assert!(seen.contains("/api/v0.2/stats/500")); + assert!(seen.contains("/api/v0.2/stats/202")); + assert!(seen.contains("/api/v0.2/traces/202")); + assert!(seen.contains("/api/v0.1/pipeline_stats/202")); + } + + #[tokio::test] + async fn dumps_never_overwrite_earlier_attempts() { + let tmp = tempfile::tempdir().expect("test: tempdir failed"); + let intake = MockIntake::start_with_options(MockIntakeOptions { + dump_dir: Some(tmp.path().to_path_buf()), + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + for _ in 0..5 { + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body.clone()).await, + StatusCode::ACCEPTED + ); + } + + let files: Vec<_> = std::fs::read_dir(tmp.path()) + .expect("test: read dump dir failed") + .filter_map(std::result::Result::ok) + .map(|e| e.path()) + .collect(); + assert_eq!(files.len(), 5, "every attempt gets its own file"); + let mut contents: Vec = files + .iter() + .map(|p| std::fs::read_to_string(p).expect("test: read dump failed")) + .collect(); + contents.sort(); + let unique: std::collections::HashSet = contents.iter().cloned().collect(); + assert_eq!( + unique.len(), + 5, + "each dump must have distinct request identity" + ); + } + + #[tokio::test] + async fn invalid_options_fail_clearly() { + // A dump directory whose parent is a file cannot be created. + let file = tempfile::NamedTempFile::new().expect("test: tempfile failed"); + let err = MockIntake::start_with_options(MockIntakeOptions { + dump_dir: Some(file.path().join("nested")), + ..MockIntakeOptions::default() + }) + .await + .expect_err("dump dir under a file must fail"); + assert!(matches!(err, MockIntakeError::DumpDir { .. })); + + // Binding to a port already in use must fail with a bind error. + let listener = TcpListener::bind(("127.0.0.1", 0)) + .await + .expect("test: bind failed"); + let port = listener + .local_addr() + .expect("test: local addr failed") + .port(); + let err = MockIntake::start_with_options(MockIntakeOptions { + port, + ..MockIntakeOptions::default() + }) + .await + .expect_err("bind conflict must fail"); + assert!(matches!(err, MockIntakeError::Bind { .. })); + } + + #[tokio::test] + async fn start_with_options_uses_configured_port() { + let listener = TcpListener::bind(("127.0.0.1", 0)) + .await + .expect("test: bind failed"); + let port = listener + .local_addr() + .expect("test: local addr failed") + .port(); + drop(listener); + + let intake = MockIntake::start_with_options(MockIntakeOptions { + port, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + assert_eq!(intake.base_url(), format!("http://127.0.0.1:{port}")); + } +} diff --git a/crates/datadog-trace-agent/tests/common/mock_server.rs b/crates/datadog-trace-agent/tests/common/mock_server.rs deleted file mode 100644 index cd0cd6b9..00000000 --- a/crates/datadog-trace-agent/tests/common/mock_server.rs +++ /dev/null @@ -1,157 +0,0 @@ -// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/ -// SPDX-License-Identifier: Apache-2.0 - -//! Simple mock HTTP server for testing flushers - -use http_body_util::BodyExt; -use hyper::{Request, Response, StatusCode, body::Incoming}; -use hyper_util::rt::TokioIo; -use libdd_common::http_common; -use std::net::SocketAddr; -use std::sync::{Arc, Mutex}; -use tokio::net::TcpListener; - -#[derive(Clone, Debug)] -pub struct ReceivedRequest { - pub method: String, - pub path: String, - pub headers: Vec<(String, String)>, - pub body: Vec, -} - -pub struct MockServer { - pub addr: SocketAddr, - pub received_requests: Arc>>, - shutdown_tx: Option>, -} - -impl MockServer { - /// Start a mock HTTP server on a random port - pub async fn start() -> Self { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("Failed to bind mock server"); - let addr = listener.local_addr().expect("Failed to get local addr"); - - let received_requests = Arc::new(Mutex::new(Vec::new())); - let requests_clone = received_requests.clone(); - - let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel(); - - tokio::spawn(async move { - loop { - tokio::select! { - result = listener.accept() => { - let (stream, _) = match result { - Ok(conn) => conn, - Err(e) => { - eprintln!("Mock server accept error: {}", e); - break; - } - }; - - let io = TokioIo::new(stream); - let requests = requests_clone.clone(); - - tokio::spawn(async move { - let service = hyper::service::service_fn(move |req: Request| { - let requests = requests.clone(); - async move { - // Capture the request - let method = req.method().to_string(); - let path = req.uri().path().to_string(); - let is_stats_intake = path.ends_with("/stats"); - let headers: Vec<(String, String)> = req - .headers() - .iter() - .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())) - .collect(); - - // Read the body - let body_bytes = req - .into_body() - .collect() - .await - .map(|collected| collected.to_bytes().to_vec()) - .unwrap_or_default(); - - // Store the request - requests.lock().unwrap().push(ReceivedRequest { - method, - path, - headers, - body: body_bytes, - }); - - // Trace intake accepts 2xx - // Stats intake accepts 202 - // see `libdd_trace_utils::stats_utils::send_stats_payload_with_client` - let (status, body) = if is_stats_intake { - (StatusCode::ACCEPTED, http_common::Body::empty()) - } else { - (StatusCode::OK, http_common::Body::from(r#"{"ok":true}"#)) - }; - Ok::<_, hyper::http::Error>( - Response::builder() - .status(status) - .body(body) - .unwrap(), - ) - } - }); - - let _ = hyper::server::conn::http1::Builder::new() - .serve_connection(io, service) - .await; - }); - } - _ = &mut shutdown_rx => { - break; - } - } - } - }); - - MockServer { - addr, - received_requests, - shutdown_tx: Some(shutdown_tx), - } - } - - /// Get the base URL of the mock server - pub fn url(&self) -> String { - format!("http://{}", self.addr) - } - - /// Get all received requests - #[allow(dead_code)] - pub fn get_requests(&self) -> Vec { - self.received_requests.lock().unwrap().clone() - } - - /// Get requests matching a path - pub fn get_requests_for_path(&self, path: &str) -> Vec { - self.received_requests - .lock() - .unwrap() - .iter() - .filter(|req| req.path == path) - .cloned() - .collect() - } - - /// Clear all received requests - #[allow(dead_code)] - pub fn clear_requests(&self) { - self.received_requests.lock().unwrap().clear(); - } -} - -impl Drop for MockServer { - fn drop(&mut self) { - if let Some(shutdown_tx) = self.shutdown_tx.take() { - let _ = shutdown_tx.send(()); - } - } -} From c5f424b49cc6307f4204e016ad57d9973bdba243 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 12:07:55 -0400 Subject: [PATCH 2/6] Serve trace-agent integration tests with datadog-mock-intake Replace the test-local MockServer with the shared datadog-mock-intake crate. Payload assertions now use the intake's decoded capture instead of re-decoding captured bytes, positive assertions wait on decoded payload arrival, and the DSM forwarding test sends a valid msgpack pipeline-stats payload so it verifies both unchanged byte forwarding and intake acceptance. --- crates/datadog-trace-agent/Cargo.toml | 2 +- .../tests/common/helpers.rs | 13 - .../datadog-trace-agent/tests/common/mod.rs | 1 - .../tests/integration_test.rs | 239 ++++++++++-------- 4 files changed, 133 insertions(+), 122 deletions(-) diff --git a/crates/datadog-trace-agent/Cargo.toml b/crates/datadog-trace-agent/Cargo.toml index 291ff87f..4e868937 100644 --- a/crates/datadog-trace-agent/Cargo.toml +++ b/crates/datadog-trace-agent/Cargo.toml @@ -42,7 +42,7 @@ bytes = "1.10.1" prost = "0.14.1" [dev-dependencies] -flate2 = "1" +datadog-mock-intake = { path = "../datadog-mock-intake" } rmp-serde = "1.1.1" serial_test = "2.0.0" duplicate = "2.0.1" diff --git a/crates/datadog-trace-agent/tests/common/helpers.rs b/crates/datadog-trace-agent/tests/common/helpers.rs index 32cd4736..8650cb48 100644 --- a/crates/datadog-trace-agent/tests/common/helpers.rs +++ b/crates/datadog-trace-agent/tests/common/helpers.rs @@ -3,14 +3,12 @@ //! Helper functions for integration tests -use flate2::read::GzDecoder; use hyper::{Request, Response}; use hyper_util::rt::TokioIo; use libdd_common::http_common; use libdd_trace_protobuf::pb; use libdd_trace_utils::test_utils::create_test_json_span; use serde_json::json; -use std::io::Read; use std::time::{Duration, UNIX_EPOCH}; use tokio::time::timeout; @@ -70,17 +68,6 @@ pub fn create_test_client_stats_payload(service: &str) -> Vec { rmp_serde::to_vec_named(&payload).expect("Failed to serialize client stats payload") } -/// Decompress a gzip+msgpack stats payload and deserialize it into a `StatsPayload`. -/// The stats flusher encodes payloads as `gzip(rmp_serde::to_vec_named(StatsPayload))`. -pub fn decode_stats_payload(body: &[u8]) -> pb::StatsPayload { - let mut decoder = GzDecoder::new(body); - let mut decompressed = Vec::new(); - decoder - .read_to_end(&mut decompressed) - .expect("Failed to decompress stats payload"); - rmp_serde::from_slice(&decompressed).expect("Failed to deserialize stats payload") -} - /// Send an HTTP request over TCP and return the response pub async fn send_tcp_request( port: u16, diff --git a/crates/datadog-trace-agent/tests/common/mod.rs b/crates/datadog-trace-agent/tests/common/mod.rs index 447946a7..3d83bbdb 100644 --- a/crates/datadog-trace-agent/tests/common/mod.rs +++ b/crates/datadog-trace-agent/tests/common/mod.rs @@ -4,5 +4,4 @@ //! Common test utilities, mocks, and helpers for integration tests pub mod helpers; -pub mod mock_server; pub mod mocks; diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index cb0117ba..0f5722ef 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -5,11 +5,10 @@ mod common; use common::helpers::{ create_client_span_with_peer_tag_payload, create_test_client_stats_payload, - create_test_trace_payload, create_trace_with_span_kind_children_payload, decode_stats_payload, - send_tcp_request, + create_test_trace_payload, create_trace_with_span_kind_children_payload, send_tcp_request, }; -use common::mock_server::MockServer; use common::mocks::{MockEnvVerifier, MockStatsFlusher, MockStatsProcessor, MockTraceFlusher}; +use datadog_mock_intake::{MockIntake, PipelineStatsPayload}; use datadog_trace_agent::{ config::{Config, Tags, test_helpers::create_tcp_test_config}, mini_agent::MiniAgent, @@ -39,15 +38,15 @@ const FLUSH_WAIT_DURATION: Duration = Duration::from_millis(1500); // the timeout exists to fail loudly rather than hang indefinitely. const VERIFY_REQUEST_TIMEOUT: Duration = Duration::from_secs(60); -async fn wait_for_request_at_path(mock_server: &common::mock_server::MockServer, path: &str) { +async fn wait_for_capture(ready: impl Fn() -> bool, what: &str) { let deadline = tokio::time::Instant::now() + VERIFY_REQUEST_TIMEOUT; while tokio::time::Instant::now() < deadline { - if !mock_server.get_requests_for_path(path).is_empty() { + if ready() { return; } tokio::time::sleep(Duration::from_millis(50)).await; } - panic!("Timed out after {VERIFY_REQUEST_TIMEOUT:?} waiting for request at {path}"); + panic!("Timed out after {VERIFY_REQUEST_TIMEOUT:?} waiting for {what}"); } /// Helper to configure a config with mock server endpoints @@ -114,14 +113,18 @@ pub fn create_mini_agent_with_real_flushers( (mini_agent, stats_concentrator_service_handle) } -/// Helper to verify trace request sent to mock server -pub async fn verify_trace_request(mock_server: &common::mock_server::MockServer) { - wait_for_request_at_path(mock_server, "/api/v0.2/traces").await; - let trace_reqs = mock_server.get_requests_for_path("/api/v0.2/traces"); +/// Helper to verify trace request sent to the mock intake +pub async fn verify_trace_request(mock_intake: &MockIntake) { + wait_for_capture( + || !mock_intake.trace_payloads().is_empty(), + "a trace request at /api/v0.2/traces", + ) + .await; + let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); assert!( !trace_reqs.is_empty(), - "Expected at least one trace request to mock server" + "Expected at least one trace request to mock intake" ); let trace_req = &trace_reqs[0]; @@ -151,14 +154,18 @@ pub async fn verify_trace_request(mock_server: &common::mock_server::MockServer) ); } -/// Helper to verify stats request sent to mock server -pub async fn verify_stats_request(mock_server: &common::mock_server::MockServer) { - wait_for_request_at_path(mock_server, "/api/v0.2/stats").await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); +/// Helper to verify stats request sent to the mock intake +pub async fn verify_stats_request(mock_intake: &MockIntake) { + wait_for_capture( + || !mock_intake.stats_payloads().is_empty(), + "a stats request at /api/v0.2/stats", + ) + .await; + let stats_reqs = mock_intake.requests_for_path("/api/v0.2/stats"); assert!( !stats_reqs.is_empty(), - "Expected at least one stats request to mock server" + "Expected at least one stats request to mock intake" ); let stats_req = &stats_reqs[0]; @@ -188,28 +195,35 @@ pub async fn verify_stats_request(mock_server: &common::mock_server::MockServer) ); } -/// Helper to verify stats request was not sent to mock server -pub fn verify_no_stats_request(mock_server: &common::mock_server::MockServer) { - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); +/// Helper to verify stats request was not sent to mock intake +pub fn verify_no_stats_request(mock_intake: &MockIntake) { + let stats_payloads = mock_intake.stats_payloads(); assert!( - stats_reqs.is_empty(), - "Expected no stats request to mock server, received {} request(s)", - stats_reqs.len() + stats_payloads.is_empty(), + "Expected no stats request to mock intake, received {} request(s)", + stats_payloads.len() ); } -/// Helper to verify a DSM request sent to the mock server +/// Helper to verify a DSM request sent to the mock intake. The forwarded body +/// must arrive byte-for-byte unchanged AND be decodable by the intake: the +/// decoded capture proves the proxy forwards a payload the real intake accepts. pub async fn verify_dsm_request( - mock_server: &common::mock_server::MockServer, + mock_intake: &MockIntake, expected_body: &[u8], + expected_edge_tags: &[&str], expected_additional_tags: &[&str], ) { - wait_for_request_at_path(mock_server, "/api/v0.1/pipeline_stats").await; - let dsm_reqs = mock_server.get_requests_for_path("/api/v0.1/pipeline_stats"); + wait_for_capture( + || !mock_intake.pipeline_stats_payloads().is_empty(), + "a DSM request at /api/v0.1/pipeline_stats", + ) + .await; + let dsm_reqs = mock_intake.requests_for_path("/api/v0.1/pipeline_stats"); assert!( !dsm_reqs.is_empty(), - "Expected at least one DSM request to mock server" + "Expected at least one DSM request to mock intake" ); let dsm_req = &dsm_reqs[0]; @@ -219,6 +233,17 @@ pub async fn verify_dsm_request( "Expected DSM payload body to be forwarded unchanged" ); + let decoded = &mock_intake.pipeline_stats_payloads()[0]; + let edge_tags: Vec<&str> = decoded.stats[0].stats[0] + .edge_tags + .iter() + .map(String::as_str) + .collect(); + assert_eq!( + edge_tags, expected_edge_tags, + "Expected decoded DSM payload to carry the fixture edge tags" + ); + let api_key = dsm_req .headers .iter() @@ -469,11 +494,10 @@ async fn test_mini_agent_named_pipe_handles_requests() { #[tokio::test] #[serial] async fn test_mini_agent_tcp_with_real_flushers() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8127); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -512,23 +536,22 @@ async fn test_mini_agent_tcp_with_real_flushers() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; + verify_trace_request(&mock_intake).await; // Trigger shutdown to force flush in progress concentrator buckets let _ = shutdown_tx.send(true); let _ = agent_handle.await; - verify_stats_request(&mock_server).await; // Stats generator should generate stats from trace payload + verify_stats_request(&mock_intake).await; // Stats generator should generate stats from trace payload } #[cfg(test)] #[tokio::test] #[serial] async fn test_mini_agent_tcp_proxies_dsm_requests() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8133); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.tags = Tags::from_env_string("env:test,service:payments"); let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -563,7 +586,23 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { "Mini agent server failed to start within timeout" ); - let dsm_payload = br#"{"Stats":[{"EdgeTags":["direction:out","type:kafka"]}]}"#.to_vec(); + // A valid DSM pipeline-stats payload: the proxy must forward it unchanged + // and the intake must accept and decode it. + let dsm_payload = rmp_serde::to_vec_named(&PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: Vec::new(), + stats: vec![datadog_mock_intake::PipelineStatsBucket { + stats: vec![datadog_mock_intake::PipelineStatsPoint { + hash: 7, + parent_hash: 0, + edge_tags: vec!["direction:out".to_string(), "type:kafka".to_string()], + }], + }], + }) + .expect("Failed to serialize DSM pipeline stats payload"); let response = send_tcp_request( test_port, "/v0.1/pipeline_stats", @@ -582,8 +621,9 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { assert_eq!(response.status(), StatusCode::OK); verify_dsm_request( - &mock_server, + &mock_intake, &dsm_payload, + &["direction:out", "type:kafka"], &[ "host:worker-1", "default_env:prod", @@ -602,11 +642,10 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { #[tokio::test] #[serial] async fn test_concentrator_task_death_shuts_down_mini_agent() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8129); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -655,11 +694,10 @@ async fn test_concentrator_task_death_shuts_down_mini_agent() { #[tokio::test] #[serial] async fn test_tracer_and_agent_stats_disabled_produces_no_stats() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8128); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = false; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -696,10 +734,10 @@ async fn test_tracer_and_agent_stats_disabled_produces_no_stats() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; + verify_trace_request(&mock_intake).await; // Bounded wait to confirm absence of stats request — neither side computed stats. tokio::time::sleep(FLUSH_WAIT_DURATION).await; - verify_no_stats_request(&mock_server); + verify_no_stats_request(&mock_intake); // Clean up agent_handle.abort(); @@ -709,11 +747,10 @@ async fn test_tracer_and_agent_stats_disabled_produces_no_stats() { #[tokio::test] #[serial] async fn test_tracer_disabled_agent_stats_enabled_uses_agent_stats() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8136); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -752,8 +789,8 @@ async fn test_tracer_disabled_agent_stats_enabled_uses_agent_stats() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; - verify_stats_request(&mock_server).await; // Agent should compute stats + verify_trace_request(&mock_intake).await; + verify_stats_request(&mock_intake).await; // Agent should compute stats // Clean up agent_handle.abort(); @@ -762,11 +799,10 @@ async fn test_tracer_disabled_agent_stats_enabled_uses_agent_stats() { #[tokio::test] #[serial] async fn test_tracer_and_agent_stats_enabled_uses_agent_stats_no_duplicates() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8135); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -827,20 +863,16 @@ async fn test_tracer_and_agent_stats_enabled_uses_agent_stats_no_duplicates() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; - verify_stats_request(&mock_server).await; + verify_trace_request(&mock_intake).await; + verify_stats_request(&mock_intake).await; // The tracer computed stats should never reach the backend. Only the agent computed stats should reach the backend. - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); - let has_marker = stats_reqs - .iter() - .map(|req| decode_stats_payload(&req.body)) - .any(|payload| { - payload - .stats - .iter() - .any(|csp| csp.service == "tracer-marker-stats") - }); + let has_marker = mock_intake.stats_payloads().iter().any(|payload| { + payload + .stats + .iter() + .any(|csp| csp.service == "tracer-marker-stats") + }); assert!( !has_marker, "Expected tracer computed stats to be dropped, not forwarded to the backend" @@ -853,11 +885,10 @@ async fn test_tracer_and_agent_stats_enabled_uses_agent_stats_no_duplicates() { #[tokio::test] #[serial] async fn test_tracer_stats_enabled_agent_stats_disabled_forwards_tracer_stats() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8134); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = false; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -898,17 +929,13 @@ async fn test_tracer_stats_enabled_agent_stats_disabled_forwards_tracer_stats() .expect("Failed to send /v0.6/stats request"); assert_eq!(stats_response.status(), StatusCode::ACCEPTED); - verify_stats_request(&mock_server).await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); - let has_marker = stats_reqs - .iter() - .map(|req| decode_stats_payload(&req.body)) - .any(|payload| { - payload - .stats - .iter() - .any(|csp| csp.service == "tracer-marker-stats") - }); + verify_stats_request(&mock_intake).await; + let has_marker = mock_intake.stats_payloads().iter().any(|payload| { + payload + .stats + .iter() + .any(|csp| csp.service == "tracer-marker-stats") + }); assert!( has_marker, "Expected tracer computed stats to be forwarded to the backend" @@ -924,11 +951,10 @@ async fn test_tracer_stats_enabled_agent_stats_disabled_forwards_tracer_stats() #[tokio::test] #[serial] async fn test_internal_span_kind_does_not_produce_stats() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8132); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -974,22 +1000,20 @@ async fn test_internal_span_kind_does_not_produce_stats() { let _ = shutdown_tx.send(true); let _ = agent_handle.await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); + let stats_payloads = mock_intake.stats_payloads(); assert!( - !stats_reqs.is_empty(), + !stats_payloads.is_empty(), "Expected a stats request from the root span" ); - let all_groups: Vec<_> = stats_reqs + let all_groups: Vec<_> = stats_payloads .iter() - .map(|req| decode_stats_payload(&req.body)) .flat_map(|payload| { payload .stats - .into_iter() - .flat_map(|csp| csp.stats.into_iter()) - .flat_map(|bucket| bucket.stats.into_iter()) - .collect::>() + .iter() + .flat_map(|csp| csp.stats.iter()) + .flat_map(|bucket| bucket.stats.iter()) }) .collect(); @@ -1015,11 +1039,10 @@ async fn test_internal_span_kind_does_not_produce_stats() { #[tokio::test] #[serial] async fn test_peer_tags_in_flushed_stats() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8131); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -1065,13 +1088,13 @@ async fn test_peer_tags_in_flushed_stats() { let _ = shutdown_tx.send(true); let _ = agent_handle.await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); + let stats_payloads = mock_intake.stats_payloads(); assert!( - !stats_reqs.is_empty(), + !stats_payloads.is_empty(), "Expected at least one stats request" ); - let payload = decode_stats_payload(&stats_reqs[0].body); + let payload = &stats_payloads[0]; let all_peer_tags: Vec<&str> = payload .stats .iter() @@ -1090,12 +1113,11 @@ async fn test_peer_tags_in_flushed_stats() { #[tokio::test] #[serial] async fn test_mini_agent_named_pipe_with_real_flushers() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let pipe_name = r"\\.\pipe\dd_trace_real_flusher_test"; let mut config = create_tcp_test_config(0); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.dd_apm_windows_pipe_name = Some(pipe_name.to_string()); config.dd_apm_receiver_port = 0; config.agent_stats_computation_enabled = true; @@ -1133,26 +1155,25 @@ async fn test_mini_agent_named_pipe_with_real_flushers() { .expect("Failed to send /v0.4/traces request over named pipe"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; + verify_trace_request(&mock_intake).await; // Trigger shutdown to force flush in progress concentrator buckets let _ = shutdown_tx.send(true); let _ = agent_handle.await; - verify_stats_request(&mock_server).await; + verify_stats_request(&mock_intake).await; } #[cfg(all(test, windows, feature = "windows-pipes"))] #[tokio::test] #[serial] async fn test_mini_agent_dual_transport_with_real_flushers() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let pipe_name = r"\\.\pipe\dd_trace_dual_transport_test"; let tcp_port: u16 = 8130; let mut config = create_tcp_test_config(tcp_port); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.dd_apm_windows_pipe_name = Some(pipe_name.to_string()); // Both transports are deliberately set on the same agent: a non-zero TCP // port AND a pipe name. They must come up concurrently. @@ -1215,12 +1236,16 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { .expect("Failed to send /v0.4/traces request over named pipe"); assert_eq!(pipe_response.status(), StatusCode::OK); - wait_for_request_at_path(&mock_server, "/api/v0.2/traces").await; + wait_for_capture( + || !mock_intake.trace_payloads().is_empty(), + "a trace request at /api/v0.2/traces", + ) + .await; // Both payloads must reach the same backend through the shared flusher // pipeline. The flusher may batch them into one POST or two; either is // fine, what matters is that both service-name needles show up. - let trace_reqs = mock_server.get_requests_for_path("/api/v0.2/traces"); + let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); assert!( !trace_reqs.is_empty(), "no trace POST reached backend; expected traces from both transports" @@ -1245,5 +1270,5 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { // a transport, agent_handle would hang or stats wouldn't arrive. let _ = shutdown_tx.send(true); let _ = agent_handle.await; - verify_stats_request(&mock_server).await; + verify_stats_request(&mock_intake).await; } From 37f5b32a853f82a00523051827af074662f3e8fa Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 13:59:34 -0400 Subject: [PATCH 3/6] Rename the standalone binary and env vars to mock-intake in crate docs The datadog-lambda-extension standalone binary and its environment variables are now named mock-intake / MOCK_INTAKE_*; align the crate documentation. --- crates/datadog-mock-intake/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 8f477e6e..a360cb5f 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -21,11 +21,11 @@ //! //! `MockIntake::start()` binds `127.0.0.1:0` with defaults and is used by the //! trace-agent integration tests in this workspace and by the APM / DSM -//! integration tests in `bottlecap`. The standalone `fake-intake` binary in +//! integration tests in `bottlecap`. The standalone `mock-intake` binary in //! the `datadog-lambda-extension` repository wraps this crate with request //! summaries, optional stats failure injection -//! (`FAKE_INTAKE_FAIL_STATS_FIRST_N`), and optional JSON dumps -//! (`FAKE_INTAKE_DUMP_DIR`) for local debugging against a live tracer. +//! (`MOCK_INTAKE_FAIL_STATS_FIRST_N`), and optional JSON dumps +//! (`MOCK_INTAKE_DUMP_DIR`) for local debugging against a live tracer. //! //! This crate is self-contained (no host config, logging framework, or trace //! processing dependencies) so it stays usable from multiple repositories. From a07419337cc4ee3ceac829e61e8cd1797739c5ae Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 18:54:18 -0400 Subject: [PATCH 4/6] Cap decompressed body size and skip unused dump serialization in mock-intake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Guards decompress() against decompression-bomb-style memory blow-up by capping gzip/zstd output at 64 MiB. Also skips building the diagnostic JSON dump when no dump_dir is configured (the default), avoiding a wasted allocation on every request, and marks DSM pipeline-stats payload fields #[serde(default)] so a payload missing an optional field still decodes. 🤖 --- crates/datadog-mock-intake/src/lib.rs | 78 +++++++++++++++++++-------- 1 file changed, 55 insertions(+), 23 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index a360cb5f..51ace914 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -62,39 +62,45 @@ use tokio::task::JoinHandle; /// before decompression. const MAX_BODY_SIZE: usize = 2_097_152; +/// Maximum accepted size for a decompressed body. `MAX_BODY_SIZE` caps the +/// wire body, but gzip/zstd can still expand a small compressed body into a +/// much larger one; this bounds that expansion so a pathological or +/// malformed payload cannot exhaust memory. +const MAX_DECOMPRESSED_SIZE: usize = 64 * 1024 * 1024; + /// A DSM pipeline-stats payload as it lands on `/api/v0.1/pipeline_stats`. /// Only the fields tests assert on are decoded; serde ignores the rest /// (including the `serde_bytes` latency sketches). JSON dumps therefore /// contain only these fields, not the full wire payload. #[derive(Clone, Debug, Deserialize, Serialize)] pub struct PipelineStatsPayload { - #[serde(rename = "Env")] + #[serde(rename = "Env", default)] pub env: String, - #[serde(rename = "Service")] + #[serde(rename = "Service", default)] pub service: String, - #[serde(rename = "TracerVersion")] + #[serde(rename = "TracerVersion", default)] pub tracer_version: String, - #[serde(rename = "Version")] + #[serde(rename = "Version", default)] pub version: String, - #[serde(rename = "Tags")] + #[serde(rename = "Tags", default)] pub tags: Vec, - #[serde(rename = "Stats")] + #[serde(rename = "Stats", default)] pub stats: Vec, } #[derive(Clone, Debug, Deserialize, Serialize)] pub struct PipelineStatsBucket { - #[serde(rename = "Stats")] + #[serde(rename = "Stats", default)] pub stats: Vec, } #[derive(Clone, Debug, Deserialize, Serialize)] pub struct PipelineStatsPoint { - #[serde(rename = "Hash")] + #[serde(rename = "Hash", default)] pub hash: u64, - #[serde(rename = "ParentHash")] + #[serde(rename = "ParentHash", default)] pub parent_hash: u64, - #[serde(rename = "EdgeTags")] + #[serde(rename = "EdgeTags", default)] pub edge_tags: Vec, } @@ -499,14 +505,16 @@ fn handle_stats( ); } - if let Some(payload) = &handled.decoded { + if state.options.dump_dir.is_some() + && let Some(payload) = &handled.decoded + { dump_request( state, handled.request_id, "/api/v0.2/stats", headers, handled.status, - serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + serde_json::to_value(payload).unwrap_or(serde_json::Value::Null), ); } @@ -570,7 +578,9 @@ fn handle_traces( ); } - if let Some(payload) = &handled.decoded { + if state.options.dump_dir.is_some() + && let Some(payload) = &handled.decoded + { dump_request( state, handled.request_id, @@ -635,14 +645,16 @@ fn handle_pipeline_stats( ); } - if let Some(payload) = &handled.decoded { + if state.options.dump_dir.is_some() + && let Some(payload) = &handled.decoded + { dump_request( state, handled.request_id, "/api/v0.1/pipeline_stats", headers, handled.status, - serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + serde_json::to_value(payload).unwrap_or(serde_json::Value::Null), ); } @@ -1004,15 +1016,16 @@ fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, Strin match encoding.as_str() { "gzip" => { - let mut decoder = flate2::read::GzDecoder::new(body.as_ref()); - let mut out = Vec::new(); - decoder - .read_to_end(&mut out) - .map_err(|e| format!("mock_intake: gzip decode failed: {e}"))?; - Ok(out) + let decoder = flate2::read::GzDecoder::new(body.as_ref()); + read_capped(decoder, MAX_DECOMPRESSED_SIZE) + .map_err(|e| format!("mock_intake: gzip decode failed: {e}")) + } + "zstd" => { + let decoder = zstd::stream::read::Decoder::new(body.as_ref()) + .map_err(|e| format!("mock_intake: zstd decoder init failed: {e}"))?; + read_capped(decoder, MAX_DECOMPRESSED_SIZE) + .map_err(|e| format!("mock_intake: zstd decode failed: {e}")) } - "zstd" => zstd::stream::decode_all(body.as_ref()) - .map_err(|e| format!("mock_intake: zstd decode failed: {e}")), _ => { if !encoding.is_empty() { eprintln!( @@ -1024,6 +1037,25 @@ fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, Strin } } +/// Read `reader` to the end, rejecting output past `cap` bytes instead of +/// buffering it. Guards decompression against small inputs that expand into +/// an unbounded amount of memory (a decompression bomb). +fn read_capped(reader: impl Read, cap: usize) -> Result, String> { + let mut out = Vec::new(); + // Ask for one byte more than the cap so an exactly-at-cap payload isn't + // mistaken for a truncated, over-cap one. + reader + .take(cap as u64 + 1) + .read_to_end(&mut out) + .map_err(|e| format!("read failed: {e}"))?; + if out.len() > cap { + return Err(format!( + "decompressed body exceeds {cap} byte cap, rejecting as a likely decompression bomb" + )); + } + Ok(out) +} + #[cfg(test)] mod tests { From b14334055dc37c6c788e1ebc5755d7114f726bbe Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 00:34:18 -0400 Subject: [PATCH 5/6] Abort accepted connections when mock intake is dropped --- crates/datadog-mock-intake/Cargo.toml | 2 +- crates/datadog-mock-intake/src/lib.rs | 101 +++++++++++++++++++++++++- 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/crates/datadog-mock-intake/Cargo.toml b/crates/datadog-mock-intake/Cargo.toml index db1cf336..894fe183 100644 --- a/crates/datadog-mock-intake/Cargo.toml +++ b/crates/datadog-mock-intake/Cargo.toml @@ -22,7 +22,7 @@ rmp-serde = { version = "1.3.1", default-features = false } serde = { version = "1.0.145", features = ["derive"] } serde_json = "1.0" thiserror = { version = "1.0.58", default-features = false } -tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "sync", "time"] } +tokio = { version = "1", features = ["io-util", "macros", "net", "rt-multi-thread", "sync", "time"] } zstd = { version = "0.13.3", default-features = false } [dev-dependencies] diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 51ace914..4ba0f4e4 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -201,6 +201,7 @@ pub struct MockIntake { state: std::sync::Arc, shutdown_tx: Option>, task: Option>, + connections: std::sync::Arc>>>, } impl MockIntake { @@ -243,7 +244,10 @@ impl MockIntake { }); let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>(); + let connections: std::sync::Arc>>> = + std::sync::Arc::default(); let task_state = std::sync::Arc::clone(&state); + let task_connections = std::sync::Arc::clone(&connections); let task = tokio::spawn(async move { let state = task_state; loop { @@ -259,7 +263,7 @@ impl MockIntake { let io = TokioIo::new(stream); let state = std::sync::Arc::clone(&state); - tokio::spawn(async move { + let handle = tokio::spawn(async move { let service = service_fn(move |req: Request| { let state = std::sync::Arc::clone(&state); async move { handle_request(state, req).await } @@ -268,6 +272,10 @@ impl MockIntake { .serve_connection(io, service) .await; }); + if let Ok(mut conns) = task_connections.lock() { + conns.retain(|h| !h.is_finished()); + conns.push(handle); + } } _ = &mut shutdown_rx => { break; @@ -281,6 +289,7 @@ impl MockIntake { state, shutdown_tx: Some(shutdown_tx), task: Some(task), + connections, }) } @@ -369,6 +378,13 @@ impl Drop for MockIntake { if let Some(task) = self.task.take() { task.abort(); } + // Abort any accepted-but-still-running connection tasks so a + // keep-alive client cannot keep sending requests after Drop. + if let Ok(mut conns) = self.connections.lock() { + for handle in conns.drain(..) { + handle.abort(); + } + } } } @@ -1248,6 +1264,89 @@ mod tests { assert!(bound, "listener port was not released after Drop"); } + #[tokio::test] + async fn drop_aborts_keep_alive_connections() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let intake = start_default().await; + let state = std::sync::Arc::clone(&intake.state); + let addr = intake + .base_url() + .trim_start_matches("http://") + .to_string(); + + // Raw TCP so the connection stays open (keep-alive) across Drop. + let mut stream = tokio::net::TcpStream::connect(&addr) + .await + .expect("test: connect to mock intake failed"); + let body = b"not-msgpack"; + let request = format!( + "POST /api/v0.2/stats HTTP/1.1\r\nHost: {addr}\r\nContent-Type: application/msgpack\r\nContent-Length: {}\r\n\r\n", + body.len() + ); + stream + .write_all(request.as_bytes()) + .await + .expect("test: write request head failed"); + stream + .write_all(body.as_slice()) + .await + .expect("test: write request body failed"); + + // The response has an empty body, so it ends at the header + // terminator; read until we have seen it. Do not read to EOF: the + // server keeps the connection open for keep-alive reuse. + let mut buf = Vec::new(); + let mut chunk = [0u8; 1024]; + loop { + let n = stream + .read(&mut chunk) + .await + .expect("test: read response failed"); + assert!(n > 0, "test: connection closed before response complete"); + buf.extend_from_slice(&chunk[..n]); + if buf.windows(4).any(|w| w == b"\r\n\r\n") { + break; + } + } + assert!( + buf.starts_with(b"HTTP/1.1"), + "test: expected an HTTP response, got {buf:?}" + ); + let captured_before = state + .captured + .lock() + .expect("test: requests mutex poisoned") + .requests + .len(); + assert_eq!(captured_before, 1, "test: first request should be captured"); + + drop(intake); + + // Reuse the still-open keep-alive connection. The connection task + // was aborted on Drop, so this request must not be captured. + stream + .write_all(request.as_bytes()) + .await + .expect("test: write second request failed"); + stream + .write_all(body.as_slice()) + .await + .expect("test: write second body failed"); + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + + let captured_after = state + .captured + .lock() + .expect("test: requests mutex poisoned") + .requests + .len(); + assert_eq!( + captured_after, 1, + "request sent after Drop must not be captured" + ); + } + #[tokio::test] async fn stats_endpoint_decodes_msgpack_and_gzip() { let intake = start_default().await; From de90579d15ad58b8db6a3c167b3133c67bf08c36 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 00:39:51 -0400 Subject: [PATCH 6/6] Assert no raw stats requests in no-stats integration test --- crates/datadog-trace-agent/tests/integration_test.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 0f5722ef..08cbdcec 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -197,11 +197,11 @@ pub async fn verify_stats_request(mock_intake: &MockIntake) { /// Helper to verify stats request was not sent to mock intake pub fn verify_no_stats_request(mock_intake: &MockIntake) { - let stats_payloads = mock_intake.stats_payloads(); + let stats_requests = mock_intake.requests_for_path("/api/v0.2/stats"); assert!( - stats_payloads.is_empty(), + stats_requests.is_empty(), "Expected no stats request to mock intake, received {} request(s)", - stats_payloads.len() + stats_requests.len() ); }