Rewrite aes encryption #257

be/CMakeLists.txt

@@ -458,7 +458,6 @@ set(DORIS_LINK_LIBS
     DorisGen
     Webserver
     TestUtil
-    AES
     ${WL_END_GROUP}
 )
 
@@ -557,7 +556,6 @@ add_subdirectory(${SRC_DIR}/exprs)
 add_subdirectory(${SRC_DIR}/udf)
 add_subdirectory(${SRC_DIR}/runtime)
 add_subdirectory(${SRC_DIR}/testutil)
-add_subdirectory(${SRC_DIR}/aes)
 add_subdirectory(${SRC_DIR}/tools)
 
 # Utility CMake function to make specifying tests and benchmarks less verbose

be/src/exprs/encryption_functions.cpp

@@ -18,7 +18,7 @@
 #include "exprs/encryption_functions.h"
 
 #include <openssl/md5.h>
-#include "aes/my_aes.h"
+#include "util/aes_util.h"
 #include "exprs/anyval_util.h"
 #include "exprs/expr.h"
 #include "util/debug_util.h"
@@ -42,8 +42,8 @@ StringVal EncryptionFunctions::aes_encrypt(FunctionContext* ctx,
     boost::scoped_array<char> p;
     p.reset(new char[cipher_len]);
 
-    int ret_code = my_aes_encrypt((unsigned char *)src.ptr, src.len,
-            (unsigned char*)p.get(), (unsigned char *)key.ptr, key.len, my_aes_128_ecb, NULL);
+    int ret_code = AesUtil::encrypt(AES_128_ECB, (unsigned char*)src.ptr, src.len,
+            (unsigned char*)key.ptr, key.len, NULL, true, (unsigned char*)p.get());
     if (ret_code < 0) {
         return StringVal::null();
     }
@@ -60,8 +60,8 @@ StringVal EncryptionFunctions::aes_decrypt(FunctionContext* ctx,
     boost::scoped_array<char> p;
     p.reset(new char[cipher_len]);
 
-    int ret_code = my_aes_decrypt((unsigned char *)src.ptr, src.len, (unsigned char*)p.get(),
-            (unsigned char *)key.ptr, key.len, my_aes_128_ecb, NULL);
+    int ret_code = AesUtil::decrypt(AES_128_ECB, (unsigned char*)src.ptr, src.len,
+            (unsigned char*)key.ptr, key.len, NULL, true,  (unsigned char*)p.get());
     if (ret_code < 0) {
         return StringVal::null();
     }

be/src/util/CMakeLists.txt

@@ -69,6 +69,7 @@ add_library(Util STATIC
   cidr.cpp
   core_local.cpp
   uid_util.cpp
+  aes_util.cpp
 )
 
 #ADD_BE_TEST(integer-array-test)

be/src/util/aes_util.cpp

@@ -0,0 +1,180 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "aes_util.h"
+#include  <cstring>
+#include  <string>
+#include  <memory>
+#include  <iostream>
+
+#include <openssl/aes.h>
+#include <openssl/evp.h>
+#include <openssl/err.h>
+
+#include "exprs/base64.h"
+
+namespace palo {
+
+static const int AES_MAX_KEY_LENGTH = 256;
+
+const EVP_CIPHER* get_evp_type(const AesMode mode) {
+    switch (mode) {
+    case AES_128_ECB:
+        return EVP_aes_128_ecb();
+    case AES_128_CBC:
+        return EVP_aes_128_cbc();
+    case AES_192_ECB:
+        return EVP_aes_192_ecb();
+    case AES_192_CBC:
+        return EVP_aes_192_cbc();
+    case AES_256_ECB:
+        return EVP_aes_256_ecb();
+    case AES_256_CBC:
+        return EVP_aes_256_cbc();
+    default:
+        return NULL;
+    }
+}
+
+static uint aes_mode_key_sizes[]= {
+    128 /* AES_128_ECB */,
+    192 /* AES_192_ECB */,
+    256 /* AES_256_ECB */,
+    128 /* AES_128_CBC */,
+    192 /* AES_192_CBC */,
+    256 /* AES_256_CBC */
+};
+
+static void aes_create_key(const unsigned char* origin_key, uint32_t key_length,
+        uint8_t* encrypt_key, AesMode mode) {
+    const uint key_size= aes_mode_key_sizes[mode] / 8;
+    uint8_t *origin_key_end= ((uint8_t*)origin_key) + key_length; /* origin key boundary*/
+
+    uint8_t *encrypt_key_end; /* encrypt key boundary */
+    encrypt_key_end= encrypt_key + key_size;
+
+    std::memset(encrypt_key, 0, key_size);          /* initialize key  */
+
+    uint8_t *ptr; /* Start of the encrypt key*/
+    uint8_t *origin_ptr; /* Start of the origin key */
+    for (ptr = encrypt_key, origin_ptr = (uint8_t*)origin_key;
+            origin_ptr < origin_key_end; ptr++, origin_ptr++) {
+        if (ptr == encrypt_key_end) {
+            /* loop over origin key until we used all key */
+            ptr = encrypt_key;
+        }
+        *ptr ^= *origin_ptr;
+    }
+}
+
+static int do_encrypt(EVP_CIPHER_CTX* aes_ctx, const EVP_CIPHER* cipher,
+        const unsigned char* source, uint32_t source_length, const unsigned char* encrypt_key,
+        const unsigned char* iv, bool padding, unsigned char* encrypt, int* length_ptr) {
+    int ret = EVP_EncryptInit(aes_ctx, cipher, encrypt_key, iv);
+    if (ret == 0) {
+        return ret;
+    }
+    ret = EVP_CIPHER_CTX_set_padding(aes_ctx, padding);
+    if (ret == 0) {
+        return ret;
+    }
+    int u_len = 0;
+    ret = EVP_EncryptUpdate(aes_ctx, encrypt, &u_len, source, source_length);
+    if (ret == 0) {
+        return ret;
+    }
+    int f_len = 0;
+    ret = EVP_EncryptFinal(aes_ctx, encrypt + u_len, &f_len);
+    *length_ptr = u_len + f_len;
+    return ret;
+}
+
+int AesUtil::encrypt(AesMode mode, const unsigned char* source, uint32_t source_length,
+        const unsigned char* key, uint32_t key_length, const unsigned char* iv,
+        bool padding, unsigned char* encrypt) {
+    EVP_CIPHER_CTX aes_ctx;
+    const EVP_CIPHER* cipher = get_evp_type(mode);
+    /* The encrypt key to be used for encryption */
+    unsigned char encrypt_key[AES_MAX_KEY_LENGTH / 8];
+    aes_create_key(key, key_length, encrypt_key, mode);
+
+    if (cipher == nullptr || (EVP_CIPHER_iv_length(cipher) > 0 && !iv)) {
+        return AES_BAD_DATA;
+    }
+    EVP_CIPHER_CTX_init(&aes_ctx);
+    int length = 0;
+    int ret = do_encrypt(&aes_ctx, cipher, source,
+            source_length, encrypt_key, iv, padding, encrypt, &length);
+    EVP_CIPHER_CTX_cleanup(&aes_ctx);
+    if (ret == 0) {
+        ERR_clear_error();
+        return AES_BAD_DATA;
+    } else {
+        return length;
+    }
+}
+
+static int do_decrypt(EVP_CIPHER_CTX* aes_ctx, const EVP_CIPHER* cipher,
+        const unsigned char* encrypt, uint32_t encrypt_length, const unsigned char* encrypt_key,
+        const unsigned char* iv, bool padding, unsigned char* decrypt_content, int* length_ptr) {
+    int ret = EVP_DecryptInit(aes_ctx, cipher, encrypt_key, iv);
+    if (ret == 0) {
+        return ret;
+    }
+    ret = EVP_CIPHER_CTX_set_padding(aes_ctx, padding);
+    if (ret == 0) {
+        return ret;
+    }
+    int u_len = 0;
+    ret = EVP_DecryptUpdate(aes_ctx, decrypt_content, &u_len, encrypt, encrypt_length);
+    if (ret == 0) {
+        return ret;
+    }
+    int f_len = 0;
+    ret = EVP_DecryptFinal_ex(aes_ctx, decrypt_content + u_len, &f_len);
+    *length_ptr = u_len + f_len;
+    return ret;
+}
+
+int AesUtil::decrypt(AesMode mode, const unsigned char* encrypt, uint32_t encrypt_length,
+        const unsigned char* key, uint32_t key_length, const unsigned char* iv,
+        bool padding, unsigned char* decrypt_content) {
+    EVP_CIPHER_CTX aes_ctx;
+    const EVP_CIPHER* cipher = get_evp_type(mode);
+
+    /* The encrypt key to be used for decryption */
+    unsigned char encrypt_key[AES_MAX_KEY_LENGTH / 8];
+    aes_create_key(key, key_length, encrypt_key, mode);
+
+    if (cipher == nullptr || (EVP_CIPHER_iv_length(cipher) > 0 && !iv)) {
+        return AES_BAD_DATA;
+    }
+
+    EVP_CIPHER_CTX_init(&aes_ctx);
+    int length = 0;
+    int ret = do_decrypt(&aes_ctx, cipher, encrypt, encrypt_length,
+            encrypt_key, iv, padding, decrypt_content, &length);
+    EVP_CIPHER_CTX_cleanup(&aes_ctx);
+    if (ret > 0) {
+        return length;
+    } else {
+        ERR_clear_error();
+        return AES_BAD_DATA;
+    }
+}
+
+}

be/src/util/aes_util.h

@@ -0,0 +1,47 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <stdint.h>
+
+namespace palo {
+
+enum AesMode {
+    AES_128_ECB,
+    AES_192_ECB,
+    AES_256_ECB,
+    AES_128_CBC,
+    AES_192_CBC,
+    AES_256_CBC
+};
+
+enum AesState {
+    AES_SUCCESS = 0,
+    AES_BAD_DATA = -1
+};
+
+class AesUtil {
+public:
+    static int encrypt(AesMode mode, const unsigned char* source, uint32_t source_length,
+            const unsigned char* key, uint32_t key_length, const unsigned char* iv,
+            bool padding, unsigned char* encrypt);
+
+    static int decrypt(AesMode mode, const unsigned char* encrypt, uint32_t encrypt_length,
+            const unsigned char* key, uint32_t key_length, const unsigned char* iv,
+            bool padding, unsigned char* decrypt_content);
+};
+
+}

be/test/util/CMakeLists.txt

@@ -36,3 +36,4 @@ ADD_BE_TEST(json_util_test)
 ADD_BE_TEST(byte_buffer_test2)
 ADD_BE_TEST(uid_util_test)
 ADD_BE_TEST(arena_test)
+ADD_BE_TEST(aes_util_test)