From 463f6467a448192fefd75b42b77505305b08894c Mon Sep 17 00:00:00 2001
From: Kouhei Sutou <kou@clear-code.com>
Date: Fri, 30 Nov 2018 14:43:24 +0900
Subject: [PATCH] [C++] Add support for decompressed buffer size check for
 Snappy

---
 cpp/src/arrow/test-util.h                | 11 +++++++++
 cpp/src/arrow/util/compression-test.cc   | 29 ++++++++++++++++++++++++
 cpp/src/arrow/util/compression_snappy.cc | 23 ++++++++++++-------
 3 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/cpp/src/arrow/test-util.h b/cpp/src/arrow/test-util.h
index 950cf02a0589..3011f287f096 100644
--- a/cpp/src/arrow/test-util.h
+++ b/cpp/src/arrow/test-util.h
@@ -65,6 +65,17 @@
     }                                                                     \
   } while (false)
 
+#define ASSERT_RAISES_WITH_MESSAGE(ENUM, message, expr)                   \
+  do {                                                                    \
+    ::arrow::Status s = (expr);                                           \
+    if (!s.Is##ENUM()) {                                                  \
+      FAIL() << "Expected '" STRINGIFY(expr) "' to fail with " STRINGIFY( \
+                    ENUM) ", but got "                                    \
+             << s.ToString();                                             \
+    }                                                                     \
+    ASSERT_EQ((message), s.ToString());                                   \
+  } while (false)
+
 #define ASSERT_OK(expr)                                               \
   do {                                                                \
     ::arrow::Status s = (expr);                                       \
diff --git a/cpp/src/arrow/util/compression-test.cc b/cpp/src/arrow/util/compression-test.cc
index f3e1106ef89e..e0e6f4837f20 100644
--- a/cpp/src/arrow/util/compression-test.cc
+++ b/cpp/src/arrow/util/compression-test.cc
@@ -345,6 +345,35 @@ TEST_P(CodecTest, CodecRoundtrip) {
   }
 }
 
+TEST_P(CodecTest, OutputBufferIsSmall) {
+  auto type = GetCompression();
+  if (type != Compression::SNAPPY) {
+    return;
+  }
+
+  std::unique_ptr<Codec> codec;
+  ASSERT_OK(Codec::Create(type, &codec));
+
+  vector<uint8_t> data = MakeRandomData(10);
+  auto max_compressed_len = codec->MaxCompressedLen(data.size(), data.data());
+  std::vector<uint8_t> compressed(max_compressed_len);
+  std::vector<uint8_t> decompressed(data.size() - 1);
+
+  int64_t actual_size;
+  ASSERT_OK(codec->Compress(data.size(), data.data(), max_compressed_len,
+                            compressed.data(), &actual_size));
+  compressed.resize(actual_size);
+
+  int64_t actual_decompressed_size;
+  std::stringstream ss;
+  ss << "Invalid: Output buffer size (" << decompressed.size() << ") must be "
+     << data.size() << " or larger.";
+  ASSERT_RAISES_WITH_MESSAGE(
+      Invalid, ss.str(),
+      codec->Decompress(compressed.size(), compressed.data(), decompressed.size(),
+                        decompressed.data(), &actual_decompressed_size));
+}
+
 TEST_P(CodecTest, StreamingCompressor) {
   if (GetCompression() == Compression::SNAPPY) {
     // SKIP: snappy doesn't support streaming compression
diff --git a/cpp/src/arrow/util/compression_snappy.cc b/cpp/src/arrow/util/compression_snappy.cc
index ab93c3ce32be..1b483e585520 100644
--- a/cpp/src/arrow/util/compression_snappy.cc
+++ b/cpp/src/arrow/util/compression_snappy.cc
@@ -19,6 +19,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <sstream>
 
 #include <snappy.h>
 
@@ -47,15 +48,21 @@ Status SnappyCodec::Decompress(int64_t input_len, const uint8_t* input,
 }
 
 Status SnappyCodec::Decompress(int64_t input_len, const uint8_t* input,
-                               int64_t ARROW_ARG_UNUSED(output_buffer_len),
-                               uint8_t* output_buffer, int64_t* output_len) {
+                               int64_t output_buffer_len, uint8_t* output_buffer,
+                               int64_t* output_len) {
+  size_t decompressed_size;
+  if (!snappy::GetUncompressedLength(reinterpret_cast<const char*>(input),
+                                     static_cast<size_t>(input_len),
+                                     &decompressed_size)) {
+    return Status::IOError("Corrupt snappy compressed data.");
+  }
+  if (output_buffer_len < static_cast<int64_t>(decompressed_size)) {
+    std::stringstream ss;
+    ss << "Output buffer size (" << output_buffer_len << ") must be " << decompressed_size
+       << " or larger.";
+    return Status::Invalid(ss.str());
+  }
   if (output_len) {
-    size_t decompressed_size;
-    if (!snappy::GetUncompressedLength(reinterpret_cast<const char*>(input),
-                                       static_cast<size_t>(input_len),
-                                       &decompressed_size)) {
-      return Status::IOError("Corrupt snappy compressed data.");
-    }
     *output_len = static_cast<int64_t>(decompressed_size);
   }
   if (!snappy::RawUncompress(reinterpret_cast<const char*>(input),