Skip to content

Validate what CompressedInput decompresses #586

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
iliaal wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Validate what CompressedInput decompresses #586

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 20 additions & 7 deletions clickhouse/base/compressed.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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<char*>(data_.data()) : &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_.data(), 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_.data(), original);
return true;
}

Expand Down
96 changes: 92 additions & 4 deletions ut/stream_ut.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
#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>

using namespace clickhouse;

namespace {
Expand All @@ -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<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};
}

TEST(CodedStreamCase, Varint64) {
Expand Down Expand Up @@ -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<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);
}
}
Loading