From 6b5187dc644de47f8cf7cf31e28d65324f1717ee Mon Sep 17 00:00:00 2001
From: Zhengda Lu <zhengda.lu@datadoghq.com>
Date: Mon, 16 Dec 2024 15:49:22 -0500
Subject: [PATCH] option to obfuscate bind parameter

---
 obfuscator.go                                 | 13 ++++++++++++
 obfuscator_test.go                            | 12 +++++++++++
 sqllexer_test.go                              | 21 +++++++++++++++++++
 sqllexer_utils.go                             |  2 +-
 .../select/select-with-bind-parameter.json    | 17 +++++++++++++++
 5 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 testdata/mssql/select/select-with-bind-parameter.json

diff --git a/obfuscator.go b/obfuscator.go
index 79a6719..44ba513 100644
--- a/obfuscator.go
+++ b/obfuscator.go
@@ -11,6 +11,7 @@ type obfuscatorConfig struct {
 	ReplaceBoolean             bool `json:"replace_boolean"`
 	ReplaceNull                bool `json:"replace_null"`
 	KeepJsonPath               bool `json:"keep_json_path"` // by default, we replace json path with placeholder
+	ReplaceBindParameter       bool `json:"replace_bind_parameter"`
 }
 
 type obfuscatorOption func(*obfuscatorConfig)
@@ -51,6 +52,12 @@ func WithKeepJsonPath(keepJsonPath bool) obfuscatorOption {
 	}
 }
 
+func WithReplaceBindParameter(replaceBindParameter bool) obfuscatorOption {
+	return func(c *obfuscatorConfig) {
+		c.ReplaceBindParameter = replaceBindParameter
+	}
+}
+
 type Obfuscator struct {
 	config *obfuscatorConfig
 }
@@ -128,6 +135,12 @@ func (o *Obfuscator) ObfuscateTokenValue(token Token, lastToken Token, lexerOpts
 		} else {
 			return token.Value
 		}
+	case BIND_PARAMETER:
+		if o.config.ReplaceBindParameter {
+			return StringPlaceholder
+		} else {
+			return token.Value
+		}
 	case IDENT, QUOTED_IDENT:
 		if o.config.ReplaceBoolean && isBoolean(token.Value) {
 			return StringPlaceholder
diff --git a/obfuscator_test.go b/obfuscator_test.go
index c8452fd..b9a0103 100644
--- a/obfuscator_test.go
+++ b/obfuscator_test.go
@@ -17,6 +17,7 @@ func TestObfuscator(t *testing.T) {
 		replaceNull                bool
 		dollarQuotedFunc           bool
 		keepJsonPath               bool
+		replaceBindParameter       bool
 		dbms                       DBMSType
 	}{
 		{
@@ -535,6 +536,16 @@ func TestObfuscator(t *testing.T) {
 			expected:     `SELECT * FROM users where data::jsonb ->> 1`,
 			keepJsonPath: true,
 		},
+		{
+			input:                `SELECT * FROM users where id = @_My_id`,
+			expected:             `SELECT * FROM users where id = @_My_id`,
+			replaceBindParameter: false,
+		},
+		{
+			input:                `SELECT * FROM users where id = @_My_id`,
+			expected:             `SELECT * FROM users where id = ?`,
+			replaceBindParameter: true,
+		},
 	}
 
 	for _, tt := range tests {
@@ -546,6 +557,7 @@ func TestObfuscator(t *testing.T) {
 				WithReplaceNull(tt.replaceNull),
 				WithDollarQuotedFunc(tt.dollarQuotedFunc),
 				WithKeepJsonPath(tt.keepJsonPath),
+				WithReplaceBindParameter(tt.replaceBindParameter),
 			)
 			got := obfuscator.Obfuscate(tt.input, WithDBMS(tt.dbms))
 			assert.Equal(t, tt.expected, got)
diff --git a/sqllexer_test.go b/sqllexer_test.go
index af3b93e..3ed42b1 100644
--- a/sqllexer_test.go
+++ b/sqllexer_test.go
@@ -529,6 +529,27 @@ func TestLexer(t *testing.T) {
 				{BIND_PARAMETER, "@1"},
 			},
 		},
+		{
+			name:  "select with bind parameter using underscore",
+			input: "SELECT * FROM users where id = @__my_id",
+			expected: []Token{
+				{IDENT, "SELECT"},
+				{WS, " "},
+				{WILDCARD, "*"},
+				{WS, " "},
+				{IDENT, "FROM"},
+				{WS, " "},
+				{IDENT, "users"},
+				{WS, " "},
+				{IDENT, "where"},
+				{WS, " "},
+				{IDENT, "id"},
+				{WS, " "},
+				{OPERATOR, "="},
+				{WS, " "},
+				{BIND_PARAMETER, "@__my_id"},
+			},
+		},
 		{
 			name:  "select with system variable",
 			input: "SELECT @@VERSION AS SqlServerVersion",
diff --git a/sqllexer_utils.go b/sqllexer_utils.go
index bbef75a..d9ede92 100644
--- a/sqllexer_utils.go
+++ b/sqllexer_utils.go
@@ -186,7 +186,7 @@ func isLetter(ch rune) bool {
 }
 
 func isAlphaNumeric(ch rune) bool {
-	return isLetter(ch) || isDigit(ch)
+	return isLetter(ch) || isDigit(ch) || ch == '_'
 }
 
 func isDoubleQuote(ch rune) bool {
diff --git a/testdata/mssql/select/select-with-bind-parameter.json b/testdata/mssql/select/select-with-bind-parameter.json
new file mode 100644
index 0000000..5fc7cf1
--- /dev/null
+++ b/testdata/mssql/select/select-with-bind-parameter.json
@@ -0,0 +1,17 @@
+{
+    "input": "SELECT[u].[USER_SETTINGS_ID], [u].[CREATED], [u].[CREATED_BY], [u].[FK_APPLICATION_SCOPE_ID], [u].[MODIFIED], [u].[MODIFIED_BY], [u].[PREFERRED_LANGUAGE], [u].[SETTINGS], [u].[USER_NAME], [u].[USER_SECRET], [u].[USER_SECRET_TMP] FROM [USER_SETTING] WHERE [u].[USER_NAME] = @__userName_0;",
+    "outputs": [
+      {
+        "expected": "SELECT [u].[USER_SETTINGS_ID], [u].[CREATED], [u].[CREATED_BY], [u].[FK_APPLICATION_SCOPE_ID], [u].[MODIFIED], [u].[MODIFIED_BY], [u].[PREFERRED_LANGUAGE], [u].[SETTINGS], [u].[USER_NAME], [u].[USER_SECRET], [u].[USER_SECRET_TMP] FROM [USER_SETTING] WHERE [u].[USER_NAME] = @__userName_0",
+        "obfuscator_config": {
+            "replace_positional_parameter": false
+        },
+        "normalizer_config": {
+            "keep_sql_alias": true,
+            "remove_space_between_parentheses": true,
+            "keep_identifier_quotation": true
+        }
+      }
+    ]
+  }
+  
\ No newline at end of file