fixed SQL Injection for #2

Author: yangyin5127

sql.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"reflect"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -16,13 +17,41 @@ var (
 	escaper            = "'"
 	nullStr            = "NULL"
 	singleQuoteEscaper = "\\"
+	escapeRegexp       = regexp.MustCompile(`[\0\t\x1a\n\r\"\'\\]`)
+
+	//see href='https://dev.mysql.com/doc/refman/8.0/en/string-literals.html#character-escape-sequences'
+	characterEscapeMap = map[string]string{
+		"\\0":  `\\0`,  //ASCII NULL
+		"\b":   `\\b`,  //backspace
+		"\t":   `\\t`,  //tab
+		"\x1a": `\\Z`,  //ASCII 26 (Control+Z);
+		"\n":   `\\n`,  //newline character
+		"\r":   `\\r`,  //return character
+		"\"":   `\\"`,  //quote (")
+		"'":    `\'`,   //quote (')
+		"\\":   `\\\\`, //backslash (\)
+		// "\\%":  `\\%`,  //% character
+		// "\\_":  `\\_`,  //_ character
+	}
 )
 
 //Escape escape the val for sql
 func Escape(val interface{}) string {
 	return EscapeInLocation(val, time.Local)
 }
 
+//toSqlString escape the string val for sql
+func toSqlString(val string) string {
+	return escapeRegexp.ReplaceAllStringFunc(val, func(s string) string {
+
+		mVal, ok := characterEscapeMap[s]
+		if ok {
+			return mVal
+		}
+		return s
+	})
+}
+
 func timeToString(t time.Time, loc *time.Location) string {
 	if t.IsZero() {
 		return escaper + tmFmtZero + escaper
@@ -70,7 +99,7 @@ func EscapeInLocation(val interface{}, loc *time.Location) string {
 		return fmt.Sprintf("%.6f", v)
 
 	case string:
-		return escaper + strings.Replace(v, escaper, singleQuoteEscaper+escaper, -1) + escaper
+		return escaper + toSqlString(v) + escaper
 	default:
 		refValue := reflect.ValueOf(v)
 		if v == nil || !refValue.IsValid() {
@@ -94,7 +123,7 @@ func EscapeInLocation(val interface{}, loc *time.Location) string {
 		if err != nil {
 			return nullStr
 		}
-		return escaper + strings.Replace(string(stringifyData), escaper, singleQuoteEscaper+escaper, -1) + escaper
+		return escaper + toSqlString(string(stringifyData)) + escaper
 
 	}
 }
@@ -144,6 +173,9 @@ func FormatInLocation(query string, loc *time.Location, args ...interface{}) str
 }
 
 //SetSingleQuoteEscaper set the singleQuoteEscaper
+//default:\' , e.g. '' 、 \'
 func SetSingleQuoteEscaper(escaper string) {
-	singleQuoteEscaper = escaper
+
+	characterEscapeMap["'"] = escaper
+	// singleQuoteEscaper = escaper
 }

sql_test.go

@@ -13,6 +13,14 @@ func TestNULLEscape(t *testing.T) {
 	}
 }
 
+func Test0Escape(t *testing.T) {
+	result := Escape(`\0`)
+	t.Logf("Test0Escape result: %s", result)
+	if result != `'\\\\0'` {
+		t.Fatalf("escape error")
+	}
+}
+
 func TestEmptyStringEscape(t *testing.T) {
 	result := Escape("")
 	t.Logf("result :%s", result)
@@ -87,9 +95,27 @@ func TestStringEscape(t *testing.T) {
 	}
 }
 
+func TestStringEscape2(t *testing.T) {
+	s := "hello world"
+	result := Escape(s)
+	if result != "'hello world'" {
+		t.Fatalf("escape string error")
+
+	}
+
+	s = `hello \' world`
+	t.Logf("TestStringEscape2 raw:%s", s)
+	result = Escape(s)
+	t.Logf("TestStringEscape2 result: %s", result)
+	if result != `'hello \\\\\' world'` {
+		t.Fatalf("escape string error")
+
+	}
+}
+
 func TestStringCustomEscape(t *testing.T) {
 	s := "hello world"
-	SetSingleQuoteEscaper("'")
+	SetSingleQuoteEscaper("''")
 	result := Escape(s)
 	if result != "'hello world'" {
 		t.Fatalf("escape string error")
@@ -103,6 +129,8 @@ func TestStringCustomEscape(t *testing.T) {
 		t.Fatalf("escape string error")
 
 	}
+	SetSingleQuoteEscaper("\\'")
+
 }
 
 func TestBytesEscape(t *testing.T) {
@@ -210,13 +238,85 @@ func TestOtherEscape(t *testing.T) {
 	result := Escape(x)
 	t.Logf("escape reuslt %s", result)
 
-	if result != "'{\"key\":\"test\",\"name\":\"asd\\'fsadf\"}'" {
+	if result != `'{\\"key\\":\\"test\\",\\"name\\":\\"asd\'fsadf\\"}'` {
 		t.Fatalf("escape map error")
 
 	}
 
 }
 
+func TestNewlineEscape(t *testing.T) {
+	s := "hello\nworld"
+	result := Escape(s)
+	t.Logf("escape newline reuslt: %s", result)
+
+	if result != "'hello\\\\nworld'" {
+		t.Fatalf("escape string error")
+
+	}
+
+}
+
+func TestReturnEscape(t *testing.T) {
+	s := "hello\rworld"
+	result := Escape(s)
+	t.Logf("escape newline reuslt: %s", result)
+
+	if result != "'hello\\\\rworld'" {
+		t.Fatalf("escape string error")
+
+	}
+
+}
+
+func TestTabEscape(t *testing.T) {
+	s := "hello\tworld"
+	result := Escape(s)
+	t.Logf("escape tab reuslt: %s", result)
+
+	if result != `'hello\\tworld'` {
+		t.Fatalf("escape string error")
+
+	}
+
+}
+
+func TestDoubleBackslashEscape(t *testing.T) {
+	s := "hello\\world"
+	result := Escape(s)
+	t.Logf("escape tab reuslt: %s", result)
+
+	if result != `'hello\\\\world'` {
+		t.Fatalf("escape string error")
+
+	}
+
+}
+
+func TestCtrlZEscape(t *testing.T) {
+	s := "hello\x1aworld"
+	result := Escape(s)
+	t.Logf("escape tab reuslt: %s", result)
+
+	if result != `'hello\\Zworld'` {
+		t.Fatalf("escape string error")
+
+	}
+
+}
+
+func TestDoubleQouteEscape(t *testing.T) {
+	s := "hello \" world"
+	result := Escape(s)
+	t.Logf("escape tab reuslt: %s", result)
+
+	if result != `'hello \\" world'` {
+		t.Fatalf("escape string error")
+
+	}
+
+}
+
 func TestFormatSql(t *testing.T) {
 
 	sql := Format("select * from users where name=? and age=? limit ?,?", "t'est", 10, 10, 10)
@@ -282,4 +382,11 @@ func TestFormatSql(t *testing.T) {
 		t.Fatalf("escape format time error")
 
 	}
+
+	sql = Format("select * from users where name=? and age=? limit ?,?", `t\'est`, 10, 10, 10)
+
+	if sql != `'select * from users where name='t\\\\\'est' and age=10 limit 10,10'` {
+
+		t.Logf("sql: %s\n", sql)
+	}
 }