diff --git a/cpp/src/gandiva/gdv_function_stubs.cc b/cpp/src/gandiva/gdv_function_stubs.cc index 94c5454cf24..8be35003e67 100644 --- a/cpp/src/gandiva/gdv_function_stubs.cc +++ b/cpp/src/gandiva/gdv_function_stubs.cc @@ -648,8 +648,14 @@ const char* mask_utf8_utf8_utf8_utf8(int64_t context, const char* data, int32_t return nullptr; } - int32_t max_length = - std::max(upper_length, std::max(lower_length, num_length)) * data_len; + int32_t max_repl_length = std::max(upper_length, std::max(lower_length, num_length)); + int32_t max_length = 0; + if (ARROW_PREDICT_FALSE(arrow::internal::MultiplyWithOverflow(max_repl_length, data_len, + &max_length))) { + gdv_fn_context_set_error_msg(context, "Would overflow maximum output size"); + *out_len = 0; + return nullptr; + } char* out = reinterpret_cast(gdv_fn_context_arena_malloc(context, max_length)); if (out == nullptr) { gdv_fn_context_set_error_msg(context, "Could not allocate memory for output string"); diff --git a/cpp/src/gandiva/gdv_function_stubs_test.cc b/cpp/src/gandiva/gdv_function_stubs_test.cc index 3067a5f2753..3cfeab6985c 100644 --- a/cpp/src/gandiva/gdv_function_stubs_test.cc +++ b/cpp/src/gandiva/gdv_function_stubs_test.cc @@ -1479,6 +1479,20 @@ TEST(TestGdvFnStubs, TestMask) { result = mask_utf8_utf8_utf8_utf8(ctx_ptr, data.c_str(), data_len, "\?", 1, "*", 1, "#", 1, &out_len); EXPECT_EQ(std::string(result, out_len), expected); + + // A replacement length that overflows int32 when multiplied by the input + // length must be rejected instead of under-allocating the output buffer. + std::string big_data(65536, 'A'); + std::string big_repl(65537, 'Z'); + data_len = static_cast(big_data.length()); + out_len = -1; + result = mask_utf8_utf8_utf8_utf8(ctx_ptr, big_data.c_str(), data_len, big_repl.c_str(), + static_cast(big_repl.length()), "x", 1, "n", + 1, &out_len); + EXPECT_EQ(result, nullptr); + EXPECT_EQ(out_len, 0); + EXPECT_TRUE(ctx.has_error()); + ctx.Reset(); } TEST(TestGdvFnStubs, TestAesEncryptDecrypt16) {