Merge pull request #63 from ShiftLeftSecurity/chipaca/escape-args-escape-question-mark

db/connection: support escaping placeholders

Author: chipaca

db/connection/connection.go

@@ -211,20 +211,34 @@ func (f *FlexibleTransaction) RollbackTransaction(ctx context.Context) error {
 }
 
 // EscapeArgs return the query and args with the argument placeholder escaped.
+//
+// The argument placeholder is `?`. If you need an actual `?` in the output, you
+// can input `\?`.If you need an actual `\` in the output, input `\\`.
 func EscapeArgs(query string, args []interface{}) (string, []interface{}, error) {
 	// TODO: make this a bit less ugly
-	// TODO: identify escaped question marks
 	queryWithArgs := &strings.Builder{}
 	argCounter := 1
+	escaped := false
 	for _, queryChar := range query {
-		if queryChar == '?' {
-			queryWithArgs.WriteRune('$')
-			queryWithArgs.WriteString(strconv.Itoa(argCounter))
-			argCounter++
-		} else {
+		if escaped {
 			queryWithArgs.WriteRune(queryChar)
+			escaped = false
+		} else {
+			switch queryChar {
+			case '\\':
+				escaped = true
+			case '?':
+				queryWithArgs.WriteRune('$')
+				queryWithArgs.WriteString(strconv.Itoa(argCounter))
+				argCounter++
+			default:
+				queryWithArgs.WriteRune(queryChar)
+			}
 		}
 	}
+	if escaped {
+		return "", nil, errors.New("the query ends with an escape")
+	}
 	if len(args) != argCounter-1 {
 		return "", nil, errors.Errorf("the query has %d args but %d were passed: \n %q \n %#v",
 			argCounter-1, len(args), queryWithArgs, args)

db/connection/connection_test.go

@@ -3,6 +3,8 @@ package connection
 import (
 	"context"
 	"testing"
+
+	"github.com/go-test/deep"
 )
 
 type fakeConn struct {
@@ -188,3 +190,25 @@ func TestFlexibleTransactionRecursive(t *testing.T) {
 		t.FailNow()
 	}
 }
+
+func TestEscapeArgsOK(t *testing.T) {
+	for in, out := range map[string]string{
+		"from ? where ?=?":     "from $1 where $2=$3",
+		"from ? where ? \\? ?": "from $1 where $2 ? $3",
+		`\\??\??`:              `\$1$2?$3`,
+	} {
+		t.Run("", func(t *testing.T) {
+			args := []interface{}{"hello", 1, 42.}
+			got, gotArgs, err := EscapeArgs(in, args)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if diff := deep.Equal(args, gotArgs); diff != nil {
+				t.Fatal(err)
+			}
+			if got != out {
+				t.Errorf("expected %q, got %q", out, got)
+			}
+		})
+	}
+}

go.mod

@@ -3,6 +3,7 @@ module github.com/ShiftLeftSecurity/gaum/v2
 go 1.15
 
 require (
+	github.com/go-test/deep v1.0.8
 	github.com/jackc/pgconn v1.8.1
 	github.com/jackc/pgproto3/v2 v2.0.7 // indirect
 	github.com/jackc/pgx/v4 v4.11.0

go.sum

@@ -67,6 +67,8 @@ github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V
 github.com/go-logfmt/logfmt v0.5.0/go.mod h1:wCYkCAKZfumFQihp8CzCvQ3paCTfi41vtzG1KdI/P7A=
 github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
+github.com/go-test/deep v1.0.8 h1:TDsG77qcSprGbC6vTN8OuXp5g+J+b5Pcguhf7Zt61VM=
+github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE=
 github.com/gofrs/uuid v3.2.0+incompatible h1:y12jRkkFxsd7GpqdSZ+/KCs/fJbqpEXSGd4+jfEaewE=
 github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gogo/googleapis v1.1.0/go.mod h1:gf4bu3Q80BeJ6H1S1vYPm8/ELATdvryBaNFGgqEef3s=