ARROW: Return Result<std::string> in base64_decode and add validation + tests

Author: Reranko05

cpp/src/arrow/flight/flight_test.cc

@@ -620,7 +620,8 @@ void ParseBasicHeader(const CallHeaders& incoming_headers, std::string& username
                       std::string& password) {
   std::string encoded_credentials =
       FindKeyValPrefixInCallHeaders(incoming_headers, kAuthHeader, kBasicPrefix);
-  std::stringstream decoded_stream(arrow::util::base64_decode(encoded_credentials));
+  ASSERT_OK_AND_ASSIGN(auto decoded, arrow::util::base64_decode(encoded_credentials));
+  std::stringstream decoded_stream(decoded);
   std::getline(decoded_stream, username, ':');
   std::getline(decoded_stream, password, ':');
 }

cpp/src/arrow/util/base64.h

@@ -20,8 +20,8 @@
 #include <string>
 #include <string_view>
 
-#include "arrow/util/visibility.h"
 #include "arrow/result.h"
+#include "arrow/util/visibility.h"
 
 namespace arrow {
 namespace util {

cpp/src/arrow/util/base64_test.cc

@@ -15,72 +15,76 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#include "arrow/testing/gtest_util.h"
 #include "arrow/util/base64.h"
-
+#include "arrow/testing/gtest_util.h"
 
 namespace arrow {
 namespace util {
 
 TEST(Base64DecodeTest, ValidInputs) {
-  ASSERT_OK_AND_ASSIGN(auto empty, arrow::util::base64_decode(""));
+  ASSERT_OK_AND_ASSIGN(auto empty, base64_decode(""));
   EXPECT_EQ(empty, "");
 
-  ASSERT_OK_AND_ASSIGN(auto two_paddings, arrow::util::base64_decode("Zg=="));
+  ASSERT_OK_AND_ASSIGN(auto two_paddings, base64_decode("Zg=="));
   EXPECT_EQ(two_paddings, "f");
 
-  ASSERT_OK_AND_ASSIGN(auto one_padding, arrow::util::base64_decode("Zm8="));
+  ASSERT_OK_AND_ASSIGN(auto one_padding, base64_decode("Zm8="));
   EXPECT_EQ(one_padding, "fo");
 
-  ASSERT_OK_AND_ASSIGN(auto no_padding, arrow::util::base64_decode("Zm9v"));
+  ASSERT_OK_AND_ASSIGN(auto no_padding, base64_decode("Zm9v"));
   EXPECT_EQ(no_padding, "foo");
 
-  ASSERT_OK_AND_ASSIGN(auto multiblock, arrow::util::base64_decode("SGVsbG8gd29ybGQ="));
+  ASSERT_OK_AND_ASSIGN(auto multiblock, base64_decode("SGVsbG8gd29ybGQ="));
   EXPECT_EQ(multiblock, "Hello world");
 }
 
 TEST(Base64DecodeTest, BinaryOutput) {
   // 'A' maps to index 0 — same zero value used for padding slots
   // verifies the 'A' bug is not present
-  ASSERT_OK_AND_ASSIGN(auto all_A, arrow::util::base64_decode("AAAA"));
+  ASSERT_OK_AND_ASSIGN(auto all_A, base64_decode("AAAA"));
   EXPECT_EQ(all_A, std::string("\x00\x00\x00", 3));
 
   // Arbitrary non-ASCII output bytes
-  ASSERT_OK_AND_ASSIGN(auto binary, arrow::util::base64_decode("AP8A"));
+  ASSERT_OK_AND_ASSIGN(auto binary, base64_decode("AP8A"));
   EXPECT_EQ(binary, std::string("\x00\xff\x00", 3));
 }
 
 TEST(Base64DecodeTest, InvalidLength) {
-  ASSERT_RAISES_WITH_MESSAGE(
-    Invalid,
-    "Invalid: Invalid base64 input: length is not a multiple of 4",
-    arrow::util::base64_decode("abc"));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode("abc"));
 }
 
 TEST(Base64DecodeTest, InvalidCharacters) {
-  ASSERT_RAISES(Invalid, arrow::util::base64_decode("ab$="));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode("ab$="));
 
   // Non-ASCII byte
   std::string non_ascii = std::string("abc") + static_cast<char>(0xFF);
-  ASSERT_RAISES(Invalid, arrow::util::base64_decode(non_ascii));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode(non_ascii));
 
   // Corruption mid-string across multiple blocks
-  ASSERT_RAISES(Invalid, arrow::util::base64_decode("aGVs$G8gd29ybGQ="));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode("aGVs$G8gd29ybGQ="));
 }
 
 TEST(Base64DecodeTest, InvalidPadding) {
   // Padding in wrong position within block
-  ASSERT_RAISES(Invalid, arrow::util::base64_decode("ab=c"));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode("ab=c"));
 
   // 3 padding characters — exceeds maximum of 2
-  ASSERT_RAISES(Invalid, arrow::util::base64_decode("a==="));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode("a==="));
 
   // 4 padding characters
-  ASSERT_RAISES(Invalid, arrow::util::base64_decode("===="));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode("===="));
 
   // Padding in non-final block across multiple blocks
-  ASSERT_RAISES(Invalid, arrow::util::base64_decode("Zm8=Zm8="));
+  ASSERT_RAISES_WITH_MESSAGE(Invalid, "Invalid: Invalid base64 input",
+                             base64_decode("Zm8=Zm8="));
 }
 
