From b74f571813453d13cbc9e0b6e414320ff223cad1 Mon Sep 17 00:00:00 2001 From: abdul rawoof Date: Thu, 16 Jul 2026 13:15:27 +0530 Subject: [PATCH 1/2] use the aes block size to size the aes_encrypt output buffer --- cpp/src/gandiva/gdv_function_stubs.cc | 21 ++++++-------- cpp/src/gandiva/gdv_function_stubs_test.cc | 33 ++++++++++++++++++++++ 2 files changed, 42 insertions(+), 12 deletions(-) diff --git a/cpp/src/gandiva/gdv_function_stubs.cc b/cpp/src/gandiva/gdv_function_stubs.cc index 6b3e9935b017..9dbf4b56ae4f 100644 --- a/cpp/src/gandiva/gdv_function_stubs.cc +++ b/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, @@ -366,10 +369,7 @@ 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(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()); @@ -377,8 +377,8 @@ const char* gdv_fn_aes_encrypt(int64_t context, const char* data, int32_t data_l return nullptr; } - *out_len = - static_cast(arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize)); + *out_len = static_cast( + arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize); char* ret = reinterpret_cast(gdv_fn_context_arena_malloc(context, *out_len)); if (ret == nullptr) { std::string err_msg = "AES_ENCRYPT: could not allocate memory for ciphertext output"; @@ -412,10 +412,7 @@ 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(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()); @@ -423,8 +420,8 @@ const char* gdv_fn_aes_decrypt(int64_t context, const char* data, int32_t data_l return nullptr; } - *out_len = - static_cast(arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize)); + *out_len = static_cast( + arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize); char* ret = reinterpret_cast(gdv_fn_context_arena_malloc(context, *out_len)); if (ret == nullptr) { std::string err_msg = "Could not allocate memory for returning aes decrypt plaintext"; diff --git a/cpp/src/gandiva/gdv_function_stubs_test.cc b/cpp/src/gandiva/gdv_function_stubs_test.cc index c0197eb815c8..0efa76c3aa84 100644 --- a/cpp/src/gandiva/gdv_function_stubs_test.cc +++ b/cpp/src/gandiva/gdv_function_stubs_test.cc @@ -1621,4 +1621,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(key16.length()); + int64_t ctx_ptr = reinterpret_cast(&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(first.length()); + auto second_len = static_cast(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 From 06d25b71321a4070b5692f26ba7211ad22e4ab78 Mon Sep 17 00:00:00 2001 From: abdul rawoof Date: Sun, 18 Oct 2026 00:49:49 +0530 Subject: [PATCH 2/2] guard the aes output size against int32 overflow --- cpp/src/gandiva/gdv_function_stubs.cc | 23 ++++++++++++++--- cpp/src/gandiva/gdv_function_stubs_test.cc | 29 ++++++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/cpp/src/gandiva/gdv_function_stubs.cc b/cpp/src/gandiva/gdv_function_stubs.cc index 9dbf4b56ae4f..a3803e498958 100644 --- a/cpp/src/gandiva/gdv_function_stubs.cc +++ b/cpp/src/gandiva/gdv_function_stubs.cc @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -377,8 +378,15 @@ const char* gdv_fn_aes_encrypt(int64_t context, const char* data, int32_t data_l return nullptr; } - *out_len = static_cast( - arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize); + int64_t alloc_len = + arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize; + if (ARROW_PREDICT_FALSE(alloc_len > std::numeric_limits::max())) { + gdv_fn_context_set_error_msg(context, + "AES_ENCRYPT: would overflow maximum output size"); + *out_len = 0; + return nullptr; + } + *out_len = static_cast(alloc_len); char* ret = reinterpret_cast(gdv_fn_context_arena_malloc(context, *out_len)); if (ret == nullptr) { std::string err_msg = "AES_ENCRYPT: could not allocate memory for ciphertext output"; @@ -420,8 +428,15 @@ const char* gdv_fn_aes_decrypt(int64_t context, const char* data, int32_t data_l return nullptr; } - *out_len = static_cast( - arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize); + int64_t alloc_len = + arrow::bit_util::RoundUpToPowerOf2(data_len, kAesBlockSize) + kAesBlockSize; + if (ARROW_PREDICT_FALSE(alloc_len > std::numeric_limits::max())) { + gdv_fn_context_set_error_msg(context, + "AES_DECRYPT: would overflow maximum output size"); + *out_len = 0; + return nullptr; + } + *out_len = static_cast(alloc_len); char* ret = reinterpret_cast(gdv_fn_context_arena_malloc(context, *out_len)); if (ret == nullptr) { std::string err_msg = "Could not allocate memory for returning aes decrypt plaintext"; diff --git a/cpp/src/gandiva/gdv_function_stubs_test.cc b/cpp/src/gandiva/gdv_function_stubs_test.cc index 0efa76c3aa84..3349bdfd1e8a 100644 --- a/cpp/src/gandiva/gdv_function_stubs_test.cc +++ b/cpp/src/gandiva/gdv_function_stubs_test.cc @@ -1654,4 +1654,33 @@ TEST(TestGdvFnStubs, TestAesEncryptBlockAlignedInput) { ASSERT_FALSE(ctx.has_error()); EXPECT_EQ(first, std::string(decrypted, decrypted_len)); } + +TEST(TestGdvFnStubs, TestAesOutputSizeOverflow) { + gandiva::ExecutionContext ctx; + std::string key16 = "12345678abcdefgh"; + auto key16_len = static_cast(key16.length()); + int64_t ctx_ptr = reinterpret_cast(&ctx); + std::string data = "a"; + + // Rounding up and adding the padding block pushes these lengths past INT32_MAX. + // The size check runs before the input is read, so the short buffer is fine. + for (int32_t data_len : + {std::numeric_limits::max() - 15, std::numeric_limits::max()}) { + int32_t out_len = -1; + const char* ret = gdv_fn_aes_encrypt(ctx_ptr, data.c_str(), data_len, key16.c_str(), + key16_len, &out_len); + EXPECT_EQ(ret, nullptr); + EXPECT_EQ(out_len, 0); + EXPECT_THAT(ctx.get_error(), ::testing::HasSubstr("would overflow")); + ctx.Reset(); + + out_len = -1; + ret = gdv_fn_aes_decrypt(ctx_ptr, data.c_str(), data_len, key16.c_str(), key16_len, + &out_len); + EXPECT_EQ(ret, nullptr); + EXPECT_EQ(out_len, 0); + EXPECT_THAT(ctx.get_error(), ::testing::HasSubstr("would overflow")); + ctx.Reset(); + } +} } // namespace gandiva