Merge pull request #1891 from bimbashrestha/oss

[fuzz] Superblock fuzz issues

lib/compress/zstd_compress.c

@@ -2450,65 +2450,70 @@ out:
     return cSize;
 }
 
-static size_t ZSTD_compressBlock_targetCBlockSize(ZSTD_CCtx* zc,
+static void ZSTD_confirmRepcodesAndEntropyTables(ZSTD_CCtx* zc)
-                               void* dst, size_t dstCapacity,
+{
-                               const void* src, size_t srcSize,
-                               U32 lastBlock) {
-    size_t cSize = 0;
-    DEBUGLOG(5, "ZSTD_compressBlock_targetCBlockSize (dstCapacity=%u, dictLimit=%u, nextToUpdate=%u, srcSize=%zu)",
-                (unsigned)dstCapacity, (unsigned)zc->blockState.matchState.window.dictLimit, (unsigned)zc->blockState.matchState.nextToUpdate, srcSize);
-
-    {   const size_t bss = ZSTD_buildSeqStore(zc, src, srcSize);
-        FORWARD_IF_ERROR(bss);
-        if (bss == ZSTDbss_compress) {
-            cSize = ZSTD_compressSuperBlock(zc, dst, dstCapacity, lastBlock);
-    }   }
-
-    /* Superblock compression may fail, in which case
-     * encode using ZSTD_noCompressSuperBlock writing sub blocks
-     * in uncompressed mode.
-     */
-    if (cSize == 0) {
-        cSize = ZSTD_noCompressSuperBlock(dst, dstCapacity, src, srcSize, zc->appliedParams.targetCBlockSize, lastBlock);
-        /* In compression, there is an assumption that a compressed block is always
-         * within the size of ZSTD_compressBound(). However, SuperBlock compression
-         * can exceed the limit due to overhead of headers from SubBlocks.
-         * This breaks in streaming mode where output buffer in compress context is
-         * allocated ZSTD_compressBound() amount of memory, which may not be big
-         * enough for SuperBlock compression.
-         * In such case, fall back to normal compression. This is possible because
-         * targetCBlockSize is best effort not a guarantee. */
-        if (cSize != ERROR(dstSize_tooSmall)) return cSize;
-        else {
-            BYTE* const ostart = (BYTE*)dst;
-            /* If ZSTD_noCompressSuperBlock fails with dstSize_tooSmall,
-             * compress normally.
-             */
-            cSize = ZSTD_compressSequences(&zc->seqStore,
-                    &zc->blockState.prevCBlock->entropy, &zc->blockState.nextCBlock->entropy,
-                    &zc->appliedParams,
-                    ostart+ZSTD_blockHeaderSize, dstCapacity-ZSTD_blockHeaderSize,
-                    srcSize,
-                    zc->entropyWorkspace, HUF_WORKSPACE_SIZE /* statically allocated in resetCCtx */,
-                    zc->bmi2);
-            if (!ZSTD_isError(cSize) && cSize != 0) {
-                U32 const cBlockHeader24 = lastBlock + (((U32)bt_compressed)<<1) + (U32)(cSize << 3);
-                MEM_writeLE24(ostart, cBlockHeader24);
-                cSize += ZSTD_blockHeaderSize;
-            }
-        }
-    }
-
-    if (!ZSTD_isError(cSize) && cSize != 0) {
-        /* confirm repcodes and entropy tables when emitting a compressed block */
     ZSTD_compressedBlockState_t* const tmp = zc->blockState.prevCBlock;
     zc->blockState.prevCBlock = zc->blockState.nextCBlock;
     zc->blockState.nextCBlock = tmp;
+}
+
+static size_t ZSTD_compressBlock_targetCBlockSize_body(ZSTD_CCtx* zc,
+                               void* dst, size_t dstCapacity,
+                               const void* src, size_t srcSize,
+                               const size_t bss, U32 lastBlock)
+{
+    /* Attempt superblock compression and return early if successful */
+    if (bss == ZSTDbss_compress) {
+        size_t const cSize = ZSTD_compressSuperBlock(zc, dst, dstCapacity, lastBlock);
+        FORWARD_IF_ERROR(cSize);
+        if (cSize != 0) {
+            ZSTD_confirmRepcodesAndEntropyTables(zc);
+            return cSize;
         }
-    /* We check that dictionaries have offset codes available for the first
+    }
-     * block. After the first block, the offcode table might not have large
+
-     * enough codes to represent the offsets in the data.
+    /* Superblock compression failed, attempt to emit noCompress superblocks
-     */
+     * and return early if that is successful and we have enough room for checksum */
+    {
+        size_t const cSize = ZSTD_noCompressSuperBlock(dst, dstCapacity, src, srcSize, zc->appliedParams.targetCBlockSize, lastBlock);
+        if (cSize != ERROR(dstSize_tooSmall) && (dstCapacity - cSize) >= 4)
+            return cSize;
+    }
+
+    /* noCompress superblock emission failed. Attempt to compress normally
+     * and return early if that is successful */
+    {
+        size_t const cSize = ZSTD_compressSequences(&zc->seqStore,
+            &zc->blockState.prevCBlock->entropy, &zc->blockState.nextCBlock->entropy,
+            &zc->appliedParams, (BYTE*)dst+ZSTD_blockHeaderSize, dstCapacity-ZSTD_blockHeaderSize,
+            srcSize, zc->entropyWorkspace, HUF_WORKSPACE_SIZE, zc->bmi2);
+        FORWARD_IF_ERROR(cSize);
+        if (cSize != 0) {
+            U32 const cBlockHeader24 = lastBlock + (((U32)bt_compressed)<<1) + (U32)(cSize << 3);
+            MEM_writeLE24((BYTE*)dst, cBlockHeader24);
+            ZSTD_confirmRepcodesAndEntropyTables(zc);
+            return cSize + ZSTD_blockHeaderSize;
+        }
+    }
+
+    /* Everything failed. Just emit a regular noCompress block */
+    return ZSTD_noCompressBlock(dst, dstCapacity, src, srcSize, lastBlock);
+}
+
+static size_t ZSTD_compressBlock_targetCBlockSize(ZSTD_CCtx* zc,
+                               void* dst, size_t dstCapacity,
+                               const void* src, size_t srcSize,
+                               U32 lastBlock)
+{
+    size_t cSize = 0;
+    const size_t bss = ZSTD_buildSeqStore(zc, src, srcSize);
+    DEBUGLOG(5, "ZSTD_compressBlock_targetCBlockSize (dstCapacity=%u, dictLimit=%u, nextToUpdate=%u, srcSize=%zu)",
+                (unsigned)dstCapacity, (unsigned)zc->blockState.matchState.window.dictLimit, (unsigned)zc->blockState.matchState.nextToUpdate, srcSize);
+    FORWARD_IF_ERROR(bss);
+
+    cSize = ZSTD_compressBlock_targetCBlockSize_body(zc, dst, dstCapacity, src, srcSize, bss, lastBlock);
+    FORWARD_IF_ERROR(cSize);
+
     if (zc->blockState.prevCBlock->entropy.fse.offcode_repeatMode == FSE_repeat_valid)
         zc->blockState.prevCBlock->entropy.fse.offcode_repeatMode = FSE_repeat_check;
 

tests/fuzzer.c

@@ -427,7 +427,6 @@ static int basicUnitTests(U32 const seed, double compressibility)
     }   }
     DISPLAYLEVEL(3, "OK \n");
 
-
     DISPLAYLEVEL(3, "test%3i : decompress with null dict : ", testNb++);
     {   ZSTD_DCtx* const dctx = ZSTD_createDCtx(); assert(dctx != NULL);
         {   size_t const r = ZSTD_decompress_usingDict(dctx,
@@ -490,6 +489,19 @@ static int basicUnitTests(U32 const seed, double compressibility)
     }
     DISPLAYLEVEL(3, "OK \n");
 
+    DISPLAYLEVEL(3, "test%3d: superblock enough room for checksum : ", testNb++)
+    {
+        /* This tests whether or not we leave enough room for the checksum at the end
+         * of the dst buffer. The bug that motivated this test was found by the
+         * stream_round_trip fuzzer but this crashes for the same reason and is
+         * far more compact than re-creating the stream_round_trip fuzzer's code path */
+        ZSTD_CCtx *cctx = ZSTD_createCCtx();
+        ZSTD_CCtx_setParameter(cctx, ZSTD_c_targetCBlockSize, 64);
+        assert(!ZSTD_isError(ZSTD_compress2(cctx, compressedBuffer, 1339, CNBuffer, 1278)));
+        ZSTD_freeCCtx(cctx);
+    }
+    DISPLAYLEVEL(3, "OK \n");
+  
     DISPLAYLEVEL(3, "test%3i : compress a NULL input with each level : ", testNb++);
     {   int level = -1;
         ZSTD_CCtx* cctx = ZSTD_createCCtx();