From 482eecb8a406bb61fa025e5e174ea12d254a4d75 Mon Sep 17 00:00:00 2001
From: Michel Laterman <82832767+michel-laterman@users.noreply.github.com>
Date: Fri, 19 Jul 2024 07:55:56 -0700
Subject: [PATCH] Add unit tests for `$$` appearing within policy inputs
 (#5097)

Add unit tests for special characters, mainly the escape sequence of $$ in the transpiler, and input parsing of policies.
---
 .../handler_action_policy_change_test.go      | 28 ++++++++++++
 internal/pkg/agent/transpiler/ast_test.go     | 14 ++++++
 internal/pkg/agent/transpiler/vars_test.go    | 44 +++++++++++++++++++
 internal/pkg/config/config_test.go            | 15 +++++++
 4 files changed, 101 insertions(+)

diff --git a/internal/pkg/agent/application/actions/handlers/handler_action_policy_change_test.go b/internal/pkg/agent/application/actions/handlers/handler_action_policy_change_test.go
index 530865fc16c..6ebc5ef783d 100644
--- a/internal/pkg/agent/application/actions/handlers/handler_action_policy_change_test.go
+++ b/internal/pkg/agent/application/actions/handlers/handler_action_policy_change_test.go
@@ -70,6 +70,34 @@ func TestPolicyChange(t *testing.T) {
 		change := <-ch
 		require.Equal(t, config.MustNewConfigFrom(conf), change.Config())
 	})
+	t.Run("Received config with $$ in inputs", func(t *testing.T) {
+		ch := make(chan coordinator.ConfigChange, 1)
+
+		conf := map[string]interface{}{
+			"inputs": []interface{}{map[string]interface{}{
+				"type": "key",
+				"key":  "$$$$",
+			}}}
+		action := &fleetapi.ActionPolicyChange{
+			ActionID:   "abc123",
+			ActionType: "POLICY_CHANGE",
+			Data: fleetapi.ActionPolicyChangeData{
+				Policy: conf,
+			},
+		}
+
+		cfg := configuration.DefaultConfiguration()
+		handler := NewPolicyChangeHandler(log, agentInfo, cfg, nullStore, ch, nilLogLevelSet(t), &coordinator.Coordinator{})
+
+		err := handler.Handle(context.Background(), action, ack)
+		require.NoError(t, err)
+
+		change := <-ch
+		m, err := change.Config().ToMapStr()
+		require.NoError(t, err)
+
+		require.Equal(t, conf, m)
+	})
 }
 
 func TestPolicyAcked(t *testing.T) {
diff --git a/internal/pkg/agent/transpiler/ast_test.go b/internal/pkg/agent/transpiler/ast_test.go
index aac95995908..3487b57b775 100644
--- a/internal/pkg/agent/transpiler/ast_test.go
+++ b/internal/pkg/agent/transpiler/ast_test.go
@@ -60,6 +60,20 @@ func TestAST(t *testing.T) {
 				},
 			},
 		},
+		"special characters in strings": {
+			hashmap: map[string]interface{}{
+				"key":            "$1$$2$$$3$$$$4",
+				"$1$$2$$$3$$$$4": "value",
+			},
+			ast: &AST{
+				root: &Dict{
+					value: []Node{
+						&Key{name: "key", value: &StrVal{value: "$1$$2$$$3$$$$4"}},
+						&Key{name: "$1$$2$$$3$$$$4", value: &StrVal{value: "value"}},
+					},
+				},
+			},
+		},
 		"integer as key": {
 			hashmap: map[string]interface{}{
 				"1": []string{"/var/log/log1", "/var/log/log2"},
diff --git a/internal/pkg/agent/transpiler/vars_test.go b/internal/pkg/agent/transpiler/vars_test.go
index 356cd47c9b1..becf267f247 100644
--- a/internal/pkg/agent/transpiler/vars_test.go
+++ b/internal/pkg/agent/transpiler/vars_test.go
@@ -34,6 +34,14 @@ func TestVars_Replace(t *testing.T) {
 		"other": map[string]interface{}{
 			"data": "info",
 		},
+		"special": map[string]interface{}{
+			"key1": "$1$$2",
+			"key2": "1$2$$",
+			"key3": "${abcd}",
+			"key4": "$${abcd}",
+			"key5": "${",
+			"key6": "$${",
+		},
 	})
 	tests := []struct {
 		Input   string
@@ -209,6 +217,42 @@ func TestVars_Replace(t *testing.T) {
 			false,
 			false,
 		},
+		{
+			`${special.key1}`,
+			NewStrVal("$1$$2"),
+			false,
+			false,
+		},
+		{
+			`${special.key2}`,
+			NewStrVal("1$2$$"),
+			false,
+			false,
+		},
+		{
+			`${special.key3}`,
+			NewStrVal("${abcd}"),
+			false,
+			false,
+		},
+		{
+			`${special.key4}`,
+			NewStrVal("$${abcd}"),
+			false,
+			false,
+		},
+		{
+			`${special.key5}`,
+			NewStrVal("${"),
+			false,
+			false,
+		},
+		{
+			`${special.key6}`,
+			NewStrVal("$${"),
+			false,
+			false,
+		},
 	}
 	for _, test := range tests {
 		t.Run(test.Input, func(t *testing.T) {
diff --git a/internal/pkg/config/config_test.go b/internal/pkg/config/config_test.go
index 5d393d3486a..3cb560f6d4e 100644
--- a/internal/pkg/config/config_test.go
+++ b/internal/pkg/config/config_test.go
@@ -143,3 +143,18 @@ func dumpToYAML(t *testing.T, out string, in interface{}) {
 	err = os.WriteFile(out, b, 0600)
 	require.NoError(t, err)
 }
+
+func TestDollarSignsInInputs(t *testing.T) {
+	in := map[string]interface{}{
+		"inputs": []interface{}{
+			map[string]interface{}{
+				"type": "logfile",
+				"what": "$$$$",
+			},
+		},
+	}
+	c := MustNewConfigFrom(in)
+	out, err := c.ToMapStr()
+	assert.NoError(t, err)
+	assert.Equal(t, in, out)
+}