Bound decompressed block sizes and read compressed blocks as they arrive

clickhouse/base/compressed.cpp

@@ -23,6 +23,12 @@ enum class CompressionMethodByte : uint8_t {
 // Documentation says that compression is faster when output buffer is larger than LZ4_compressBound/ZSTD_compressBound estimation.
 constexpr size_t EXTRA_COMPRESS_BUFFER_SIZE = 4096;
 constexpr size_t DBMS_MAX_COMPRESSED_SIZE = 0x40000000ULL;   // 1GB
+// The same limits as ClickHouse applies to the sizes in a block header (src/Compression/CompressionInfo.h).
+constexpr size_t DBMS_MAX_DECOMPRESSED_SIZE = 0x40000000ULL;  // 1GB
+// The compressed payload is read in steps starting at this size, so that a declared
+// size is not allocated before the bytes behind it arrive. It covers the server's
+// default 1 MiB blocks even when they do not compress.
+constexpr size_t COMPRESSED_READ_STEP = 4 << 20;
 }
 
 namespace clickhouse {
@@ -45,7 +51,7 @@ CompressedInput::~CompressedInput() {
 }
 
 size_t CompressedInput::DoNext(const void** ptr, size_t len) {
-    if (mem_.Exhausted()) {
+    while (mem_.Exhausted()) {
         if (!Decompress()) {
             return 0;
         }
@@ -86,7 +92,12 @@ bool CompressedInput::Decompress() {
         throw CompressionError("compressed block size " + std::to_string(compressed) + " is smaller than header, data is likely corrupted");
     }
 
-    Buffer tmp(compressed);
+    if (original > DBMS_MAX_DECOMPRESSED_SIZE) {
+        throw CompressionError("decompressed block size " + std::to_string(original) +
+                               " exceeds maximum allowed, data is likely corrupted");
+    }
+
+    Buffer tmp;
 
     // Data header
     {
@@ -97,34 +108,62 @@ bool CompressedInput::Decompress() {
         out.Flush();
     }
 
-    if (!WireFormat::ReadBytes(*input_, tmp.data() + HEADER_SIZE, compressed - HEADER_SIZE)) {
-        return false;
-    } else {
-        if (hash != cityhash::CityHash128((const char*)tmp.data(), compressed)) {
-            throw CompressionError("data was corrupted");
+    while (tmp.size() < compressed) {
+        const size_t loaded = tmp.size();
+        if (tmp.capacity() <= loaded) {
+            // Eightfold, without overflowing a 32-bit size_t near the 1 GiB limit.
+            const size_t target = loaded > compressed / 8 ? compressed : std::max(COMPRESSED_READ_STEP, loaded * 8);
+            tmp.reserve(std::min<size_t>(compressed, target));
         }
+        const size_t step = std::min({compressed - loaded, tmp.capacity() - loaded, std::max(loaded, COMPRESSED_READ_STEP)});
+        tmp.resize(loaded + step);
+        if (!WireFormat::ReadBytes(*input_, tmp.data() + loaded, step)) {
+            return false;
+        }
+    }
+
+    if (hash != cityhash::CityHash128((const char*)tmp.data(), compressed)) {
+        throw CompressionError("data was corrupted");
+    }
+
+    // Left uninitialized: zero-filling would commit the whole declared size even when
+    // the payload decompresses to far less, which the length checks below reject.
+    if (original > data_capacity_) {
+        data_.reset(new uint8_t[original]);
+        data_capacity_ = original;
     }
 
-    data_ = Buffer(original);
+    // data_ is null until a non-empty block arrives, and the decompressors do
+    // pointer arithmetic on the destination even when its capacity is zero.
+    char empty_output  = 0;
+    char* const output = original > 0 ? reinterpret_cast<char*>(data_.get()) : &empty_output;
 
     switch (method) {
     case static_cast<uint8_t>(CompressionMethodByte::LZ4): {
-        if (LZ4_decompress_safe((const char*)tmp.data() + HEADER_SIZE, (char*)data_.data(), static_cast<int>(compressed - HEADER_SIZE), original) < 0) {
+        const int res =
+            LZ4_decompress_safe((const char*)tmp.data() + HEADER_SIZE, output, static_cast<int>(compressed - HEADER_SIZE), original);
+        if (res < 0) {
             throw CompressionError("can't decompress LZ4-encoded data");
-        } else {
-            mem_.Reset(data_.data(), original);
         }
+        if (static_cast<uint32_t>(res) != original) {
+            throw CompressionError("LZ4-encoded block decompressed to " + std::to_string(res) + " bytes, expected " +
+                                   std::to_string(original));
+        }
+        mem_.Reset(data_.get(), original);
         return true;
     }
 
     case static_cast<uint8_t>(CompressionMethodByte::ZSTD): {
-        size_t res = ZSTD_decompress((char*)data_.data(), original, (const char*)tmp.data() + HEADER_SIZE, static_cast<int>(compressed - HEADER_SIZE));
+        size_t res = ZSTD_decompress(output, original, (const char*)tmp.data() + HEADER_SIZE, static_cast<int>(compressed - HEADER_SIZE));
 
         if (ZSTD_isError(res)) {
             throw CompressionError("can't decompress ZSTD-encoded data, ZSTD error: " + std::string(ZSTD_getErrorName(res)));
-        } else {
-            mem_.Reset(data_.data(), original);
         }
+        if (res != original) {
+            throw CompressionError("ZSTD-encoded block decompressed to " + std::to_string(res) + " bytes, expected " +
+                                   std::to_string(original));
+        }
+        mem_.Reset(data_.get(), original);
         return true;
     }
 

clickhouse/base/compressed.h

@@ -21,7 +21,8 @@ class CompressedInput : public ZeroCopyInput {
 private:
     InputStream* const input_;
 
-    Buffer data_;
+    std::unique_ptr<uint8_t[]> data_;
+    size_t data_capacity_ = 0;
     ArrayInput mem_;
 };
 

ut/BUILD.bazel

@@ -51,6 +51,7 @@ cc_test(
         "types_ut.cpp",
         "utils_ut.cpp",
         # Test entry point and shared support code.
+        "allocation_tracker.cpp",
         "main.cpp",
         "tcp_server.cpp",
         "tcp_server.h",

ut/CMakeLists.txt

@@ -31,6 +31,7 @@ SET ( clickhouse-cpp-ut-src
 
     bignum_samples.cpp
     utils.cpp
+    allocation_tracker.cpp
     value_generators.cpp
     low_cardinality_nullable_tests.cpp
 

ut/allocation_tracker.cpp

@@ -0,0 +1,51 @@
+#include <clickhouse/columns/numeric.h>
+
+#include <cstdlib>
+#include <new>
+
+#include "utils.h"
+
+namespace {
+thread_local bool track_allocations    = false;
+thread_local size_t largest_allocation = 0;
+}  // namespace
+
+// The default operator new[], operator delete[], and nothrow forms forward to these.
+void* operator new(std::size_t size) {
+    if (track_allocations && size > largest_allocation) {
+        largest_allocation = size;
+    }
+    if (void* ptr = std::malloc(size ? size : 1)) {
+        return ptr;
+    }
+    throw std::bad_alloc();
+}
+
+void operator delete(void* ptr) noexcept {
+    std::free(ptr);
+}
+
+void operator delete(void* ptr, std::size_t) noexcept {
+    std::free(ptr);
+}
+
+LargestAllocationScope::LargestAllocationScope() {
+    largest_allocation = 0;
+    track_allocations  = true;
+}
+
+LargestAllocationScope::~LargestAllocationScope() {
+    track_allocations = false;
+}
+
+size_t LargestAllocationScope::Get() const {
+    return largest_allocation;
+}
+
+bool LargestAllocationScope::ObservesLibraryAllocations() {
+    constexpr size_t rows = 1 << 20;
+    LargestAllocationScope allocations;
+    clickhouse::ColumnUInt64 column;
+    column.Reserve(rows);
+    return allocations.Get() >= rows * sizeof(uint64_t);
+}

ut/stream_ut.cpp

@@ -1,10 +1,15 @@
-#include <clickhouse/base/wire_format.h>
-#include <clickhouse/base/output.h>
-#include <clickhouse/base/input.h>
+#include <city.h>
 #include <clickhouse/base/compressed.h>
-
+#include <clickhouse/base/input.h>
+#include <clickhouse/base/output.h>
+#include <clickhouse/base/wire_format.h>
 #include <gtest/gtest.h>
 
+#include <string>
+#include <vector>
+
+#include "utils.h"
+
 using namespace clickhouse;
 
 namespace {
@@ -26,6 +31,43 @@ Buffer MakeCompressedBlockHeader(uint32_t compressed_size) {
     return buffer;
 }
 
+void WriteCompressedFrame(OutputStream& output, uint8_t method, uint32_t original_size, const std::vector<uint8_t>& payload) {
+    Buffer frame;
+    {
+        BufferOutput frame_output(&frame);
+        const uint32_t compressed_size = static_cast<uint32_t>(9 + payload.size());
+        WireFormat::WriteFixed(frame_output, method);
+        WireFormat::WriteFixed(frame_output, compressed_size);
+        WireFormat::WriteFixed(frame_output, original_size);
+        WireFormat::WriteBytes(frame_output, payload.data(), payload.size());
+        frame_output.Flush();
+    }
+
+    WireFormat::WriteFixed(output, cityhash::CityHash128(reinterpret_cast<const char*>(frame.data()), frame.size()));
+    WireFormat::WriteBytes(output, frame.data(), frame.size());
+}
+
+constexpr uint8_t kLZ4  = 0x82;
+constexpr uint8_t kZSTD = 0x90;
+
+// Encodings of an empty input: a single LZ4 token with no literals, and a
+// ZSTD frame with Frame_Content_Size 0 holding one empty raw block.
+const std::vector<uint8_t> kEmptyLZ4  = {0x00};
+const std::vector<uint8_t> kEmptyZSTD = {0x28, 0xB5, 0x2F, 0xFD, 0x20, 0x00, 0x01, 0x00, 0x00};
+
+constexpr size_t kMaxAllocationForTruncatedBlock = 8 * 1024 * 1024;
+
+Buffer MakePseudoRandomData(size_t size) {
+    Buffer data(size);
+    uint64_t state = 0x9E3779B97F4A7C15ULL;
+    for (auto& byte : data) {
+        state ^= state << 13;
+        state ^= state >> 7;
+        state ^= state << 17;
+        byte = static_cast<uint8_t>(state);
+    }
+    return data;
+}
 }
 
 TEST(CodedStreamCase, Varint64) {
@@ -56,3 +98,127 @@ TEST(CompressedInputCase, RejectsBlocksShorterThanHeader) {
             << "compressed size: " << compressed_size;
     }
 }
+
+TEST(CompressedInputCase, SkipsEmptyBlocks) {
+    for (const auto& [method, empty_payload] : {std::make_pair(kLZ4, kEmptyLZ4), std::make_pair(kZSTD, kEmptyZSTD)}) {
+        Buffer buffer;
+        {
+            BufferOutput output(&buffer);
+            WriteCompressedFrame(output, method, 0, empty_payload);
+            CompressedOutput compressed_output(&output, 0, method == kLZ4 ? CompressionMethod::LZ4 : CompressionMethod::ZSTD);
+            WireFormat::WriteBytes(compressed_output, "abc", 3);
+            compressed_output.Flush();
+        }
+
+        ArrayInput input(buffer.data(), buffer.size());
+        CompressedInput compressed_input(&input);
+        std::string data(3, '\0');
+
+        ASSERT_TRUE(WireFormat::ReadBytes(compressed_input, data.data(), data.size())) << "method: " << int(method);
+        EXPECT_EQ(data, "abc");
+    }
+}
+
+TEST(CompressedInputCase, RejectsEmptyBlockWithNonEmptyContent) {
+    // LZ4: a 4-byte literal run. ZSTD: Frame_Content_Size 1 with a 1-byte raw block.
+    const std::vector<uint8_t> lz4_payload  = {0x40, 0x61, 0x62, 0x63, 0x64};
+    const std::vector<uint8_t> zstd_payload = {0x28, 0xB5, 0x2F, 0xFD, 0x20, 0x01, 0x09, 0x00, 0x00, 0x61};
+
+    for (const auto& [method, payload] : {std::make_pair(kLZ4, lz4_payload), std::make_pair(kZSTD, zstd_payload)}) {
+        Buffer buffer;
+        {
+            BufferOutput output(&buffer);
+            WriteCompressedFrame(output, method, 0, payload);
+            output.Flush();
+        }
+
+        ArrayInput input(buffer.data(), buffer.size());
+        CompressedInput compressed_input(&input);
+        const void* data = nullptr;
+
+        EXPECT_THROW(compressed_input.Next(&data, 1), CompressionError) << "method: " << int(method);
+    }
+}
+
+TEST(CompressedInputCase, RejectsBlockShorterThanDeclaredOriginalSize) {
+    // Both payloads decode to "abc", but the frame header declares 4 bytes.
+    const std::vector<uint8_t> lz4_payload  = {0x30, 0x61, 0x62, 0x63};
+    const std::vector<uint8_t> zstd_payload = {0x28, 0xB5, 0x2F, 0xFD, 0x20, 0x03, 0x19, 0x00, 0x00, 0x61, 0x62, 0x63};
+
+    for (const auto& [method, payload] : {std::make_pair(kLZ4, lz4_payload), std::make_pair(kZSTD, zstd_payload)}) {
+        Buffer buffer;
+        {
+            BufferOutput output(&buffer);
+            WriteCompressedFrame(output, method, 4, payload);
+            output.Flush();
+        }
+
+        ArrayInput input(buffer.data(), buffer.size());
+        CompressedInput compressed_input(&input);
+        const void* data = nullptr;
+
+        EXPECT_THROW(compressed_input.Next(&data, 4), CompressionError) << "method: " << int(method);
+    }
+}
+
+TEST(CompressedInputCase, RejectsDecompressedSizeAboveLimitBeforeAllocating) {
+    SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED();
+    Buffer buffer;
+    {
+        BufferOutput output(&buffer);
+        // An LZ4 literal 'A', declared as 1 GiB + 1 byte.
+        WriteCompressedFrame(output, kLZ4, 0x40000001U, {0x10, 'A'});
+        output.Flush();
+    }
+    ArrayInput input(buffer.data(), buffer.size());
+    CompressedInput compressed_input(&input);
+    const void* data = nullptr;
+
+    LargestAllocationScope allocations;
+    EXPECT_THROW(compressed_input.Next(&data, 1), CompressionError);
+    EXPECT_LT(allocations.Get(), kMaxAllocationForTruncatedBlock);
+}
+
+TEST(CompressedInputCase, TruncatedBlockDoesNotAllocateDeclaredCompressedSize) {
+    SKIP_UNLESS_LIBRARY_ALLOCATIONS_OBSERVED();
+    Buffer buffer = MakeCompressedBlockHeader(128 * 1024 * 1024);
+    buffer.insert(buffer.end(), 16, 'A');
+    ArrayInput input(buffer.data(), buffer.size());
+    CompressedInput compressed_input(&input);
+    const void* data = nullptr;
+
+    LargestAllocationScope allocations;
+    EXPECT_EQ(compressed_input.Next(&data, 1), 0u);
+    EXPECT_LT(allocations.Get(), kMaxAllocationForTruncatedBlock);
+}
+
+TEST(CompressedInputCase, LargeBlocksRoundTrip) {
+    // Incompressible, so the large block is compressed to more than twice the first
+    // read step and the compressed payload has to be read in several growing steps.
+    const Buffer large = MakePseudoRandomData(9 * 1024 * 1024 + 17);
+    const Buffer small = MakePseudoRandomData(1000);
+
+    for (const auto method : {CompressionMethod::LZ4, CompressionMethod::ZSTD}) {
+        Buffer buffer;
+        {
+            BufferOutput output(&buffer);
+            CompressedOutput compressed_output(&output, 0, method);
+            WireFormat::WriteBytes(compressed_output, small.data(), small.size());
+            WireFormat::WriteBytes(compressed_output, large.data(), large.size());
+            WireFormat::WriteBytes(compressed_output, small.data(), small.size());
+            compressed_output.Flush();
+            output.Flush();
+        }
+
+        ArrayInput input(buffer.data(), buffer.size());
+        CompressedInput compressed_input(&input);
+        Buffer first(small.size()), second(large.size()), third(small.size());
+        ASSERT_TRUE(WireFormat::ReadBytes(compressed_input, first.data(), first.size()));
+        ASSERT_TRUE(WireFormat::ReadBytes(compressed_input, second.data(), second.size()));
+        ASSERT_TRUE(WireFormat::ReadBytes(compressed_input, third.data(), third.size()));
+        EXPECT_EQ(first, small);
+        EXPECT_EQ(second, large);
+        EXPECT_EQ(third, small);
+        EXPECT_TRUE(input.Exhausted());
+    }
+}