Merge pull request #51 from ShiftLeftSecurity/preetam/support-escaped-placeholders

Support escaped '?', which may be JSONB operators

Author: Preetam

CONTRIBUTORS

@@ -1,3 +1,4 @@
 Horacio Duran <[email protected]>
 William Cody Laeder <[email protected]>
-Alex Hornbake <[email protected]>
+Alex Hornbake <[email protected]>
+Preetam Jinka <[email protected]>

db/chain/placeholders.go

@@ -76,14 +76,25 @@ func MarksToPlaceholders(q string, args []interface{}) (string, []interface{}, e
 	args = otherArgs
 
 	// TODO: make this a bit less ugly
-	// TODO: identify escaped question marks
 	// TODO: use an actual parser <3
 	// TODO: structure query segments around SQL-Standard AST
 	queryWithArgs := &strings.Builder{}
 	argCounter := 1
 	argPositioner := 0
 	expandedArgs := []interface{}{}
-	for _, queryChar := range q {
+	skip := false
+	for i, queryChar := range q {
+		if skip {
+			skip = false
+			continue
+		}
+
+		if queryChar == '\\' && i < len(q)-1 && q[i+1] == '?' {
+			// Escaped '?'
+			queryWithArgs.WriteRune('?')
+			skip = true
+			continue
+		}
 		if queryChar == '?' {
 			arg := args[argPositioner]
 			switch reflect.TypeOf(arg).Kind() {
@@ -126,7 +137,6 @@ func MarksToPlaceholders(q string, args []interface{}) (string, []interface{}, e
 
 // PlaceholdersToPositional converts ? in a query into $<argument number> which postgres expects
 func PlaceholdersToPositional(q *strings.Builder, argCount int) (*strings.Builder, int, error) {
-	// TODO: identify escaped question marks
 	// TODO: use an actual parser <3
 	// TODO: structure query segments around SQL-Standard AST
 	newQ := &strings.Builder{}
@@ -136,8 +146,22 @@ func PlaceholdersToPositional(q *strings.Builder, argCount int) (*strings.Builde
 		newQ.Grow(renderedLength - newQ.Len())
 	}
 
+	queryString := q.String()
 	argCounter := 1
-	for _, queryChar := range q.String() {
+	skip := false
+	for i, queryChar := range queryString {
+		if skip {
+			skip = false
+			continue
+		}
+
+		if queryChar == '\\' && i < len(queryString)-1 && queryString[i+1] == '?' {
+			// Escaped '?'
+			newQ.WriteRune('?')
+			skip = true
+			continue
+		}
+
 		if queryChar == '?' {
 			newQ.WriteRune('$')
 			newQ.WriteString(strconv.Itoa(argCounter))

db/chain/placeholders_test.go

@@ -1,6 +1,9 @@
 package chain
 
-import "testing"
+import (
+	"fmt"
+	"testing"
+)
 
 func Test_digitSize(t *testing.T) {
 	type args struct {
@@ -56,3 +59,45 @@ func Test_digitSize(t *testing.T) {
 		})
 	}
 }
+
+func TestPlaceholderEscaping(t *testing.T) {
+	tests := []struct {
+		q    string
+		want string
+		args []interface{}
+	}{
+		{
+			q:    "? = 1",
+			want: "$1 = 1",
+			args: []interface{}{1},
+		},
+		{
+			q:    "\\? = 1",
+			want: "? = 1",
+			args: []interface{}{},
+		},
+		{
+			q:    "? = ? AND \\? = 1",
+			want: "$1 = $2 AND ? = 1",
+			args: []interface{}{1, 1},
+		},
+		{
+			q:    `'["a", "b"]'::jsonb \?& array['a', 'b']`,
+			want: `'["a", "b"]'::jsonb ?& array['a', 'b']`,
+			args: []interface{}{},
+		},
+		{
+			q:    `'["a", "b"]'::jsonb \?& array[?]`,
+			want: `'["a", "b"]'::jsonb ?& array[$1]`,
+			args: []interface{}{"a"},
+		},
+	}
+	for i, tt := range tests {
+		t.Run(fmt.Sprint(i), func(t *testing.T) {
+			result, _, _ := MarksToPlaceholders(tt.q, tt.args)
+			if result != tt.want {
+				t.Errorf("got %v, want %v", result, tt.want)
+			}
+		})
+	}
+}