Trace Tag Replacer: functionality to scrub sensitive data from spans

.github/CODEOWNERS

@@ -11,3 +11,4 @@ NOTICE              @Datadog/libdatadog
 rustfmt.toml        @Datadog/libdatadog
 README.md           @Datadog/libdatadog
 trace-normalization @Datadog/serverless
+trace-obfuscation   @Datadog/serverless

Cargo.toml

@@ -11,6 +11,7 @@ members = [
     "ddtelemetry-ffi",
     "tools",
     "trace-normalization",
+    "trace-obfuscation",
 ]
 # https://doc.rust-lang.org/cargo/reference/resolver.html#feature-resolver-version-2
 resolver = "2"

LICENSE-3rdparty.yml

@@ -1,5 +1,5 @@
 ---
-root_name: "datadog-profiling, ddcommon, datadog-profiling-ffi, ddcommon-ffi, ddtelemetry, ddtelemetry-ffi, tools"
+root_name: "datadog-profiling, ddcommon, datadog-profiling-ffi, ddcommon-ffi, ddtelemetry, ddtelemetry-ffi, tools, datadog-trace-normalization"
 third_party_libraries:
   - package_name: aho-corasick
     package_version: 0.7.20

tools/docker/Dockerfile.build

@@ -78,6 +78,7 @@ COPY "ddtelemetry-ffi/Cargo.toml" "ddtelemetry-ffi/"
 COPY "profiling/Cargo.toml" "profiling/"
 COPY "profiling-ffi/Cargo.toml" "profiling-ffi/"
 COPY "trace-normalization/Cargo.toml" "trace-normalization/"
+COPY "trace-obfuscation/Cargo.toml" "trace-obfuscation/"
 COPY "tools/Cargo.toml" "tools/" 
 RUN find -name "Cargo.toml" | sed -e s#Cargo.toml#src/lib.rs#g | xargs -n 1 sh -c 'mkdir -p $(dirname $1); touch $1; echo $1' create_stubs
 RUN echo ddtelemetry/benches/ipc.rs tools/src/bin/dedup_headers.rs ddtelemetry/examples/tm-worker-test.rs | xargs -n 1 sh -c 'mkdir -p $(dirname $1); touch $1; echo $1' create_stubs

trace-obfuscation/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "datadog-trace-obfuscation"
+version = "2.0.0"
+authors = ["David Lee <david.lee@datadoghq.com>"]
+edition = "2021"
+
+[dependencies]
+prost = "0.11.6"
+anyhow = "1.0"
+regex = "1"
+
+[dev-dependencies]
+duplicate = "0.4.1"

trace-obfuscation/src/lib.rs

@@ -0,0 +1,12 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0. This product includes software
+// developed at Datadog (https://www.datadoghq.com/). Copyright 2023-Present
+// Datadog, Inc.
+
+#![deny(clippy::all)]
+
+pub mod pb {
+    include!("../../trace-normalization/src/pb/pb.rs");
+}
+
+pub mod replacer;

trace-obfuscation/src/replacer.rs

