Made it possible to implement driver.Valuer for byte slices

This commit is contained in:
Viktor Nikolaiev 2017-08-30 22:52:45 +03:00 committed by Jinzhu
parent fd15156d39
commit 3b6d790e93
3 changed files with 128 additions and 32 deletions

View File

@ -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"`
}

View File

@ -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)
}
}

View File

@ -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()
}
}