Made it possible to implement driver.Valuer for byte slices

sjw1357 · Feb 11, 2018 · 3b6d790 · 3b6d790
1 parent fd15156
commit 3b6d790
Show file tree

Hide file tree

Showing 3 changed files with 128 additions and 32 deletions.
diff --git a/migration_test.go b/migration_test.go
@@ -33,6 +33,7 @@ type User struct {
 	CompanyID         *int
 	Company           Company
 	Role              Role
+	Password          EncryptedData
 	PasswordHash      []byte
 	IgnoreMe          int64                 `sql:"-"`
 	IgnoreStringSlice []string              `sql:"-"`
@@ -116,6 +117,31 @@ type Company struct {
 	Owner *User `sql:"-"`
 }
 
+type EncryptedData []byte
+
+func (data *EncryptedData) Scan(value interface{}) error {
+	if b, ok := value.([]byte); ok {
+		if len(b) < 3 || b[0] != '*' || b[1] != '*' || b[2] != '*'{
+			return errors.New("Too short")
+		}
+
+		*data = b[3:]
+		return nil
+	} else {
+		return errors.New("Bytes expected")
+	}
+}
+
+func (data EncryptedData) Value() (driver.Value, error) {
+	if len(data) > 0 && data[0] == 'x' {
+		//needed to test failures
+		return nil, errors.New("Should not start with 'x'")
+	}
+
+	//prepend asterisks
+	return append([]byte("***"), data...), nil
+}
+
 type Role struct {
 	Name string `gorm:"size:256"`
 }

diff --git a/scope.go b/scope.go
@@ -557,22 +557,29 @@ func (scope *Scope) buildWhereCondition(clause map[string]interface{}) (str stri
 
 	args := clause["args"].([]interface{})
 	for _, arg := range args {
-		switch reflect.ValueOf(arg).Kind() {
-		case reflect.Slice: // For where("id in (?)", []int64{1,2})
-			if bytes, ok := arg.([]byte); ok {
-				str = strings.Replace(str, "?", scope.AddToVars(bytes), 1)
-			} else if values := reflect.ValueOf(arg); values.Len() > 0 {
-				var tempMarks []string
-				for i := 0; i < values.Len(); i++ {
-					tempMarks = append(tempMarks, scope.AddToVars(values.Index(i).Interface()))
+		rArg := reflect.ValueOf(arg)
+		rArgType := reflect.TypeOf(arg)
+		vArg, isValuer := arg.(driver.Valuer)
+		var err error
+
+		//non byte slice and non driver.Valuer
+		if rArgType.Kind() == reflect.Slice && rArgType.Elem().Kind() != reflect.Uint8 && !isValuer {
+			if rArg.Len() > 0 {
+				tempMarks := make([]string, 0, rArg.Len())
+				for i := 0; i < rArg.Len(); i++ {
+					tempMarks = append(tempMarks, scope.AddToVars(rArg.Index(i).Interface()))
 				}
+
 				str = strings.Replace(str, "?", strings.Join(tempMarks, ","), 1)
 			} else {
 				str = strings.Replace(str, "?", scope.AddToVars(Expr("NULL")), 1)
 			}
-		default:
-			if valuer, ok := interface{}(arg).(driver.Valuer); ok {
-				arg, _ = valuer.Value()
+		} else {
+			if isValuer {
+				arg, err = vArg.Value()
+				if err != nil {
+					scope.Err(err)
+				}
 			}
 
 			str = strings.Replace(str, "?", scope.AddToVars(arg), 1)
@@ -629,23 +636,31 @@ func (scope *Scope) buildNotCondition(clause map[string]interface{}) (str string
 
 	args := clause["args"].([]interface{})
 	for _, arg := range args {
-		switch reflect.ValueOf(arg).Kind() {
-		case reflect.Slice: // For where("id in (?)", []int64{1,2})
-			if bytes, ok := arg.([]byte); ok {
-				str = strings.Replace(str, "?", scope.AddToVars(bytes), 1)
-			} else if values := reflect.ValueOf(arg); values.Len() > 0 {
-				var tempMarks []string
-				for i := 0; i < values.Len(); i++ {
-					tempMarks = append(tempMarks, scope.AddToVars(values.Index(i).Interface()))
+		rArg := reflect.ValueOf(arg)
+		rArgType := reflect.TypeOf(arg)
+		vArg, isValuer := arg.(driver.Valuer)
+		var err error
+
+		//non byte slice and non driver.Valuer
+		if rArgType.Kind() == reflect.Slice && rArgType.Elem().Kind() != reflect.Uint8 && !isValuer {
+			if rArg.Len() > 0 {
+				tempMarks := make([]string, 0, rArg.Len())
+				for i := 0; i < rArg.Len(); i++ {
+					tempMarks = append(tempMarks, scope.AddToVars(rArg.Index(i).Interface()))
 				}
+
 				str = strings.Replace(str, "?", strings.Join(tempMarks, ","), 1)
 			} else {
 				str = strings.Replace(str, "?", scope.AddToVars(Expr("NULL")), 1)
 			}
-		default:
-			if scanner, ok := interface{}(arg).(driver.Valuer); ok {
-				arg, _ = scanner.Value()
+		} else {
+			if isValuer {
+				arg, err = vArg.Value()
+				if err != nil {
+					scope.Err(err)
+				}
 			}
+
 			str = strings.Replace(notEqualSQL, "?", scope.AddToVars(arg), 1)
 		}
 	}
@@ -662,18 +677,31 @@ func (scope *Scope) buildSelectQuery(clause map[string]interface{}) (str string)
 
 	args := clause["args"].([]interface{})
 	for _, arg := range args {
-		switch reflect.ValueOf(arg).Kind() {
-		case reflect.Slice:
-			values := reflect.ValueOf(arg)
-			var tempMarks []string
-			for i := 0; i < values.Len(); i++ {
-				tempMarks = append(tempMarks, scope.AddToVars(values.Index(i).Interface()))
+		rArg := reflect.ValueOf(arg)
+		rArgType := reflect.TypeOf(arg)
+		vArg, isValuer := arg.(driver.Valuer)
+		var err error
+
+		//non byte slice and non driver.Valuer
+		if rArgType.Kind() == reflect.Slice && rArgType.Elem().Kind() != reflect.Uint8 && !isValuer {
+			if rArg.Len() > 0 {
+				tempMarks := make([]string, 0, rArg.Len())
+				for i := 0; i < rArg.Len(); i++ {
+					tempMarks = append(tempMarks, scope.AddToVars(rArg.Index(i).Interface()))
+				}
+
+				str = strings.Replace(str, "?", strings.Join(tempMarks, ","), 1)
+			} else {
+				str = strings.Replace(str, "?", scope.AddToVars(Expr("NULL")), 1)
 			}
-			str = strings.Replace(str, "?", strings.Join(tempMarks, ","), 1)
-		default:
-			if valuer, ok := interface{}(arg).(driver.Valuer); ok {
-				arg, _ = valuer.Value()
+		} else {
+			if isValuer {
+				arg, err = vArg.Value()
+				if err != nil {
+					scope.Err(err)
+				}
 			}
+
 			str = strings.Replace(str, "?", scope.AddToVars(arg), 1)
 		}
 	}

diff --git a/scope_test.go b/scope_test.go
@@ -1,7 +1,10 @@
 package gorm_test
 
 import (
+	"encoding/hex"
 	"github.com/jinzhu/gorm"
+	"math/rand"
+	"strings"
 	"testing"
 )
 
@@ -41,3 +44,42 @@ func TestScopes(t *testing.T) {
 		t.Errorf("Should found two users's name in 1, 3")
 	}
 }
+
+func randName() string {
+	data := make([]byte, 8)
+	rand.Read(data)
+
+	return "n-" + hex.EncodeToString(data)
+}
+
+func TestValuer(t *testing.T) {
+	name := randName()
+
+	origUser := User{Name: name, Age: 1, Password: EncryptedData("pass1"), PasswordHash: []byte("abc")}
+	err := DB.Save(&origUser).Error
+	if err != nil {
+		t.Log(err)
+		t.FailNow()
+	}
+
+	var user2 User
+	err = DB.Where("name=? AND password=? AND password_hash=?", name, EncryptedData("pass1"), []byte("abc")).First(&user2).Error
+	if err != nil {
+		t.Log(err)
+		t.FailNow()
+	}
+}
+
+func TestFailedValuer(t *testing.T) {
+	name := randName()
+
+	err := DB.Exec("INSERT INTO users(name, password) VALUES(?, ?)", name, EncryptedData("xpass1")).Error
+	if err == nil {
+		t.FailNow()
+	}
+
+	if !strings.HasPrefix(err.Error(), "Should not start with") {
+		t.FailNow()
+	}
+
+}