diff --git a/cpp/src/gandiva/gdv_function_stubs.cc b/cpp/src/gandiva/gdv_function_stubs.cc index 6b3e9935b017..c9d9e5d5ca58 100644 --- a/cpp/src/gandiva/gdv_function_stubs.cc +++ b/cpp/src/gandiva/gdv_function_stubs.cc @@ -28,6 +28,7 @@ #include "arrow/util/base64.h" #include "arrow/util/bit_util.h" #include "arrow/util/double_conversion_internal.h" +#include "arrow/util/int_util_overflow.h" #include "arrow/util/value_parsing.h" #include "gandiva/encrypt_utils.h" @@ -663,8 +664,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_length; + if (ARROW_PREDICT_FALSE(arrow::internal::MultiplyWithOverflow( + std::max(upper_length, std::max(lower_length, num_length)), data_len, + &max_length))) { + gdv_fn_context_set_error_msg(context, "Could not allocate memory for output string"); + *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 c0197eb815c8..a85370fdc7fb 100644 --- a/cpp/src/gandiva/gdv_function_stubs_test.cc +++ b/cpp/src/gandiva/gdv_function_stubs_test.cc @@ -805,7 +805,7 @@ TEST(TestGdvFnStubs, TestInitCap) { EXPECT_EQ(std::string(out_str, out_len), "{Õhp,Pqśv}Ń+"); EXPECT_FALSE(ctx.has_error()); - out_str = gdv_fn_initcap_utf8(ctx_ptr, "sɦasasdsɦsd\"sdsdɦ", 19, &out_len); + out_str = gdv_fn_initcap_utf8(ctx_ptr, "sɦasasdsɦsd\"sdsdɦ", 20, &out_len); EXPECT_EQ(std::string(out_str, out_len), "Sɦasasdsɦsd\"Sdsdɦ"); EXPECT_FALSE(ctx.has_error()); @@ -1188,6 +1188,44 @@ TEST(TestGdvFnStubs, TestMaskTruncatedUtf8NoOverread) { EXPECT_TRUE(ctx.has_error()); } +TEST(TestGdvFnStubs, TestMaskOutputLengthOverflow) { + gandiva::ExecutionContext ctx; + auto ctx_ptr = reinterpret_cast(&ctx); + int32_t out_len = -1; + const std::string data(65536, 'A'); + const std::string replacement(65536, 'X'); + auto result = mask_utf8_utf8_utf8_utf8( + ctx_ptr, data.data(), static_cast(data.size()), replacement.data(), + static_cast(replacement.size()), "x", 1, "n", 1, &out_len); + EXPECT_EQ(result, nullptr); + EXPECT_EQ(out_len, 0); + EXPECT_TRUE(ctx.has_error()); +} + +TEST(TestGdvFnStubs, TestCaseConversionTruncatedUtf8) { + const std::string inputs[] = {"a\xc2\xa2", "a\xe2\x82\xac", "a\xf0\x9f\x98\x80"}; + for (auto convert : {gdv_fn_lower_utf8, gdv_fn_upper_utf8, gdv_fn_initcap_utf8}) { + for (const auto& input : inputs) { + const auto full_len = static_cast(input.size()); + for (int32_t len = 2; len < full_len; ++len) { + gandiva::ExecutionContext ctx; + auto ctx_ptr = reinterpret_cast(&ctx); + int32_t out_len = -1; + // Bytes beyond len complete the glyph, but are outside the input slice. + convert(ctx_ptr, input.data(), len, &out_len); + EXPECT_EQ(out_len, 0); + EXPECT_TRUE(ctx.has_error()); + } + gandiva::ExecutionContext ctx; + auto ctx_ptr = reinterpret_cast(&ctx); + int32_t out_len = -1; + convert(ctx_ptr, input.data(), full_len, &out_len); + EXPECT_GT(out_len, 0); + EXPECT_FALSE(ctx.has_error()); + } + } +} + TEST(TestGdvFnStubs, TestTranslate) { gandiva::ExecutionContext ctx; int64_t ctx_ptr = reinterpret_cast(&ctx); diff --git a/cpp/src/gandiva/gdv_string_function_stubs.cc b/cpp/src/gandiva/gdv_string_function_stubs.cc index 55f5f13ac312..fc18bc0fac01 100644 --- a/cpp/src/gandiva/gdv_string_function_stubs.cc +++ b/cpp/src/gandiva/gdv_string_function_stubs.cc @@ -289,6 +289,14 @@ const char* gdv_fn_lower_utf8(int64_t context, const char* data, int32_t data_le } // Control reaches here when we encounter a multibyte character + // Ensure the multibyte sequence fits within the buffer to avoid + // reading past data_len (truncated trailing multibyte sequence). + if (char_len > data_len - i) { + gdv_fn_set_error_for_invalid_utf8(context, data[i]); + *out_len = 0; + return ""; + } + const auto* in_char = (const uint8_t*)(data + i); // Decode the multibyte character @@ -366,6 +374,14 @@ const char* gdv_fn_upper_utf8(int64_t context, const char* data, int32_t data_le } // Control reaches here when we encounter a multibyte character + // Ensure the multibyte sequence fits within the buffer to avoid + // reading past data_len (truncated trailing multibyte sequence). + if (char_len > data_len - i) { + gdv_fn_set_error_for_invalid_utf8(context, data[i]); + *out_len = 0; + return ""; + } + const auto* in_char = (const uint8_t*)(data + i); // Decode the multibyte character @@ -584,6 +600,14 @@ const char* gdv_fn_initcap_utf8(int64_t context, const char* data, int32_t data_ char_len = gdv_fn_utf8_char_length(data[i]); // Control reaches here when we encounter a multibyte character + // Ensure the multibyte sequence fits within the buffer to avoid + // reading past data_len (truncated trailing multibyte sequence). + if (char_len > data_len - i) { + gdv_fn_set_error_for_invalid_utf8(context, data[i]); + *out_len = 0; + return ""; + } + const auto* in_char = (const uint8_t*)(data + i); // Decode the multibyte character