From 3b6d790e93e9715cafbc66179f9435994e7413a2 Mon Sep 17 00:00:00 2001 From: Viktor Nikolaiev Date: Wed, 30 Aug 2017 22:52:45 +0300 Subject: [PATCH] Made it possible to implement driver.Valuer for byte slices --- migration_test.go | 26 ++++++++++++++ scope.go | 92 ++++++++++++++++++++++++++++++----------------- scope_test.go | 42 ++++++++++++++++++++++ 3 files changed, 128 insertions(+), 32 deletions(-) diff --git a/migration_test.go b/migration_test.go index 6b4470a6..7c3436ca 100644 --- 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 index ae98d251..65d35461 100644 --- 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 index 42458995..71e80225 100644 --- 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() + } + +}