GH-49614: [C++] Fix silent truncation in base64_decode on invalid input

Author: Reranko05

cpp/src/arrow/util/base64.h

@@ -21,6 +21,8 @@
 #include <string_view>
 
 #include "arrow/util/visibility.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
 
 namespace arrow {
 namespace util {
@@ -29,7 +31,7 @@ ARROW_EXPORT
 std::string base64_encode(std::string_view s);
 
 ARROW_EXPORT
-std::string base64_decode(std::string_view s);
+arrow::Result<std::string> base64_decode(std::string_view s);
 
 }  // namespace util
 }  // namespace arrow

cpp/src/arrow/util/string_test.cc

@@ -28,6 +28,7 @@
 #include "arrow/testing/gtest_util.h"
 #include "arrow/util/regex.h"
 #include "arrow/util/string.h"
+#include "arrow/util/base64.h"
 
 namespace arrow {
 namespace internal {
@@ -232,12 +233,79 @@ TEST(ToChars, FloatingPoint) {
     // to std::to_string which may make ad hoc formatting choices, so we cannot
     // really test much about the result.
     auto result = ToChars(0.0f);
-    ASSERT_TRUE(result.starts_with("0")) << result;
+    ASSERT_TRUE(result.rfind("0", 0) == 0) << result;
     result = ToChars(0.25);
-    ASSERT_TRUE(result.starts_with("0.25")) << result;
+    ASSERT_TRUE(result.rfind("0.25", 0) == 0) << result;
   }
 }
 
+TEST(Base64DecodeTest, ValidInputs) {
+  auto r1 = arrow::util::base64_decode("Zg==");
+  ASSERT_TRUE(r1.ok());
+  EXPECT_EQ(r1.ValueOrDie(), "f");
+  auto r2 = arrow::util::base64_decode("Zm8=");
+  ASSERT_TRUE(r2.ok());
+  EXPECT_EQ(r2.ValueOrDie(), "fo");
+  auto r3 = arrow::util::base64_decode("Zm9v");
+  ASSERT_TRUE(r3.ok());
+  EXPECT_EQ(r3.ValueOrDie(), "foo");
+  auto r4 = arrow::util::base64_decode("aGVsbG8gd29ybGQ=");
+  ASSERT_TRUE(r4.ok());
+  EXPECT_EQ(r4.ValueOrDie(), "hello world");
+}
+
+TEST(Base64DecodeTest, InvalidLength) {
+  auto r1 = arrow::util::base64_decode("abc");
+  ASSERT_FALSE(r1.ok());
+  auto r2 = arrow::util::base64_decode("abcde");
+  ASSERT_FALSE(r2.ok());
+}
+
+TEST(Base64DecodeTest, InvalidCharacters) {
+  auto r1 = arrow::util::base64_decode("ab$=");
+  ASSERT_FALSE(r1.ok());
+  auto r2 = arrow::util::base64_decode("Zm9v*");
+  ASSERT_FALSE(r2.ok());
+  auto r3 = arrow::util::base64_decode("abcd$AAA");
+  ASSERT_FALSE(r3.ok());
+}
+
+TEST(Base64DecodeTest, InvalidPadding) {
+  auto r1 = arrow::util::base64_decode("ab=c");
+  ASSERT_FALSE(r1.ok());
+  auto r2 = arrow::util::base64_decode("abc===");
+  ASSERT_FALSE(r2.ok());
+  auto r3 = arrow::util::base64_decode("abcd=AAA");
+  ASSERT_FALSE(r3.ok());
+  auto r4 = arrow::util::base64_decode("Zm=9v");
+  ASSERT_FALSE(r4.ok());
+}
+
+TEST(Base64DecodeTest, EdgeCases) {
+  auto r1 = arrow::util::base64_decode("====");
+  ASSERT_FALSE(r1.ok());
+  auto r2 = arrow::util::base64_decode("TQ==");
+  ASSERT_TRUE(r2.ok());
+  EXPECT_EQ(r2.ValueOrDie(), "M");
+}
+
+TEST(Base64DecodeTest, EmptyInput) {
+  auto r = arrow::util::base64_decode("");
+  ASSERT_TRUE(r.ok());
+  EXPECT_EQ(r.ValueOrDie(), "");
+}
+
+TEST(Base64DecodeTest, NonAsciiInput) {
+  std::string input = std::string("abcd") + char(0xFF) + "==";
+  auto r = arrow::util::base64_decode(input);
+  ASSERT_FALSE(r.ok());
+}
+
+TEST(Base64DecodeTest, PartialCorruption) {
+  auto r = arrow::util::base64_decode("aGVs$G8gd29ybGQ=");
+  ASSERT_FALSE(r.ok());
+}
+
 #if !defined(_WIN32) || defined(NDEBUG)
 
 TEST(ToChars, LocaleIndependent) {

cpp/src/arrow/vendored/base64.cpp

@@ -30,7 +30,11 @@
 */
 
 #include "arrow/util/base64.h"
+#include "arrow/util/logging.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
 #include <iostream>
+#include <cctype>
 
 namespace arrow {
 namespace util {
@@ -93,18 +97,51 @@ std::string base64_encode(std::string_view string_to_encode) {
   return base64_encode(bytes_to_encode, in_len);
 }
 
-std::string base64_decode(std::string_view encoded_string) {
+arrow::Result<std::string> base64_decode(std::string_view encoded_string) {
   size_t in_len = encoded_string.size();
   int i = 0;
   int j = 0;
   int in_ = 0;
   unsigned char char_array_4[4], char_array_3[3];
   std::string ret;
 
-  while (in_len-- && ( encoded_string[in_] != '=') && is_base64(encoded_string[in_])) {
+  static const std::string base64_chars =
+               "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+               "abcdefghijklmnopqrstuvwxyz"
+               "0123456789+/";
+  
+  auto is_base64 = [](unsigned char c) -> bool {
+    return (std::isalnum(c) || (c == '+') || (c == '/'));
+  };
+
+  if (encoded_string.size() % 4 != 0) {
+    return arrow::Status::Invalid("Invalid base64 input: length is not a multiple of 4");
+  }
+
+  size_t padding_start = encoded_string.find('=');
+  if (padding_start != std::string::npos) {
+    for (size_t k = padding_start; k < encoded_string.size(); ++k) {
+      if (encoded_string[k] != '=') {
+        return arrow::Status::Invalid("Invalid base64 input: padding character '=' found at invalid position");
+      }
+    }
+
+    size_t padding_count = encoded_string.size() - padding_start;
+    if (padding_count > 2) {
+      return arrow::Status::Invalid("Invalid base64 input: too many padding characters");
+    }
+  }
+
+    for (char c : encoded_string) {
+      if (c != '=' && !is_base64(c)) {
+        return arrow::Status::Invalid("Invalid base64 input: contains non-base64 character '" + std::string(1, c) + "'");
+      }
+    }
+
+  while (in_len-- && encoded_string[in_] != '=') {
     char_array_4[i++] = encoded_string[in_]; in_++;
-    if (i ==4) {
-      for (i = 0; i <4; i++)
+    if (i == 4) {
+      for (i = 0; i < 4; i++)
         char_array_4[i] = base64_chars.find(char_array_4[i]) & 0xff;
 
       char_array_3[0] = ( char_array_4[0] << 2       ) + ((char_array_4[1] & 0x30) >> 4);