Merge pull request #110 from appoptics/AO-7048-sanitization-plus

[SQLSanitizationPlus] Sanitize more literal formats
solarwinds · Jun 24, 2019 · 3f5f9f8 · 3f5f9f8
2 parents 2b28157 + d899047
commit 3f5f9f8
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 6 deletions.
diff --git a/v1/ao/internal/reporter/sql_sanitizer.go b/v1/ao/internal/reporter/sql_sanitizer.go
@@ -143,6 +143,11 @@ func NewSQLSanitizer(dbType string, sanitizeFlag int) *SQLSanitizer {
 		sanitizer.identifierQuotes['['] = ']'
 	}
 
+	// For double-dollar quoted literals: $tag$literal$tag$
+	if dbType == PostgreSQL {
+		sanitizer.literalQuotes['$'] = '$'
+	}
+
 	return &sanitizer
 }
 
@@ -191,7 +196,11 @@ func (s *SQLSanitizer) Sanitize(sql string) string {
 		return runeCnt
 	}
 
-	for _, currRune := range sql {
+	// For PostgreSQL's dollar-quoted literal only.
+	// see https://www.postgresql.org/docs/9.0/sql-syntax-lexical.html
+	tag := []rune{'$'}
+
+	for i, currRune := range sql {
 		if StackSize() >= maxRuneCnt {
 			StackPush(Ellipsis)
 			break
@@ -203,6 +212,19 @@ func (s *SQLSanitizer) Sanitize(sql string) string {
 
 		switch currState {
 		case FSMStringStart:
+			// Handle PostgreSQL's double-dollar quoted literal
+			if s.dbType == PostgreSQL && closingQuote == '$' {
+				if currRune == '$' {
+					currState = FSMStringBody
+					StackPush(ReplacementRune)
+				}
+				// Record the tag: '$tag$'. It will be used to compare with the
+				// literal when a '$' is seen to identify the end of string.
+				tag = append(tag, currRune)
+
+				break // break out of switch
+			}
+
 			StackPush(ReplacementRune)
 
 			if currRune == closingQuote {
@@ -215,15 +237,40 @@ func (s *SQLSanitizer) Sanitize(sql string) string {
 
 		case FSMStringBody:
 			if currRune == closingQuote {
-				currState = FSMStringEnd
-			} else if currRune == EscapeRune {
+				if s.dbType == PostgreSQL &&
+					closingQuote == '$' &&
+					i+len(tag) <= utf8.RuneCountInString(sql) && // slice bounds check
+					strings.Compare(string(tag), string([]rune(sql)[i:i+len(tag)])) != 0 {
+					// Do nothing - It's only a '$' inside a string, rather than
+					// the literal end. The nested dollar-quoted string is not handled
+					// specially as they are all replaced with the placeholder.
+				} else {
+					currState = FSMStringEnd
+					if s.dbType == PostgreSQL && closingQuote == '$' {
+						tag = []rune{'$'} // reset the tag
+					}
+				}
+			} else if currRune == EscapeRune && closingQuote != '$' {
+				// The escape rune, e.g., a '\', should not be escaped inside a
+				// dollar-quoted string of PostgreSQL. However, it doesn't have
+				// to be handle specially either as the entire string body will
+				// be removed (replaced with the placeholder) anyways.
 				currState = FSMStringEscape
 			}
 
 		case FSMStringEscape:
 			currState = FSMStringBody
 
 		case FSMStringEnd:
+			// Handle PostgreSQL's double-dollar quoted literal
+			if s.dbType == PostgreSQL && closingQuote == '$' {
+				if currRune == '$' {
+					// The real end of closing quote
+					currState = FSMCopy
+				}
+				break // break out of switch
+			}
+
 			if currRune == closingQuote {
 				currState = FSMStringBody
 			} else {
@@ -251,9 +298,11 @@ func (s *SQLSanitizer) Sanitize(sql string) string {
 		case FSMIdentifier:
 			if c, ok := s.literalQuotes[currRune]; ok {
 				// PostgreSQL has literals like X'FEFF' or U&'\0441'
-				if StackPeek(1) == 'X' {
+				top1 := unicode.ToUpper(StackPeek(1))
+				top2 := unicode.ToUpper(StackPeek(2))
+				if top1 == 'X' || top1 == 'B' || top1 == 'U' || top1 == 'N' {
 					StackPop()
-				} else if StackPeek(1) == '&' && StackPeek(2) == 'U' {
+				} else if top1 == '&' && top2 == 'U' {
 					StackPop()
 					StackPop()
 				}

diff --git a/v1/ao/internal/reporter/sql_sanitizer_test.go b/v1/ao/internal/reporter/sql_sanitizer_test.go
@@ -373,6 +373,54 @@ func TestSQLSanitize(t *testing.T) {
 			"SELECT",
 			"SELECT",
 		},
+		{ // PostgreSQL's special literal formats
+			EnabledAuto,
+			PostgreSQL,
+			`SELECT * FROM employees WHERE name = B'0101'`,
+			`SELECT * FROM employees WHERE name = ?`,
+		},
+		{ // PostgreSQL's double-dollar quoted literal with optional tag
+			EnabledAuto,
+			PostgreSQL,
+			`SELECT * FROM employees WHERE name = $tag$Eric$tag$`,
+			`SELECT * FROM employees WHERE name = ?`,
+		},
+		{ // PostgreSQL's double-dollar quoted literal with optional tag and dollar in string
+			EnabledAuto,
+			PostgreSQL,
+			`SELECT * FROM employees WHERE name = $tag$Eric spends $99 '$tag$`,
+			`SELECT * FROM employees WHERE name = ?`,
+		},
+		{ // PostgreSQL's double-dollar quoted literal
+			EnabledAuto,
+			PostgreSQL,
+			`SELECT * FROM employees WHERE name = $tag\$Eric spends $99 '\$tag\$ AND gender = $tag$male$tag$`,
+			`SELECT * FROM employees WHERE name = ? AND gender = ?`,
+		},
+		{ // PostgreSQL's double-dollar quoted literal
+			EnabledAuto,
+			PostgreSQL,
+			`SELECT * FROM employees WHERE name = $$Eric$$`,
+			`SELECT * FROM employees WHERE name = ?`,
+		},
+		{ // PostgreSQL's double-dollar quoted literal with escape rune
+			EnabledAuto,
+			PostgreSQL,
+			`SELECT * FROM employees WHERE name = $$Eric$$$`,
+			`SELECT * FROM employees WHERE name = ?`,
+		},
+		{ // Oracle's special treatment when literal replacement is turned on
+			EnabledAuto,
+			Oracle,
+			`SELECT * FROM employees WHERE name = U'500 Oracle Parkway'`,
+			`SELECT * FROM employees WHERE name = ?`,
+		},
+		{ // Oracle's N function
+			EnabledAuto,
+			Oracle,
+			`SELECT * FROM employees WHERE name = N'500 Oracle Parkway'`,
+			`SELECT * FROM employees WHERE name = ?`,
+		},
 	}
 
 	for _, c := range cases {

diff --git a/v1/ao/internal/utils/version.go b/v1/ao/internal/utils/version.go
@@ -7,7 +7,7 @@ import (
 
 var (
 	// The AppOptics Go agent version
-	version = "1.8.0"
+	version = "1.8.1"
 
 	// The Go version
 	goVersion = strings.TrimPrefix(runtime.Version(), "go")