From f65b9e27ce0b6e4ed096126659021359d004d1ab Mon Sep 17 00:00:00 2001
From: Elliot Gorokhovsky <embg@fb.com>
Date: Tue, 12 Mar 2024 17:07:06 -0400
Subject: [PATCH] Exercise ZSTD_findDecompressedSize() in the simple
 decompression fuzzer (#3959)

* Improve decompression fuzzer

* Fix legacy frame header fuzzer crash, add unit test
---
 lib/decompress/zstd_decompress.c | 9 +++++++++
 tests/fuzz/simple_decompress.c   | 8 +++++++-
 tests/zstreamtest.c              | 9 +++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/lib/decompress/zstd_decompress.c b/lib/decompress/zstd_decompress.c
index 17305908540..42636d5780c 100644
--- a/lib/decompress/zstd_decompress.c
+++ b/lib/decompress/zstd_decompress.c
@@ -1093,6 +1093,15 @@ size_t ZSTD_decompressMultiFrame(ZSTD_DCtx* dctx,
             decodedSize = ZSTD_decompressLegacy(dst, dstCapacity, src, frameSize, dict, dictSize);
             if (ZSTD_isError(decodedSize)) return decodedSize;
 
+            {
+                unsigned long long const expectedSize = ZSTD_getFrameContentSize(src, srcSize);
+                RETURN_ERROR_IF(expectedSize == ZSTD_CONTENTSIZE_ERROR, corruption_detected, "Corrupted frame header!");
+                if (expectedSize != ZSTD_CONTENTSIZE_UNKNOWN) {
+                    RETURN_ERROR_IF(expectedSize != decodedSize, corruption_detected,
+                        "Frame header size does not match decoded size!");
+                }
+            }
+
             assert(decodedSize <= dstCapacity);
             dst = (BYTE*)dst + decodedSize;
             dstCapacity -= decodedSize;
diff --git a/tests/fuzz/simple_decompress.c b/tests/fuzz/simple_decompress.c
index ce5f9f0981b..0ee61902cdb 100644
--- a/tests/fuzz/simple_decompress.c
+++ b/tests/fuzz/simple_decompress.c
@@ -37,7 +37,13 @@ int LLVMFuzzerTestOneInput(const uint8_t *src, size_t size)
     size_t const bufSize = FUZZ_dataProducer_uint32Range(producer, 0, 10 * size);
     void *rBuf = FUZZ_malloc(bufSize);
 
-    ZSTD_decompressDCtx(dctx, rBuf, bufSize, src, size);
+    size_t const dSize = ZSTD_decompressDCtx(dctx, rBuf, bufSize, src, size);
+    if (!ZSTD_isError(dSize)) {
+        /* If decompression was successful, the content size from the frame header(s) should be valid. */
+        size_t const expectedSize = ZSTD_findDecompressedSize(src, size);
+        FUZZ_ASSERT(expectedSize != ZSTD_CONTENTSIZE_ERROR);
+        FUZZ_ASSERT(expectedSize == ZSTD_CONTENTSIZE_UNKNOWN || expectedSize == dSize);
+    }
     free(rBuf);
 
     FUZZ_dataProducer_free(producer);
diff --git a/tests/zstreamtest.c b/tests/zstreamtest.c
index 82aaf3db50c..7cc4068bc09 100644
--- a/tests/zstreamtest.c
+++ b/tests/zstreamtest.c
@@ -2408,6 +2408,15 @@ static int basicUnitTests(U32 seed, double compressibility, int bigTests)
     }
     DISPLAYLEVEL(3, "OK \n");
 
+    DISPLAYLEVEL(3, "test%3i : Decoder should reject invalid frame header on legacy frames: ", testNb++);
+    {
+        const unsigned char compressed[] = { 0x26,0xb5,0x2f,0xfd,0x50,0x91,0xfd,0xd8,0xb5 };
+        const size_t compressedSize = 9;
+        size_t const dSize = ZSTD_decompress(NULL, 0, compressed, compressedSize);
+        CHECK(!ZSTD_isError(dSize), "must reject when legacy frame header is invalid");
+    }
+    DISPLAYLEVEL(3, "OK \n");
+
 _end:
     FUZ_freeDictionary(dictionary);
     ZSTD_freeCStream(zc);