lz4 decompressor has been fixed (#6414)

Author: dorooleg

ydb/library/yql/providers/s3/compressors/lz4io.cpp

@@ -22,7 +22,7 @@ constexpr ui32 LegacyMagicNumber = 0x184C2102U;
 constexpr size_t LegacyBlockSize = 8_MB;
 constexpr size_t FrameMaxBlockSize = 4_MB;
 
-void WriteLE32 (void* p, ui32 value32)
+void WriteLE32(void* p, ui32 value32)
 {
     const auto dstPtr = static_cast<unsigned char*>(p);
     dstPtr[0] = (unsigned char)value32;
@@ -31,7 +31,7 @@ void WriteLE32 (void* p, ui32 value32)
     dstPtr[3] = (unsigned char)(value32 >> 24U);
 }
 
-ui32 ReadLE32 (const void* s) {
+ui32 ReadLE32(const void* s) {
     const auto srcPtr = static_cast<const unsigned char*>(s);
     ui32 value32 = srcPtr[0];
     value32 += (ui32)srcPtr[1] <<  8U;
@@ -112,29 +112,35 @@ bool TReadBuffer::nextImpl() {
 }
 
 size_t TReadBuffer::DecompressFrame() {
-    if (NextToLoad > InBuffer.size()) {
-        InBuffer.resize(NextToLoad);
-    }
-
-    if (Pos >= Remaining) {
-        for (auto toRead = NextToLoad; toRead > 0U;) {
-            const auto sizeCheck = Source.read(InBuffer.data() + NextToLoad - toRead, toRead);
-            YQL_ENSURE(sizeCheck > 0U && sizeCheck <= toRead, "Cannot access compressed block.");
-            toRead -= sizeCheck;
+    while (NextToLoad) {
+        if (NextToLoad > InBuffer.size()) {
+            InBuffer.resize(NextToLoad);
         }
 
-        Pos = 0ULL;
-        Remaining = NextToLoad;
-    }
+        if (Pos >= Remaining) {
+            for (auto toRead = NextToLoad; toRead > 0U;) {
+                const auto sizeCheck = Source.read(InBuffer.data() + NextToLoad - toRead, toRead);
+                YQL_ENSURE(sizeCheck > 0U && sizeCheck <= toRead, "Cannot access compressed block.");
+                toRead -= sizeCheck;
+            }
+
+            Pos = 0ULL;
+            Remaining = NextToLoad;
+        }
 
-    if (Pos < Remaining) {
         auto decodedBytes = OutBuffer.size();
-        NextToLoad = LZ4F_decompress_usingDict(Ctx, OutBuffer.data(), &decodedBytes, InBuffer.data() + Pos, &Remaining, nullptr, 0ULL, nullptr);
-        YQL_ENSURE(!LZ4F_isError(NextToLoad), "Decompression error: " << LZ4F_getErrorName(NextToLoad));
-        Pos += Remaining;
+        while (Pos < Remaining || (decodedBytes == OutBuffer.size())) {
+            decodedBytes = OutBuffer.size();
+            NextToLoad = LZ4F_decompress_usingDict(Ctx, OutBuffer.data(), &decodedBytes, InBuffer.data() + Pos, &Remaining, nullptr, 0ULL, nullptr);
+            YQL_ENSURE(!LZ4F_isError(NextToLoad), "Decompression error: " << LZ4F_getErrorName(NextToLoad));
+            Pos += Remaining;
+
+            if (decodedBytes)
+                return decodedBytes;
 
-        if (decodedBytes)
-            return decodedBytes;
+            if (!NextToLoad)
+                return decodedBytes;
+        }
     }
 
     return 0ULL;

ydb/library/yql/providers/s3/compressors/ut/decompressor_ut.cpp

@@ -0,0 +1,45 @@
+#include <ydb/library/yql/providers/s3/compressors/lz4io.h>
+#include <ydb/library/yql/udfs/common/clickhouse/client/src/IO/ReadBufferFromFile.h>
+
+#include <library/cpp/scheme/scheme.h>
+#include <library/cpp/testing/common/env.h>
+#include <library/cpp/testing/unittest/registar.h>
+
+namespace NYql::NCompressors {
+
+namespace {
+    TString GetResourcePath(const TString& path) {
+        return ArcadiaSourceRoot() + "/ydb/library/yql/providers/s3/compressors/ut/test_compression_data/" + path;
+    }
+}
+
+Y_UNIT_TEST_SUITE(TCompressorTests) {
+    Y_UNIT_TEST(SuccessLz4) {
+        NDB::ReadBufferFromFile buffer(GetResourcePath("test.json.lz4"));
+        auto decompressorBuffer = std::make_unique<NLz4::TReadBuffer>(buffer);
+
+        char str[256] = {};
+        decompressorBuffer->read(str, 256);
+        UNIT_ASSERT_VALUES_EQUAL(NSc::TValue::FromJsonThrow(str), NSc::TValue::FromJsonThrow(R"([
+            {
+                "id": 0,
+                "description": "yq",
+                "info": "abc"
+            }
+            ])"));
+    }
+
+    Y_UNIT_TEST(WrongMagicLz4) {
+        NDB::ReadBufferFromFile buffer(GetResourcePath("test.json"));
+        UNIT_ASSERT_EXCEPTION_CONTAINS(std::make_unique<NLz4::TReadBuffer>(buffer), yexception, "TReadBuffer(): requirement StreamType != EStreamType::Unknown failed, message: Wrong magic.");
+    }
+
+    Y_UNIT_TEST(ErrorLz4) {
+        NDB::ReadBufferFromFile buffer(GetResourcePath("test.broken.lz4"));
+        auto decompressorBuffer = std::make_unique<NLz4::TReadBuffer>(buffer);
+        char str[256] = {};
+        UNIT_ASSERT_EXCEPTION_CONTAINS(decompressorBuffer->read(str, 256), yexception, "DecompressFrame(): requirement !LZ4F_isError(NextToLoad) failed, message: Decompression error: ERROR_reservedFlag_set");
+    }
+}
+
+}

ydb/library/yql/providers/s3/compressors/ut/test_compression_data/test.broken.lz4

@@ -0,0 +1,8 @@
+[
+{
+    "id": 0,
+    "description": "yq",
+    "info": "abc"
+}
+]
+

ydb/library/yql/providers/s3/compressors/ut/test_compression_data/test.json

@@ -0,0 +1,24 @@
+IF (NOT OS_WINDOWS AND CLANG AND NOT WITH_VALGRIND)
+
+UNITTEST_FOR(ydb/library/yql/providers/s3/compressors)
+
+SRCS(
+    decompressor_ut.cpp
+)
+
+PEERDIR(
+    library/cpp/scheme
+    ydb/library/yql/public/udf/service/stub
+    ydb/library/yql/udfs/common/clickhouse/client
+)
+
+ADDINCL(
+    ydb/library/yql/udfs/common/clickhouse/client/base
+    ydb/library/yql/udfs/common/clickhouse/client/base/pcg-random
+    ydb/library/yql/udfs/common/clickhouse/client/src
+)
+
+END()
+
+ENDIF()
+

ydb/library/yql/providers/s3/compressors/ut/test_compression_data/test.json.lz4

@@ -36,3 +36,6 @@ YQL_LAST_ABI_VERSION()
 
 END()
 
+RECURSE_FOR_TESTS(
+    ut
+)