feat!: integrate obfuscation to the stats exporter [APMSP-2764]

Cargo.toml

@@ -76,6 +76,7 @@ license = "Apache-2.0"
 authors = ["Datadog Inc. <info@datadoghq.com>"]
 
 [workspace.dependencies]
+arc-swap = "1.7.1"
 hyper = { version = "1.6", features = [
   "http1",
   "client",

libdd-common-ffi/src/vec.rs

@@ -181,7 +181,7 @@ mod tests {
         let vec = vec![0, 2, 4, 6];
         let ffi_vec: Vec<u8> = Vec::from(vec.clone());
 
-        for (a, b) in vec.iter().zip(ffi_vec.into_iter()) {
+        for (a, b) in vec.iter().zip(&ffi_vec) {
             assert_eq!(a, b)
         }
     }

libdd-data-pipeline/Cargo.toml

@@ -12,8 +12,8 @@ autobenches = false
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+arc-swap.workspace = true
 anyhow = { version = "1.0" }
-arc-swap = "1.7.1"
 async-trait = "0.1"
 http = "1"
 http-body-util = "0.1"
@@ -36,6 +36,7 @@ libdd-telemetry = { version = "4.0.0", path = "../libdd-telemetry", default-feat
 libdd-trace-protobuf = { version = "3.0.1", path = "../libdd-trace-protobuf" }
 libdd-trace-stats = { version = "2.0.0", path = "../libdd-trace-stats" }
 libdd-trace-utils = { version = "3.0.1", path = "../libdd-trace-utils", default-features = false }
+libdd-trace-obfuscation = { version = "2.0.0", path = "../libdd-trace-obfuscation", optional = true }
 libdd-ddsketch = { version = "1.0.1", path = "../libdd-ddsketch" }
 libdd-dogstatsd-client = { version = "2.0.0", path = "../libdd-dogstatsd-client", default-features = false }
 libdd-tinybytes = { version = "1.1.0", path = "../libdd-tinybytes", features = [
@@ -81,4 +82,8 @@ https = [
     "libdd-trace-utils/https",
     "libdd-dogstatsd-client/https",
 ]
+stats-obfuscation = [
+    "libdd-trace-obfuscation",
+    "libdd-trace-stats/stats-obfuscation"
+]
 test-utils = []

libdd-data-pipeline/src/agent_info/fetcher.rs

@@ -392,12 +392,12 @@ mod single_threaded_tests {
                         },
                         "remove_stack_traces": false,
                         "redis": {
-                                "Enabled": true,
-                                "RemoveAllArgs": false
+                                "enabled": true,
+                                "remove_all_args": false
                         },
                         "memcached": {
-                                "Enabled": true,
-                                "KeepCommand": false
+                                "enabled": true,
+                                "keep_command": false
                         }
                 }
         },
@@ -408,7 +408,7 @@ mod single_threaded_tests {
         format!("{:x}", Sha256::digest(json.as_bytes()))
     }
 
-    const TEST_INFO_HASH: &str = "b7709671827946c15603847bca76c90438579c038ec134eae19c51f1f3e3dfea";
+    const TEST_INFO_HASH: &str = "cce54bf6e7d1bf38088a3ec809bfeec160bc52d37f70bd6b581ce3c2f7be5a65";
 
     #[cfg_attr(miri, ignore)]
     #[tokio::test]

libdd-data-pipeline/src/agent_info/schema.rs

@@ -35,6 +35,8 @@ pub struct AgentInfoStruct {
     pub peer_tags: Option<Vec<String>>,
     /// List of span kinds eligible for stats computation
     pub span_kinds_stats_computed: Option<Vec<String>>,
+    /// Obfuscation version supported by the agent for client-side stats
+    pub obfuscation_version: Option<u32>,
     /// Container tags hash from HTTP response header
     pub container_tags_hash: Option<String>,
 }
@@ -54,37 +56,41 @@ pub struct Config {
     pub max_memory: Option<f64>,
     pub max_cpu: Option<f64>,
     pub analyzed_spans_by_service: Option<HashMap<String, HashMap<String, f64>>>,
+    pub obfuscation: Option<ObfuscationConfig>,
 }
 
 #[allow(missing_docs)]
-#[derive(Clone, Deserialize, Default, Debug, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Default, Debug, PartialEq)]
 pub struct ObfuscationConfig {
     pub elastic_search: bool,
     pub mongo: bool,
     pub sql_exec_plan: bool,
     pub sql_exec_plan_normalize: bool,
+    #[cfg(feature = "stats-obfuscation")]
+    #[serde(default)]
+    pub sql_obfuscation_mode: libdd_trace_obfuscation::sql::SqlObfuscationMode,
     pub http: HttpObfuscationConfig,
     pub remove_stack_traces: bool,
     pub redis: RedisObfuscationConfig,
     pub memcached: MemcachedObfuscationConfig,
 }
 
 #[allow(missing_docs)]
-#[derive(Clone, Deserialize, Default, Debug, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Default, Debug, PartialEq)]
 pub struct HttpObfuscationConfig {
     pub remove_query_string: bool,
     pub remove_path_digits: bool,
 }
 
 #[allow(missing_docs)]
-#[derive(Clone, Deserialize, Default, Debug, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Default, Debug, PartialEq)]
 pub struct RedisObfuscationConfig {
     pub enabled: bool,
     pub remove_all_args: bool,
 }
 
 #[allow(missing_docs)]
-#[derive(Clone, Deserialize, Default, Debug, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Default, Debug, PartialEq)]
 pub struct MemcachedObfuscationConfig {
     pub enabled: bool,
     pub keep_command: bool,

libdd-data-pipeline/src/lib.rs

@@ -3,6 +3,7 @@
 #![cfg_attr(not(test), deny(clippy::panic))]
 #![cfg_attr(not(test), deny(clippy::unwrap_used))]
 #![cfg_attr(not(test), deny(clippy::expect_used))]
+#![cfg_attr(not(test), deny(clippy::unreachable))]
 #![cfg_attr(not(test), deny(clippy::todo))]
 #![cfg_attr(not(test), deny(clippy::unimplemented))]
 

libdd-data-pipeline/src/stats_exporter.rs

@@ -16,6 +16,8 @@ use libdd_capabilities::{HttpClientTrait, MaybeSend};
 use libdd_common::Endpoint;
 use libdd_shared_runtime::Worker;
 use libdd_trace_protobuf::pb;
+#[cfg(feature = "stats-obfuscation")]
+use libdd_trace_stats::span_concentrator::SharedStatsComputationObfuscationConfig;
 use libdd_trace_stats::span_concentrator::SpanConcentrator;
 use libdd_trace_utils::send_with_retry::{send_with_retry, RetryStrategy};
 use tracing::error;
@@ -34,6 +36,8 @@ pub struct StatsExporter<H: HttpClientTrait> {
     meta: TracerMetadata,
     sequence_id: AtomicU64,
     client: H,
+    #[cfg(feature = "stats-obfuscation")]
+    obfuscation_config: SharedStatsComputationObfuscationConfig,
 }
 
 impl<H: HttpClientTrait> StatsExporter<H> {
@@ -43,14 +47,14 @@ impl<H: HttpClientTrait> StatsExporter<H> {
     /// - `concentrator` SpanConcentrator storing the stats to be sent to the agent
     /// - `meta` metadata used in ClientStatsPayload and as headers to send stats to the agent
     /// - `endpoint` the Endpoint used to send stats to the agent
-    /// - `cancellation_token` Token used to safely shutdown the exporter by force flushing the
-    ///   concentrator
     pub fn new(
         flush_interval: time::Duration,
         concentrator: Arc<Mutex<SpanConcentrator>>,
         meta: TracerMetadata,
         endpoint: Endpoint,
         client: H,
+        #[cfg(feature = "stats-obfuscation")]
+        obfuscation_config: SharedStatsComputationObfuscationConfig,
     ) -> Self {
         Self {
             flush_interval,
@@ -59,6 +63,8 @@ impl<H: HttpClientTrait> StatsExporter<H> {
             meta,
             sequence_id: AtomicU64::new(0),
             client,
+            #[cfg(feature = "stats-obfuscation")]
+            obfuscation_config,
         }
     }
 
@@ -91,6 +97,16 @@ impl<H: HttpClientTrait> StatsExporter<H> {
             libdd_common::header::APPLICATION_MSGPACK,
         );
 
+        #[cfg(feature = "stats-obfuscation")]
+        if self.obfuscation_config.load().enabled {
+            headers.insert(
+                http::HeaderName::from_static("datadog-obfuscation-version"),
+                http::HeaderValue::from_static(
+                    crate::trace_exporter::stats::SUPPORTED_OBFUSCATION_VERSION_STR,
+                ),
+            );
+        }
+
         let result = send_with_retry(
             &self.client,
             &self.endpoint,
@@ -194,6 +210,8 @@ mod tests {
     use httpmock::MockServer;
     use libdd_capabilities_impl::NativeCapabilities;
     use libdd_shared_runtime::SharedRuntime;
+    #[cfg(feature = "stats-obfuscation")]
+    use libdd_trace_stats::span_concentrator::StatsComputationObfuscationConfig;
     use libdd_trace_utils::span::{trace_utils, v04::SpanSlice};
     use libdd_trace_utils::test_utils::poll_for_mock_hit;
     use time::Duration;
@@ -231,6 +249,8 @@ mod tests {
             SystemTime::now() - BUCKETS_DURATION * 3,
             vec![],
             vec![],
+            #[cfg(feature = "stats-obfuscation")]
+            None,
         );
         let mut trace = vec![];
 
@@ -272,6 +292,8 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             NativeCapabilities::new_client(),
+            #[cfg(feature = "stats-obfuscation")]
+            StatsComputationObfuscationConfig::disabled(),
         );
 
         let send_status = stats_exporter.send(true).await;
@@ -299,6 +321,8 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             NativeCapabilities::new_client(),
+            #[cfg(feature = "stats-obfuscation")]
+            StatsComputationObfuscationConfig::disabled(),
         );
 
         let send_status = stats_exporter.send(true).await;
@@ -333,6 +357,8 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             NativeCapabilities::new_client(),
+            #[cfg(feature = "stats-obfuscation")]
+            StatsComputationObfuscationConfig::disabled(),
         );
 
         let _handle = shared_runtime
@@ -374,6 +400,8 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             NativeCapabilities::new_client(),
+            #[cfg(feature = "stats-obfuscation")]
+            StatsComputationObfuscationConfig::disabled(),
         );
 
         let _handle = shared_runtime
@@ -413,4 +441,40 @@ mod tests {
             "Non-empty env should be preserved"
         );
     }
+    #[cfg(feature = "stats-obfuscation")]
+    #[cfg_attr(miri, ignore)]
+    #[tokio::test]
+    async fn test_send_stats_with_obfuscation_header() {
+        use arc_swap::ArcSwap;
+
+        let server = MockServer::start_async().await;
+
+        let mock = server
+            .mock_async(|when, then| {
+                when.method(POST)
+                    .header("Content-type", "application/msgpack")
+                    .header("datadog-obfuscation-version", "1")
+                    .path("/v0.6/stats")
+                    .body_includes("libdatadog-test");
+                then.status(200).body("");
+            })
+            .await;
+
+        let stats_exporter = StatsExporter::new(
+            BUCKETS_DURATION,
+            Arc::new(Mutex::new(get_test_concentrator())),
+            get_test_metadata(),
+            Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
+            NativeCapabilities::new_client(),
+            Arc::new(ArcSwap::from_pointee(StatsComputationObfuscationConfig {
+                enabled: true,
+                ..Default::default()
+            })),
+        );
+
+        let send_status = stats_exporter.send(true).await;
+        send_status.unwrap();
+
+        mock.assert_async().await;
+    }
 }

libdd-data-pipeline/src/trace_exporter/builder.rs

@@ -299,6 +299,11 @@ impl TraceExporterBuilder {
 
         #[cfg(not(target_arch = "wasm32"))]
         {
+            #[cfg(feature = "stats-obfuscation")]
+            use libdd_trace_stats::span_concentrator::StatsComputationObfuscationConfig;
+
+            use crate::trace_exporter::stats::StatsComputationConfig;
+
             let info_endpoint = Endpoint::from_url(add_path(&agent_url, INFO_ENDPOINT));
             let (info_fetcher, info_response_observer) =
                 AgentInfoFetcher::<H>::new(info_endpoint.clone(), Duration::from_secs(5 * 60));
@@ -378,7 +383,13 @@ impl TraceExporterBuilder {
                 shared_runtime,
                 dogstatsd,
                 common_stats_tags: vec![libdatadog_version],
-                client_side_stats: ArcSwap::new(stats.into()),
+                client_side_stats: StatsComputationConfig {
+                    status: ArcSwap::new(stats.into()),
+                    #[cfg(feature = "stats-obfuscation")]
+                    obfuscation_config: Arc::new(ArcSwap::from_pointee(
+                        StatsComputationObfuscationConfig::default(),
+                    )),
+                },
                 previous_info_state: arc_swap::ArcSwapOption::new(None),
                 info_response_observer,
                 #[cfg(feature = "telemetry")]
@@ -427,6 +438,11 @@ impl TraceExporterBuilder {
 
         #[cfg(target_arch = "wasm32")]
         {
+            #[cfg(feature = "stats-obfuscation")]
+            use libdd_trace_stats::span_concentrator::StatsComputationObfuscationConfig;
+
+            use crate::trace_exporter::stats::StatsComputationConfig;
+
             let info_endpoint = Endpoint::from_url(add_path(&agent_url, INFO_ENDPOINT));
             let (_info_fetcher, info_response_observer) =
                 AgentInfoFetcher::<H>::new(info_endpoint, Duration::from_secs(5 * 60));
@@ -462,7 +478,13 @@ impl TraceExporterBuilder {
                 shared_runtime,
                 dogstatsd,
                 common_stats_tags: vec![libdatadog_version],
-                client_side_stats: ArcSwap::new(stats.into()),
+                client_side_stats: StatsComputationConfig {
+                    status: ArcSwap::new(stats.into()),
+                    #[cfg(feature = "stats-obfuscation")]
+                    obfuscation_config: Arc::new(ArcSwap::from_pointee(
+                        StatsComputationObfuscationConfig::default(),
+                    )),
+                },
                 previous_info_state: arc_swap::ArcSwapOption::new(None),
                 info_response_observer,
                 health_metrics_enabled: self.health_metrics_enabled,