From f0f292754631679d47584ee167b9b32614498002 Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Fri, 18 Sep 2026 23:13:25 +0200
Subject: [PATCH] Fix compressed block size validation

---
 clickhouse/base/compressed.cpp |  4 ++++
 ut/stream_ut.cpp               | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/clickhouse/base/compressed.cpp b/clickhouse/base/compressed.cpp
index 06d72376c..f0a79d009 100644
--- a/clickhouse/base/compressed.cpp
+++ b/clickhouse/base/compressed.cpp
@@ -82,6 +82,10 @@ bool CompressedInput::Decompress() {
         throw CompressionError("compressed data too big");
     }
 
+    if (compressed < HEADER_SIZE) {
+        throw CompressionError("compressed data too small");
+    }
+
     Buffer tmp(compressed);
 
     // Data header
diff --git a/ut/stream_ut.cpp b/ut/stream_ut.cpp
index e139d8bce..0bcf669c8 100644
--- a/ut/stream_ut.cpp
+++ b/ut/stream_ut.cpp
@@ -1,11 +1,33 @@
 #include <clickhouse/base/wire_format.h>
 #include <clickhouse/base/output.h>
 #include <clickhouse/base/input.h>
+#include <clickhouse/base/compressed.h>
 
 #include <gtest/gtest.h>
 
 using namespace clickhouse;
 
+namespace {
+
+Buffer MakeCompressedBlockHeader(uint32_t compressed_size) {
+    Buffer buffer;
+    BufferOutput output(&buffer);
+
+    uint8_t hash[16] = {};
+    const uint8_t method = 0x82;
+    const uint32_t original_size = 0;
+
+    WireFormat::WriteFixed(output, hash);
+    WireFormat::WriteFixed(output, method);
+    WireFormat::WriteFixed(output, compressed_size);
+    WireFormat::WriteFixed(output, original_size);
+    output.Flush();
+
+    return buffer;
+}
+
+}
+
 TEST(CodedStreamCase, Varint64) {
     Buffer buf;
 
@@ -22,3 +44,15 @@ TEST(CodedStreamCase, Varint64) {
         ASSERT_EQ(value, 18446744071965638648ULL);
     }
 }
+
+TEST(CompressedInputCase, RejectsBlocksShorterThanHeader) {
+    for (const uint32_t compressed_size : {0U, 1U, 8U}) {
+        Buffer buffer = MakeCompressedBlockHeader(compressed_size);
+        ArrayInput input(buffer.data(), buffer.size());
+        CompressedInput compressed_input(&input);
+        const void* data = nullptr;
+
+        EXPECT_THROW(compressed_input.Next(&data, 1), CompressionError)
+            << "compressed size: " << compressed_size;
+    }
+}