From b5705a96155c252f23a06b9a9e832887c463edff Mon Sep 17 00:00:00 2001
From: Nayana Naik <nayana@labs.digiscrypt.com>
Date: Thu, 15 Oct 2026 15:14:03 +0530
Subject: [PATCH] MINOR: [C++][Gandiva] Fix output buffer size in mask()

---
 cpp/src/gandiva/gdv_function_stubs.cc      | 14 ++++++++++++--
 cpp/src/gandiva/gdv_function_stubs_test.cc | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/cpp/src/gandiva/gdv_function_stubs.cc b/cpp/src/gandiva/gdv_function_stubs.cc
index 6b3e9935b017..cf6319ae1dad 100644
--- a/cpp/src/gandiva/gdv_function_stubs.cc
+++ b/cpp/src/gandiva/gdv_function_stubs.cc
@@ -19,6 +19,7 @@
 
 #include <utf8proc.h>
 
+#include <algorithm>
 #include <boost/crc.hpp>
 #include <cstdio>
 #include <sstream>
@@ -28,6 +29,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 +665,16 @@ 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;
+  // Characters that are not replaced are copied as-is, so reserve at least one
+  // output byte per input byte even when all replacements are empty.
+  int32_t max_length;
+  if (ARROW_PREDICT_FALSE(arrow::internal::MultiplyWithOverflow(
+          std::max({upper_length, lower_length, num_length, 1}), 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<char*>(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..8549aa4c077d 100644
--- a/cpp/src/gandiva/gdv_function_stubs_test.cc
+++ b/cpp/src/gandiva/gdv_function_stubs_test.cc
@@ -1539,6 +1539,26 @@ 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);
+
+  // With empty replacements the unmasked characters must still get their own
+  // output space, so a later allocation cannot overwrite them.
+  data = "QwErTy:4)";
+  expected = ":)";
+  data_len = static_cast<int32_t>(data.length());
+  result = mask_utf8_utf8_utf8_utf8(ctx_ptr, data.c_str(), data_len, "", 0, "", 0, "", 0,
+                                    &out_len);
+  int32_t other_len = 0;
+  mask_utf8_utf8_utf8_utf8(ctx_ptr, "#$", 2, "", 0, "", 0, "", 0, &other_len);
+  EXPECT_EQ(std::string(result, out_len), expected);
+
+  data = std::string(46341, 'A');
+  data_len = static_cast<int32_t>(data.length());
+  result = mask_utf8_utf8_utf8_utf8(ctx_ptr, data.c_str(), data_len, data.c_str(),
+                                    data_len, "x", 1, "n", 1, &out_len);
+  EXPECT_EQ(out_len, 0);
+  EXPECT_THAT(ctx.get_error(),
+              ::testing::HasSubstr("Would overflow maximum output size"));
+  ctx.Reset();
 }
 
 TEST(TestGdvFnStubs, TestAesEncryptDecrypt16) {