-}
-}
+}  // namespace util
+}  // namespace arrow

cpp/src/arrow/util/string_test.cc

@@ -28,7 +28,6 @@
 #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 {

cpp/src/arrow/vendored/base64.cpp

@@ -32,7 +32,6 @@
 #include "arrow/util/base64.h"
 #include "arrow/result.h"
 #include <iostream>
-#include <cctype>
 
 namespace arrow {
 namespace util {
@@ -93,15 +92,15 @@ std::string base64_encode(std::string_view string_to_encode) {
 Result<std::string> base64_decode(std::string_view encoded_string) {
   size_t in_len = encoded_string.size();
   int i = 0;
-  int in_ = 0;
+  std::string_view::size_type in_ = 0;
   int padding_count = 0;
   int block_padding = 0;
   bool padding_started = false;
   unsigned char char_array_4[4], char_array_3[3];
   std::string ret;
 
   if (encoded_string.size() % 4 != 0) {
-    return Status::Invalid("Invalid base64 input: length is not a multiple of 4");
+    return Status::Invalid("Invalid base64 input");
   }
 
   while (in_len--) {
@@ -112,19 +111,17 @@ Result<std::string> base64_decode(std::string_view encoded_string) {
       padding_count++;
 
       if (padding_count > 2) {
-        return Status::Invalid("Invalid base64 input: too many padding characters");
+        return Status::Invalid("Invalid base64 input");
       }
 
       char_array_4[i++] = 0;
     } else {
       if (padding_started) {
-        return Status::Invalid("Invalid base64 input: padding characters must be at the end");
+        return Status::Invalid("Invalid base64 input");
       }
 
       if (base64_chars.find(c) == std::string::npos) {
-        return Status::Invalid(
-            "Invalid base64 input: contains non-base64 byte at position " +
-            std::to_string(in_));
+        return Status::Invalid("Invalid base64 input");
       }
 
       char_array_4[i++] = c;

cpp/src/gandiva/gdv_function_stubs.cc

@@ -269,7 +269,15 @@ const char* gdv_fn_base64_decode_utf8(int64_t context, const char* in, int32_t i
     return "";
   }
   // use arrow method to decode base64 string
-  std::string decoded_str = arrow::util::base64_decode(std::string_view(in, in_len));
+  auto result = arrow::util::base64_decode(std::string_view(in, in_len));
+  if (!result.ok()) {
+    gdv_fn_context_set_error_msg(context, result.status().message().c_str());
+    *out_len = 0;
+    return "";
+  }
+
+  std::string decoded_str = *result;
+
   *out_len = static_cast<int32_t>(decoded_str.length());
   // allocate memory for response
   char* ret = reinterpret_cast<char*>(

cpp/src/parquet/arrow/fuzz_internal.cc

@@ -83,8 +83,11 @@ class FuzzDecryptionKeyRetriever : public DecryptionKeyRetriever {
     }
     // Is it a key generated by MakeEncryptionKey?
     if (key_id.starts_with(kInlineKeyPrefix)) {
-      return SecureString(
+      PARQUET_ASSIGN_OR_THROW(
+          auto decoded_key,
           ::arrow::util::base64_decode(key_id.substr(kInlineKeyPrefix.length())));
+
+      return SecureString(std::move(decoded_key));
     }
     throw ParquetException("Unknown fuzz encryption key_id");
   }

cpp/src/parquet/arrow/schema.cc

@@ -953,7 +953,8 @@ Status GetOriginSchema(const std::shared_ptr<const KeyValueMetadata>& metadata,
   // The original Arrow schema was serialized using the store_schema option.
   // We deserialize it here and use it to inform read options such as
   // dictionary-encoded fields.
-  auto decoded = ::arrow::util::base64_decode(metadata->value(schema_index));
+  ARROW_ASSIGN_OR_RAISE(auto decoded,
+                        ::arrow::util::base64_decode(metadata->value(schema_index)));
   auto schema_buf = std::make_shared<Buffer>(decoded);
 
   ::arrow::ipc::DictionaryMemo dict_memo;

cpp/src/parquet/encryption/file_key_unwrapper.cc

@@ -122,7 +122,12 @@ KeyWithMasterId FileKeyUnwrapper::GetDataEncryptionKey(const KeyMaterial& key_ma
         });
 
     // Decrypt the data key
-    std::string aad = ::arrow::util::base64_decode(encoded_kek_id);
+    auto result = ::arrow::util::base64_decode(encoded_kek_id);
+    if (!result.ok()) {
+      throw ParquetException(result.status().message());
+    }
+
+    std::string aad = std::move(result).ValueOrDie();
     data_key = internal::DecryptKeyLocally(encoded_wrapped_dek, kek_bytes, aad);
   }
 

cpp/src/parquet/encryption/key_toolkit_internal.cc

@@ -52,7 +52,12 @@ std::string EncryptKeyLocally(const SecureString& key_bytes,
 
 SecureString DecryptKeyLocally(const std::string& encoded_encrypted_key,
                                const SecureString& master_key, const std::string& aad) {
-  std::string encrypted_key = ::arrow::util::base64_decode(encoded_encrypted_key);
+  auto result = ::arrow::util::base64_decode(encoded_encrypted_key);
+  if (!result.ok()) {
+    throw ParquetException(result.status().message());
+  }
+
+  std::string encrypted_key = std::move(result).ValueOrDie();
 
   AesDecryptor key_decryptor(ParquetCipher::AES_GCM_V1,
                              static_cast<int>(master_key.size()), false,