Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

expression: implement vectorized evaluation for builtinAesEncryptIVSig #13521

Merged
merged 4 commits into from
Nov 18, 2019
Merged

expression: implement vectorized evaluation for builtinAesEncryptIVSig #13521

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3cbb340
implement vectorized evaluation for `builtinAesEncryptIVSig`
js00070 Nov 17, 2019
9ebc8a5
add key const check to optimize
js00070 Nov 17, 2019
78c106d
Merge branch 'master' into builtinAesEncryptIVSig
js00070 Nov 18, 2019
9fccea1
Merge branch 'master' into builtinAesEncryptIVSig
sre-bot Nov 18, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 7 additions & 9 deletions expression/bench_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

type vecExprBenchCases map[string][]vecExprBenchCase
Expand Down Expand Up @@ -992,10 +994,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.User {
ctx.GetSessionVars().User = &auth.UserIdentity{
Username: "tidb",
Expand Down Expand Up @@ -1195,11 +1195,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.User {
ctx.GetSessionVars().User = &auth.UserIdentity{
Expand Down
91 changes: 89 additions & 2 deletions expression/builtin_encryption_vec.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ package expression

import (
"bytes"
"crypto/aes"
"crypto/md5"
"crypto/rand"
"crypto/sha1"
Expand Down Expand Up @@ -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 {
Expand Down
11 changes: 7 additions & 4 deletions expression/builtin_encryption_vec_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}},
Expand Down