database/tidb
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

expression: implement vectorized evaluation for builtinAesEncryptIVSig (#13521)

Browse Source
This commit is contained in:
Zhang Zhiyi
2019-11-18 22:00:24 +08:00
committed by pingcap-github-bot
parent 768ec30a88
commit 871c19bb46
3 changed files with 103 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
expression/bench_test.go
View File

@ -693,6 +693,8 @@ type vecExprBenchCase struct {
// geners[gen1, gen2] will be regarded as geners[gen1, gen2, nil].
// This field is optional.
geners []dataGenerator
// aesModeAttr information, needed by encryption functions
aesModes string
// constants are used to generate constant data for children[i].
constants []*Constant
}
@ -1002,10 +1004,8 @@ func testVectorizedBuiltinFunc(c *C, vecExprCases vecExprBenchCases) {
for funcName, testCases := range vecExprCases {
for _, testCase := range testCases {
ctx := mock.NewContext()
if funcName == ast.AesDecrypt || funcName == ast.AesEncrypt {
err := ctx.GetSessionVars().SetSystemVar(variable.BlockEncryptionMode, "aes-128-ecb")
c.Assert(err, IsNil)
}
err := ctx.GetSessionVars().SetSystemVar(variable.BlockEncryptionMode, testCase.aesModes)
c.Assert(err, IsNil)
if funcName == ast.CurrentUser || funcName == ast.User {
ctx.GetSessionVars().User = &auth.UserIdentity{
Username: "tidb",
@ -1205,11 +1205,9 @@ func benchmarkVectorizedBuiltinFunc(b *testing.B, vecExprCases vecExprBenchCases
}
for funcName, testCases := range vecExprCases {
for _, testCase := range testCases {
if funcName == ast.AesDecrypt || funcName == ast.AesEncrypt {
err := ctx.GetSessionVars().SetSystemVar(variable.BlockEncryptionMode, "aes-128-ecb")
if err != nil {
panic(err)
}
err := ctx.GetSessionVars().SetSystemVar(variable.BlockEncryptionMode, testCase.aesModes)
if err != nil {
panic(err)
}
if funcName == ast.CurrentUser || funcName == ast.User {
ctx.GetSessionVars().User = &auth.UserIdentity{

91
expression/builtin_encryption_vec.go
View File

@ -15,6 +15,7 @@ package expression
import (
"bytes"
"crypto/aes"
"crypto/md5"
"crypto/rand"
"crypto/sha1"
@ -101,11 +102,97 @@ func (b *builtinAesDecryptSig) vecEvalString(input *chunk.Chunk, result *chunk.C
}
func (b *builtinAesEncryptIVSig) vectorized() bool {
return false
return true
}
// evalString evals AES_ENCRYPT(str, key_str, iv).
// See https://dev.mysql.com/doc/refman/5.7/en/encryption-functions.html#function_aes-decrypt
func (b *builtinAesEncryptIVSig) vecEvalString(input *chunk.Chunk, result *chunk.Column) error {
return errors.Errorf("not implemented")
n := input.NumRows()
strBuf, err := b.bufAllocator.get(types.ETString, n)
if err != nil {
return err
}
defer b.bufAllocator.put(strBuf)
if err := b.args[0].VecEvalString(b.ctx, input, strBuf); err != nil {
return err
}
keyBuf, err := b.bufAllocator.get(types.ETString, n)
if err != nil {
return err
}
defer b.bufAllocator.put(keyBuf)
if err := b.args[1].VecEvalString(b.ctx, input, keyBuf); err != nil {
return err
}
ivBuf, err := b.bufAllocator.get(types.ETString, n)
if err != nil {
return err
}
defer b.bufAllocator.put(ivBuf)
if err := b.args[2].VecEvalString(b.ctx, input, ivBuf); err != nil {
return err
}
isCBC := false
isOFB := false
isCFB := false
switch b.modeName {
case "cbc":
isCBC = true
case "ofb":
isOFB = true
case "cfb":
isCFB = true
default:
return errors.Errorf("unsupported block encryption mode - %v", b.modeName)
}
isConst := b.args[1].ConstItem()
var key []byte
if isConst {
key = encrypt.DeriveKeyMySQL(keyBuf.GetBytes(0), b.keySize)
}
result.ReserveString(n)
for i := 0; i < n; i++ {
// According to doc: If either function argument is NULL, the function returns NULL.
if strBuf.IsNull(i) || keyBuf.IsNull(i) || ivBuf.IsNull(i) {
result.AppendNull()
continue
}
iv := ivBuf.GetBytes(i)
if len(iv) < aes.BlockSize {
return errIncorrectArgs.GenWithStack("The initialization vector supplied to aes_decrypt is too short. Must be at least %d bytes long", aes.BlockSize)
}
// init_vector must be 16 bytes or longer (bytes in excess of 16 are ignored)
iv = iv[0:aes.BlockSize]
if !isConst {
key = encrypt.DeriveKeyMySQL(keyBuf.GetBytes(i), b.keySize)
}
var cipherText []byte
// ANNOTATION:
// we can't use GetBytes here because GetBytes return raw memory in strBuf,
// and the memory will be modified in AESEncryptWithCBC & AESEncryptWithOFB & AESEncryptWithCFB
if isCBC {
cipherText, err = encrypt.AESEncryptWithCBC([]byte(strBuf.GetString(i)), key, iv)
}
if isOFB {
cipherText, err = encrypt.AESEncryptWithOFB([]byte(strBuf.GetString(i)), key, iv)
}
if isCFB {
cipherText, err = encrypt.AESEncryptWithCFB([]byte(strBuf.GetString(i)), key, iv)
}
if err != nil {
result.AppendNull()
}
result.AppendBytes(cipherText)
}
return nil
}
func (b *builtinDecodeSig) vectorized() bool {

11
expression/builtin_encryption_vec_test.go
View File

@ -23,15 +23,18 @@ import (
var vecBuiltinEncryptionCases = map[string][]vecExprBenchCase{
ast.AesEncrypt: {
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString}},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETString}, geners: []dataGenerator{nil, nil, &constStrGener{"iv"}}},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString}, aesModes: "aes-128-ecb"},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETString}, geners: []dataGenerator{nil, nil, &constStrGener{"iv"}}, aesModes: "aes-128-ecb"},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETString}, geners: []dataGenerator{nil, nil, &randLenStrGener{16, 17}}, aesModes: "aes-128-cbc"},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETString}, geners: []dataGenerator{nil, nil, &randLenStrGener{16, 17}}, aesModes: "aes-128-ofb"},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETString}, geners: []dataGenerator{nil, nil, &randLenStrGener{16, 17}}, aesModes: "aes-128-cfb"},
},
ast.Uncompress: {
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}},
},
ast.AesDecrypt: {
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString}},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETString}, geners: []dataGenerator{nil, nil, &constStrGener{"iv"}}},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString}, aesModes: "aes-128-ecb"},
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETString}, geners: []dataGenerator{nil, nil, &constStrGener{"iv"}}, aesModes: "aes-128-ecb"},
},
ast.Compress: {
{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}},