@@ -0,0 +1,198 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0. This product includes software
+// developed at Datadog (https://www.datadoghq.com/). Copyright 2023-Present
+// Datadog, Inc.
+
+use crate::pb;
+use regex::Regex;
+
+pub trait TraceTagReplacer {
+    fn replace_trace_tags(trace: &mut [pb::Span], rules: &[ReplaceRule]);
+}
+
+#[derive(Debug)]
+pub struct ReplaceRule<'a> {
+    // name specifies the name of the tag that the replace rule addresses. However,
+    // some exceptions apply such as:
+    // • "resource.name" will target the resource
+    // • "*" will target all tags and the resource
+    name: &'a str,
+
+    // re holds the regex pattern for matching.
+    re: regex::Regex,
+
+    // repl specifies the replacement string to be used when Pattern matches.
+    repl: &'a str,
+}
+
+struct DefaultTraceTagReplacer {}
+
+impl TraceTagReplacer for DefaultTraceTagReplacer {
+    /// replace_trace_tags replaces the tag values of all spans within a trace with a given set of rules.
+    fn replace_trace_tags(trace: &mut [pb::Span], rules: &[ReplaceRule]) {
+        for rule in rules {
+            for span in &mut *trace {
+                match rule.name {
+                    "*" => {
+                        for (_, val) in span.meta.iter_mut() {
+                            *val = rule.re.replace_all(val, rule.repl).to_string();
+                        }
+                    }
+                    "resource.name" => {
+                        span.resource = rule.re.replace_all(&span.resource, rule.repl).to_string();
+                    }
+                    _ => {
+                        if let Some(val) = span.meta.get_mut(rule.name) {
+                            let replaced_tag = rule.re.replace_all(val, rule.repl).to_string();
+                            *val = replaced_tag;
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+/// parse_rules_from_string takes an array of rules, represented as an array of length 3 arrays
+/// holding the tag name, regex pattern, and replacement string as strings.
+/// * returns a vec of ReplaceRules
+pub fn parse_rules_from_string<'a>(
+    rules: &'a [[&'a str; 3]],
+) -> anyhow::Result<Vec<ReplaceRule<'a>>> {
+    let mut vec: Vec<ReplaceRule> = Vec::with_capacity(rules.len());
+
+    for [name, pattern, repl] in rules {
+        let compiled_regex = match Regex::new(pattern) {
+            Ok(res) => res,
+            Err(err) => {
+                anyhow::bail!(format!(
+                    "Obfuscator Error: Error while parsing rule: {}",
+                    err
+                ))
+            }
+        };
+        vec.push(ReplaceRule {
+            name,
+            re: compiled_regex,
+            repl,
+        });
+    }
+    Ok(vec)
+}
+
+#[cfg(test)]
+mod tests {
+
+    use crate::pb;
+    use crate::replacer;
+    use duplicate::duplicate_item;
+    use std::collections::HashMap;
+
+    use super::TraceTagReplacer;
+
+    fn new_test_span_with_tags(tags: HashMap<&str, &str>) -> pb::Span {
+        let mut span = pb::Span {
+            duration: 10000000,
+            error: 0,
+            resource: "GET /some/raclette".to_string(),
+            service: "django".to_string(),
+            name: "django.controller".to_string(),
+            span_id: 123,
+            start: 1448466874000000000,
+            trace_id: 424242,
+            meta: HashMap::new(),
+            metrics: HashMap::from([("cheese_weight".to_string(), 100000.0)]),
+            parent_id: 1111,
+            r#type: "http".to_string(),
+            meta_struct: HashMap::new(),
+        };
+        for (key, val) in tags {
+            match key {
+                "resource.name" => {
+                    span.resource = val.to_string();
+                }
+                _ => {
+                    span.meta.insert(key.to_string(), val.to_string());
+                }
+            }
+        }
+        span
+    }
+
+    #[duplicate_item(
+        [
+        test_name   [test_replace_tags]
+        rules       [&[
+                        ["http.url", "(token/)([^/]*)", "${1}?"],
+                        ["http.url", "guid", "[REDACTED]"],
+                        ["custom.tag", "(/foo/bar/).*", "${1}extra"],
+                    ]]
+        input       [
+                        HashMap::from([
+                            ("http.url", "some/guid/token/abcdef/abc"),
+                            ("custom.tag", "/foo/bar/foo"),
+                        ])
+                    ]
+        expected    [
+                        HashMap::from([
+                            ("http.url", "some/[REDACTED]/token/?/abc"),
+                            ("custom.tag", "/foo/bar/extra"),
+                        ])
+                    ];
+        ]
+        [
+        test_name   [test_replace_tags_with_exceptions]
+        rules       [&[
+                        ["*", "(token/)([^/]*)", "${1}?"],
+                        ["*", "this", "that"],
+                        ["http.url", "guid", "[REDACTED]"],
+                        ["custom.tag", "(/foo/bar/).*", "${1}extra"],
+                        ["resource.name", "prod", "stage"],
+                    ]]
+        input       [
+                        HashMap::from([
+                            ("resource.name", "this is prod"),
+                            ("http.url", "some/[REDACTED]/token/abcdef/abc"),
+                            ("other.url", "some/guid/token/abcdef/abc"),
+                            ("custom.tag", "/foo/bar/foo"),
+                        ])
+                    ]
+        expected    [
+                        HashMap::from([
+                            ("resource.name", "this is stage"),
+                            ("http.url", "some/[REDACTED]/token/?/abc"),
+                            ("other.url", "some/guid/token/?/abc"),
+                            ("custom.tag", "/foo/bar/extra"),
+                        ])
+                    ];
+        ]
+    )]
+    #[test]
+    fn test_name() {
+        let parsed_rules = replacer::parse_rules_from_string(rules);
+        let root_span = new_test_span_with_tags(input);
+        let child_span = new_test_span_with_tags(input);
+        let mut trace = [root_span, child_span];
+
+        replacer::DefaultTraceTagReplacer::replace_trace_tags(&mut trace, &parsed_rules.unwrap());
+
+        for (key, val) in expected {
+            match key {
+                "resource.name" => {
+                    assert_eq!(val, trace[0].resource);
+                    assert_eq!(val, trace[1].resource);
+                }
+                _ => {
+                    assert_eq!(val, trace[0].meta.get(key).unwrap());
+                    assert_eq!(val, trace[1].meta.get(key).unwrap());
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn test_parse_rules_invalid_regex() {
+        let result = replacer::parse_rules_from_string(&[["http.url", ")", "${1}?"]]);
+        assert!(result.is_err());
+    }
+}