use the aes block size to size the aes_encrypt output buffer

Author: Arawoof06

cpp/src/gandiva/gdv_function_stubs.cc

@@ -44,6 +44,9 @@ extern "C" {
 
 ARROW_SUPPRESS_MISSING_DECLARATIONS_WARNING
 
+// AES operates on 16 byte blocks for every supported key length.
+static constexpr int64_t kAesBlockSize = 16;
+
 static char mask_array[256] = {
     (char)0,  (char)1,  (char)2,  (char)3,   (char)4,   (char)5,   (char)6,   (char)7,
     (char)8,  (char)9,  (char)10, (char)11,  (char)12,  (char)13,  (char)14,  (char)15,
@@ -361,19 +364,16 @@ const char* gdv_fn_aes_encrypt(int64_t context, const char* data, int32_t data_l
     return "";
   }
 
-  int64_t kAesBlockSize = 0;
-  if (key_data_len == 16 || key_data_len == 24 || key_data_len == 32) {
-    kAesBlockSize = static_cast<int64_t>(key_data_len);
-  } else {
+  if (key_data_len != 16 && key_data_len != 24 && key_data_len != 32) {
     std::ostringstream oss;
     oss << "invalid key length: " << key_data_len;
     gdv_fn_context_set_error_msg(context, oss.str().c_str());
     *out_len = 0;
     return nullptr;
   }
 
-  *out_len =
-      static_cast<int32_t>(arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize));
+  *out_len = static_cast<int32_t>(
+      arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize);
   char* ret = reinterpret_cast<char*>(gdv_fn_context_arena_malloc(context, *out_len));
   if (ret == nullptr) {
     std::string err_msg = "AES_ENCRYPT: could not allocate memory for ciphertext output";
@@ -407,19 +407,16 @@ const char* gdv_fn_aes_decrypt(int64_t context, const char* data, int32_t data_l
     return "";
   }
 
-  int64_t kAesBlockSize = 0;
-  if (key_data_len == 16 || key_data_len == 24 || key_data_len == 32) {
-    kAesBlockSize = static_cast<int64_t>(key_data_len);
-  } else {
+  if (key_data_len != 16 && key_data_len != 24 && key_data_len != 32) {
     std::ostringstream oss;
     oss << "invalid key length: " << key_data_len;
     gdv_fn_context_set_error_msg(context, oss.str().c_str());
     *out_len = 0;
     return nullptr;
   }
 
-  *out_len =
-      static_cast<int32_t>(arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize));
+  *out_len = static_cast<int32_t>(
+      arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize);
   char* ret = reinterpret_cast<char*>(gdv_fn_context_arena_malloc(context, *out_len));
   if (ret == nullptr) {
     std::string err_msg = "Could not allocate memory for returning aes decrypt plaintext";

cpp/src/gandiva/gdv_function_stubs_test.cc

@@ -1584,4 +1584,37 @@ TEST(TestGdvFnStubs, TestAesEncryptDecryptValidation) {
   EXPECT_THAT(ctx.get_error(), ::testing::HasSubstr("invalid key length"));
   ctx.Reset();
 }
+
+TEST(TestGdvFnStubs, TestAesEncryptBlockAlignedInput) {
+  gandiva::ExecutionContext ctx;
+  std::string key16 = "12345678abcdefgh";
+  auto key16_len = static_cast<int32_t>(key16.length());
+  int64_t ctx_ptr = reinterpret_cast<int64_t>(&ctx);
+
+  // A block-aligned input still gets a whole block of padding appended, so the
+  // ciphertext is one block longer than the input. Encrypting twice from the same
+  // context makes an undersized output buffer observable, because the second
+  // ciphertext lands on top of the first one.
+  std::string first = "aaaaaaaaaaaaaaaa";
+  std::string second = "bbbbbbbbbbbbbbbb";
+  auto first_len = static_cast<int32_t>(first.length());
+  auto second_len = static_cast<int32_t>(second.length());
+
+  int32_t first_cipher_len = 0;
+  const char* first_cipher = gdv_fn_aes_encrypt(
+      ctx_ptr, first.c_str(), first_len, key16.c_str(), key16_len, &first_cipher_len);
+  ASSERT_FALSE(ctx.has_error());
+  EXPECT_EQ(first_cipher_len, first_len + 16);
+
+  int32_t second_cipher_len = 0;
+  gdv_fn_aes_encrypt(ctx_ptr, second.c_str(), second_len, key16.c_str(), key16_len,
+                     &second_cipher_len);
+  ASSERT_FALSE(ctx.has_error());
+
+  int32_t decrypted_len = 0;
+  const char* decrypted = gdv_fn_aes_decrypt(ctx_ptr, first_cipher, first_cipher_len,
+                                             key16.c_str(), key16_len, &decrypted_len);
+  ASSERT_FALSE(ctx.has_error());
+  EXPECT_EQ(first, std::string(decrypted, decrypted_len));
+}
 }  // namespace gandiva