diff --git a/clickhouse/base/compressed.cpp b/clickhouse/base/compressed.cpp index 7820662d..1d1b49e9 100644 --- a/clickhouse/base/compressed.cpp +++ b/clickhouse/base/compressed.cpp @@ -45,7 +45,7 @@ CompressedInput::~CompressedInput() { } size_t CompressedInput::DoNext(const void** ptr, size_t len) { - if (mem_.Exhausted()) { + while (mem_.Exhausted()) { if (!Decompress()) { return 0; } @@ -107,24 +107,37 @@ bool CompressedInput::Decompress() { data_ = Buffer(original); + // Buffer(0).data() is null, 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(data_.data()) : &empty_output; + switch (method) { case static_cast(CompressionMethodByte::LZ4): { - if (LZ4_decompress_safe((const char*)tmp.data() + HEADER_SIZE, (char*)data_.data(), static_cast(compressed - HEADER_SIZE), original) < 0) { + const int res = + LZ4_decompress_safe((const char*)tmp.data() + HEADER_SIZE, output, static_cast(compressed - HEADER_SIZE), original); + if (res < 0) { throw CompressionError("can't decompress LZ4-encoded data"); - } else { - mem_.Reset(data_.data(), original); } + if (static_cast(res) != original) { + throw CompressionError("LZ4-encoded block decompressed to " + std::to_string(res) + " bytes, expected " + + std::to_string(original)); + } + mem_.Reset(data_.data(), original); return true; } case static_cast(CompressionMethodByte::ZSTD): { - size_t res = ZSTD_decompress((char*)data_.data(), original, (const char*)tmp.data() + HEADER_SIZE, static_cast(compressed - HEADER_SIZE)); + size_t res = ZSTD_decompress(output, original, (const char*)tmp.data() + HEADER_SIZE, static_cast(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_.data(), original); return true; } diff --git a/ut/stream_ut.cpp b/ut/stream_ut.cpp index 0bcf669c..f1112bbe 100644 --- a/ut/stream_ut.cpp +++ b/ut/stream_ut.cpp @@ -1,10 +1,13 @@ -#include -#include -#include +#include #include - +#include +#include +#include #include +#include +#include + using namespace clickhouse; namespace { @@ -26,6 +29,29 @@ Buffer MakeCompressedBlockHeader(uint32_t compressed_size) { return buffer; } +void WriteCompressedFrame(OutputStream& output, uint8_t method, uint32_t original_size, const std::vector& payload) { + Buffer frame; + { + BufferOutput frame_output(&frame); + const uint32_t compressed_size = static_cast(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(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 kEmptyLZ4 = {0x00}; +const std::vector kEmptyZSTD = {0x28, 0xB5, 0x2F, 0xFD, 0x20, 0x00, 0x01, 0x00, 0x00}; } TEST(CodedStreamCase, Varint64) { @@ -56,3 +82,65 @@ 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 lz4_payload = {0x40, 0x61, 0x62, 0x63, 0x64}; + const std::vector 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 lz4_payload = {0x30, 0x61, 0x62, 0x63}; + const std::vector 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); + } +}