Support obfuscation for valkey spans (#33187)

Signed-off-by: keisku <[email protected]>
Co-authored-by: maxime mouial <[email protected]>

comp/trace/config/config_test.go

@@ -1661,6 +1661,34 @@ func TestLoadEnv(t *testing.T) {
 		assert.True(t, cfg.Obfuscation.Redis.RemoveAllArgs)
 	})
 
+	env = "DD_APM_OBFUSCATION_VALKEY_ENABLED"
+	t.Run(env, func(t *testing.T) {
+		t.Setenv(env, "true")
+
+		c := buildConfigComponent(t, true, fx.Replace(corecomp.MockParams{
+			Params: corecomp.Params{ConfFilePath: "./testdata/full.yaml"},
+		}))
+		cfg := c.Object()
+
+		assert.NotNil(t, cfg)
+		assert.True(t, pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.valkey.enabled"))
+		assert.True(t, cfg.Obfuscation.Valkey.Enabled)
+	})
+
+	env = "DD_APM_OBFUSCATION_VALKEY_REMOVE_ALL_ARGS"
+	t.Run(env, func(t *testing.T) {
+		t.Setenv(env, "true")
+
+		c := buildConfigComponent(t, true, fx.Replace(corecomp.MockParams{
+			Params: corecomp.Params{ConfFilePath: "./testdata/full.yaml"},
+		}))
+		cfg := c.Object()
+
+		assert.NotNil(t, cfg)
+		assert.True(t, pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.valkey.remove_all_args"))
+		assert.True(t, cfg.Obfuscation.Valkey.RemoveAllArgs)
+	})
+
 	env = "DD_APM_OBFUSCATION_REMOVE_STACK_TRACES"
 	t.Run(env, func(t *testing.T) {
 		t.Setenv(env, "true")

comp/trace/config/setup.go

@@ -441,6 +441,8 @@ func applyDatadogConfig(c *config.AgentConfig, core corecompcfg.Component) error
 	c.Obfuscation.Memcached.KeepCommand = pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.memcached.keep_command")
 	c.Obfuscation.Redis.Enabled = pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.redis.enabled")
 	c.Obfuscation.Redis.RemoveAllArgs = pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.redis.remove_all_args")
+	c.Obfuscation.Valkey.Enabled = pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.valkey.enabled")
+	c.Obfuscation.Valkey.RemoveAllArgs = pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.valkey.remove_all_args")
 	c.Obfuscation.CreditCards.Enabled = pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.credit_cards.enabled")
 	c.Obfuscation.CreditCards.Luhn = pkgconfigsetup.Datadog().GetBool("apm_config.obfuscation.credit_cards.luhn")
 	c.Obfuscation.CreditCards.KeepValues = pkgconfigsetup.Datadog().GetStringSlice("apm_config.obfuscation.credit_cards.keep_values")

comp/trace/config/testdata/full.yaml

@@ -84,6 +84,9 @@ apm_config:
     redis:
       enabled: true
       remove_all_args: true
+    valkey:
+      enabled: true
+      remove_all_args: true
     memcached:
       enabled: true
       keep_command: true

pkg/collector/python/test_datadog_agent.go

@@ -154,6 +154,10 @@ func testObfuscaterConfig(t *testing.T) {
 			Enabled:       true,
 			RemoveAllArgs: false,
 		},
+		Valkey: obfuscate.ValkeyConfig{
+			Enabled:       true,
+			RemoveAllArgs: false,
+		},
 		Memcached: obfuscate.MemcachedConfig{
 			Enabled:     true,
 			KeepCommand: false,

pkg/config/config_template.yaml

@@ -1230,6 +1230,15 @@ api_key:
   ##        When true, replaces all arguments of a redis command with a single "?". Disabled by default.
   #         remove_all_args: false
   #
+  #     valkey:
+  ##        @param DD_APM_OBFUSCATION_VALKEY_ENABLED - boolean - optional
+  ##        Enables obfuscation rules for spans of type "valkey". Enabled by default.
+  #         enabled: true
+  #
+  ##        @param DD_APM_OBFUSCATION_VALKEY_REMOVE_ALL_ARGS - boolean - optional
+  ##        When true, replaces all arguments of a valkey command with a single "?". Disabled by default.
+  #         remove_all_args: false
+  #
   ##    @param DD_APM_OBFUSCATION_REMOVE_STACK_TRACES - boolean - optional
   ##    Enables removing stack traces to replace them with "?". Disabled by default.
   #     remove_stack_traces: false

pkg/config/setup/apm.go

@@ -41,6 +41,8 @@ func setupAPM(config pkgconfigmodel.Setup) {
 	config.BindEnvAndSetDefault("apm_config.obfuscation.remove_stack_traces", false, "DD_APM_OBFUSCATION_REMOVE_STACK_TRACES")
 	config.BindEnvAndSetDefault("apm_config.obfuscation.redis.enabled", true, "DD_APM_OBFUSCATION_REDIS_ENABLED")
 	config.BindEnvAndSetDefault("apm_config.obfuscation.redis.remove_all_args", false, "DD_APM_OBFUSCATION_REDIS_REMOVE_ALL_ARGS")
+	config.BindEnvAndSetDefault("apm_config.obfuscation.valkey.enabled", true, "DD_APM_OBFUSCATION_VALKEY_ENABLED")
+	config.BindEnvAndSetDefault("apm_config.obfuscation.valkey.remove_all_args", false, "DD_APM_OBFUSCATION_VALKEY_REMOVE_ALL_ARGS")
 	config.BindEnvAndSetDefault("apm_config.obfuscation.memcached.enabled", true, "DD_APM_OBFUSCATION_MEMCACHED_ENABLED")
 	config.BindEnvAndSetDefault("apm_config.obfuscation.memcached.keep_command", false, "DD_APM_OBFUSCATION_MEMCACHED_KEEP_COMMAND")
 	config.BindEnvAndSetDefault("apm_config.obfuscation.cache.enabled", true, "DD_APM_OBFUSCATION_CACHE_ENABLED")

pkg/obfuscate/obfuscate.go

@@ -91,6 +91,9 @@ type Config struct {
 	// Redis holds the obfuscation settings for Redis commands.
 	Redis RedisConfig `mapstructure:"redis"`
 
+	// Valkey holds the obfuscation settings for Valkey commands.
+	Valkey ValkeyConfig `mapstructure:"valkey"`
+
 	// Memcached holds the obfuscation settings for Memcached commands.
 	Memcached MemcachedConfig `mapstructure:"memcached"`
 
@@ -232,6 +235,16 @@ type RedisConfig struct {
 	RemoveAllArgs bool `mapstructure:"remove_all_args"`
 }
 
+// ValkeyConfig holds the configuration settings for Valkey obfuscation
+type ValkeyConfig struct {
+	// Enabled specifies whether this feature should be enabled.
+	Enabled bool `mapstructure:"enabled"`
+
+	// RemoveAllArgs specifies whether all arguments to a given Valkey
+	// command should be obfuscated.
+	RemoveAllArgs bool `mapstructure:"remove_all_args"`
+}
+
 // MemcachedConfig holds the configuration settings for Memcached obfuscation
 type MemcachedConfig struct {
 	// Enabled specifies whether this feature should be enabled.

pkg/trace/agent/fuzz_test.go

@@ -89,6 +89,11 @@ func FuzzObfuscateSpan(f *testing.F) {
 			Resource: "SET k v\nGET k",
 			Meta:     map[string]string{"redis.raw_command": "SET k v\nGET k"},
 		},
+		{
+			Type:     "valkey",
+			Resource: "SET k v\nGET k",
+			Meta:     map[string]string{"valkey.raw_command": "SET k v\nGET k"},
+		},
 		{
 			Type:     "sql",
 			Resource: "UPDATE users(name) SET ('Jim')",

pkg/trace/agent/obfuscate.go

@@ -62,12 +62,16 @@ func (a *Agent) obfuscateSpan(span *pb.Span) {
 			// no error was thrown but no query was found/sanitized either
 			return
 		}
-	case "redis":
-		// if a span is redis type, it should be quantized regardless of obfuscation setting
+	case "redis", "valkey":
+		// if a span is redis/valkey type, it should be quantized regardless of obfuscation setting.
+		// valkey is a folk of redis, so we can use the same logic for both.
 		span.Resource = o.QuantizeRedisString(span.Resource)
-		if a.conf.Obfuscation.Redis.Enabled {
+		if span.Type == "redis" && a.conf.Obfuscation.Redis.Enabled {
 			transform.ObfuscateRedisSpan(o, span, a.conf.Obfuscation.Redis.RemoveAllArgs)
 		}
+		if span.Type == "valkey" && a.conf.Obfuscation.Valkey.Enabled {
+			transform.ObfuscateValkeySpan(o, span, a.conf.Obfuscation.Valkey.RemoveAllArgs)
+		}
 	case "memcached":
 		if !a.conf.Obfuscation.Memcached.Enabled {
 			return
@@ -164,7 +168,7 @@ func (a *Agent) obfuscateStatsGroup(b *pb.ClientGroupedStats) {
 		} else {
 			b.Resource = oq.Query
 		}
-	case "redis":
+	case "redis", "valkey":
 		b.Resource = o.QuantizeRedisString(b.Resource)
 	}
 }

pkg/trace/agent/obfuscate_test.go

@@ -44,6 +44,7 @@ func TestObfuscateStatsGroup(t *testing.T) {
 		{statsGroup("sql", "SELECT 1 FROM db"), "SELECT ? FROM db"},
 		{statsGroup("sql", "SELECT 1\nFROM Blogs AS [b\nORDER BY [b]"), textNonParsable},
 		{statsGroup("redis", "ADD 1, 2"), "ADD"},
+		{statsGroup("valkey", "ADD 1, 2"), "ADD"},
 		{statsGroup("other", "ADD 1, 2"), "ADD 1, 2"},
 	} {
 		agnt, stop := agentWithDefaults()
@@ -70,6 +71,20 @@ func TestObfuscateDefaults(t *testing.T) {
 		assert.Equal(t, "SET GET", span.Resource)
 	})
 
+	t.Run("valkey", func(t *testing.T) {
+		cmd := "SET k v\nGET k"
+		span := &pb.Span{
+			Type:     "valkey",
+			Resource: cmd,
+			Meta:     map[string]string{"valkey.raw_command": cmd},
+		}
+		agnt, stop := agentWithDefaults()
+		defer stop()
+		agnt.obfuscateSpan(span)
+		assert.Equal(t, cmd, span.Meta["valkey.raw_command"])
+		assert.Equal(t, "SET GET", span.Resource)
+	})
+
 	t.Run("sql", func(t *testing.T) {
 		query := "UPDATE users(name) SET ('Jim')"
 		span := &pb.Span{
@@ -143,6 +158,33 @@ func TestObfuscateConfig(t *testing.T) {
 		&config.ObfuscationConfig{},
 	))
 
+	t.Run("valkey/enabled", testConfig(
+		"valkey",
+		"valkey.raw_command",
+		"SET key val",
+		"SET key ?",
+		&config.ObfuscationConfig{Valkey: obfuscate.ValkeyConfig{Enabled: true}},
+	))
+
+	t.Run("valkey/remove_all_args", testConfig(
+		"valkey",
+		"valkey.raw_command",
+		"SET key val",
+		"SET ?",
+		&config.ObfuscationConfig{Valkey: obfuscate.ValkeyConfig{
+			Enabled:       true,
+			RemoveAllArgs: true,
+		}},
+	))
+
+	t.Run("valkey/disabled", testConfig(
+		"valkey",
+		"valkey.raw_command",
+		"SET key val",
+		"SET key val",
+		&config.ObfuscationConfig{},
+	))
+
 	t.Run("http/enabled", testConfig(
 		"http",
 		"http.url",

pkg/trace/api/info.go

@@ -35,6 +35,7 @@ func (r *HTTPReceiver) makeInfoHandler() (hash string, handler http.HandlerFunc)
 		HTTP                 obfuscate.HTTPConfig      `json:"http"`
 		RemoveStackTraces    bool                      `json:"remove_stack_traces"`
 		Redis                obfuscate.RedisConfig     `json:"redis"`
+		Valkey               obfuscate.ValkeyConfig    `json:"valkey"`
 		Memcached            obfuscate.MemcachedConfig `json:"memcached"`
 	}
 	type reducedConfig struct {
@@ -61,6 +62,7 @@ func (r *HTTPReceiver) makeInfoHandler() (hash string, handler http.HandlerFunc)
 		oconf.HTTP = o.HTTP
 		oconf.RemoveStackTraces = o.RemoveStackTraces
 		oconf.Redis = o.Redis
+		oconf.Valkey = o.Valkey
 		oconf.Memcached = o.Memcached
 	}
 

pkg/trace/api/info_test.go

@@ -230,6 +230,7 @@ func TestInfoHandler(t *testing.T) {
 		},
 		RemoveStackTraces: false,
 		Redis:             obfuscate.RedisConfig{Enabled: true},
+		Valkey:            obfuscate.ValkeyConfig{Enabled: true},
 		Memcached:         obfuscate.MemcachedConfig{Enabled: false},
 	}
 	conf := &config.AgentConfig{
@@ -328,6 +329,7 @@ func TestInfoHandler(t *testing.T) {
 				},
 				"remove_stack_traces": nil,
 				"redis":               nil,
+				"valkey":              nil,
 				"memcached":           nil,
 			},
 		},

pkg/trace/config/config.go

@@ -110,6 +110,10 @@ type ObfuscationConfig struct {
 	// for spans of type "redis".
 	Redis obfuscate.RedisConfig `mapstructure:"redis"`
 
+	// Valkey holds the configuration for obfuscating the "valkey.raw_command" tag
+	// for spans of type "valkey".
+	Valkey obfuscate.ValkeyConfig `mapstructure:"valkey"`
+
 	// Memcached holds the configuration for obfuscating the "memcached.command" tag
 	// for spans of type "memcached".
 	Memcached obfuscate.MemcachedConfig `mapstructure:"memcached"`
@@ -152,6 +156,7 @@ func (o *ObfuscationConfig) Export(conf *AgentConfig) obfuscate.Config {
 		SQLExecPlanNormalize: o.SQLExecPlanNormalize,
 		HTTP:                 o.HTTP,
 		Redis:                o.Redis,
+		Valkey:               o.Valkey,
 		Memcached:            o.Memcached,
 		CreditCard:           o.CreditCards,
 		Logger:               new(debugLogger),

pkg/trace/transform/obfuscate.go

@@ -14,6 +14,8 @@ import (
 const (
 	// TagRedisRawCommand represents a redis raw command tag
 	TagRedisRawCommand = "redis.raw_command"
+	// TagValkeyRawCommand represents a redis raw command tag
+	TagValkeyRawCommand = "valkey.raw_command"
 	// TagMemcachedCommand represents a memcached command tag
 	TagMemcachedCommand = "memcached.command"
 	// TagMongoDBQuery represents a MongoDB query tag
@@ -66,3 +68,15 @@ func ObfuscateRedisSpan(o *obfuscate.Obfuscator, span *pb.Span, removeAllArgs bo
 	}
 	span.Meta[TagRedisRawCommand] = o.ObfuscateRedisString(span.Meta[TagRedisRawCommand])
 }
+
+// ObfuscateValkeySpan obfuscates a Valkey span using pkg/obfuscate logic
+func ObfuscateValkeySpan(o *obfuscate.Obfuscator, span *pb.Span, removeAllArgs bool) {
+	if span.Meta == nil || span.Meta[TagValkeyRawCommand] == "" {
+		return
+	}
+	if removeAllArgs {
+		span.Meta[TagValkeyRawCommand] = o.RemoveAllRedisArgs(span.Meta[TagValkeyRawCommand])
+		return
+	}
+	span.Meta[TagValkeyRawCommand] = o.ObfuscateRedisString(span.Meta[TagValkeyRawCommand])
+}

releasenotes/notes/add-valkey-support-c800f16cfb220811.yaml

@@ -0,0 +1,14 @@
+# Each section from every release note are combined when the
+# CHANGELOG.rst is rendered. So the text needs to be worded so that
+# it does not depend on any information only available in another
+# section. This may mean repeating some details, but each section
+# must be readable independently of the other.
+#
+# Each section note must be formatted as reStructuredText.
+---
+features:
+  - |
+    Added support for obfuscation for valkey command. This feature is enabled by default.
+    To disable it, set ``DD_APM_OBFUSCATION_VALKEY_ENABLED=false``.
+    To replace all valkey command arguments with a single ``?``,
+    set ``DD_APM_OBFUSCATION_VALKEY_REMOVE_ALL_ARGS=true`` (default: false).