Fixed : ZBUFF_compressEnd() called multiple times with too small dst buffer (#206)

diff --git a/lib/compress/zbuff_compress.c b/lib/compress/zbuff_compress.c
index 9c9c8059f0971e38dcdbd4ef1dbfbf97a521a3e1..d4ea217c25468b4b8a668f51f8d1bd7decc06af3 100644 (file)
--- a/lib/compress/zbuff_compress.c
+++ b/lib/compress/zbuff_compress.c
@@ -79,7 +79,7 @@ static size_t const ZBUFF_endFrameSize = ZSTD_BLOCKHEADERSIZE;
 *  output : ZSTD_compressBound(ZSTD_BLOCKSIZE_MAX) + ZSTD_blockHeaderSize + ZBUFF_endFrameSize : ensures it's always possible to write/flush/end a full block at best speed.
 * **************************************************/
 
-typedef enum { ZBUFFcs_init, ZBUFFcs_load, ZBUFFcs_flush } ZBUFF_cStage;
+typedef enum { ZBUFFcs_init, ZBUFFcs_load, ZBUFFcs_flush, ZBUFFcs_final } ZBUFF_cStage;
 
 /* *** Resources *** */
 struct ZBUFF_CCtx_s {
@@ -196,7 +196,7 @@ MEM_STATIC size_t ZBUFF_limitCopy(void* dst, size_t dstCapacity, const void* src
 static size_t ZBUFF_compressContinue_generic(ZBUFF_CCtx* zbc,
                               void* dst, size_t* dstCapacityPtr,
                         const void* src, size_t* srcSizePtr,
-                              int flush)   /* aggregate : wait for full block before compressing */
+                              int flush)
 {
     U32 notDone = 1;
     const char* const istart = (const char*)src;
@@ -243,15 +243,15 @@ static size_t ZBUFF_compressContinue_generic(ZBUFF_CCtx* zbc,
             }
 
         case ZBUFFcs_flush:
-            /* flush into dst */
+        case ZBUFFcs_final: /* flush into dst */
             {   size_t const toFlush = zbc->outBuffContentSize - zbc->outBuffFlushedSize;
                 size_t const flushed = ZBUFF_limitCopy(op, oend-op, zbc->outBuff + zbc->outBuffFlushedSize, toFlush);
                 op += flushed;
                 zbc->outBuffFlushedSize += flushed;
-                if (toFlush!=flushed) { notDone = 0; break; } /* not enough space within dst to store compressed block : stop there */
-                zbc->outBuffContentSize = 0;
-                zbc->outBuffFlushedSize = 0;
-                zbc->stage = ZBUFFcs_load;
+                if (toFlush!=flushed) { notDone = 0; break; } /* dst too small to store flushed data : stop there */
+                zbc->outBuffContentSize = zbc->outBuffFlushedSize = 0;
+                if (zbc->stage==ZBUFFcs_flush) { zbc->stage = ZBUFFcs_load; break; }
+                notDone=0;
                 break;
             }
         default:
@@ -291,19 +291,28 @@ size_t ZBUFF_compressEnd(ZBUFF_CCtx* zbc, void* dst, size_t* dstCapacityPtr)
     BYTE* const ostart = (BYTE*)dst;
     BYTE* const oend = ostart + *dstCapacityPtr;
     BYTE* op = ostart;
-    size_t outSize = *dstCapacityPtr;
-    size_t epilogueSize, remaining;
-    ZBUFF_compressFlush(zbc, dst, &outSize);     /* flush any remaining inBuff */
-    op += outSize;
-    epilogueSize = ZSTD_compressEnd(zbc->zc, zbc->outBuff + zbc->outBuffContentSize, zbc->outBuffSize - zbc->outBuffContentSize);   /* epilogue into outBuff */
-    zbc->outBuffContentSize += epilogueSize;
-    outSize = oend-op;
-    zbc->stage = ZBUFFcs_flush;
-    remaining = ZBUFF_compressFlush(zbc, op, &outSize);   /* attempt to flush epilogue into dst */
-    op += outSize;
-    if (!remaining) zbc->stage = ZBUFFcs_init;   /* close only if nothing left to flush */
-    *dstCapacityPtr = op-ostart;                 /* tells how many bytes were written */
-    return remaining;
+
+    {   size_t outSize = *dstCapacityPtr;
+        size_t const remainingToFlush = ZBUFF_compressFlush(zbc, dst, &outSize);
+        op += outSize;
+        if (remainingToFlush) {
+            *dstCapacityPtr = op-ostart;
+            return remainingToFlush;
+    }   }
+
+    if (zbc->stage == ZBUFFcs_final) { zbc->stage = ZBUFFcs_init; *dstCapacityPtr = op-ostart; return 0; }
+
+    /* outBuff is flushed */
+    {   size_t outSize = oend-op;
+        size_t remainingToFlush;
+        zbc->outBuffContentSize = ZSTD_compressEnd(zbc->zc, zbc->outBuff, zbc->outBuffSize); /* epilogue into outBuff */
+        zbc->stage = ZBUFFcs_flush;
+        remainingToFlush = ZBUFF_compressFlush(zbc, op, &outSize);
+        op += outSize;
+        zbc->stage = remainingToFlush ? ZBUFFcs_final : ZBUFFcs_init; /* close only if nothing left to flush */
+        *dstCapacityPtr = op-ostart;   /* how many bytes were written */
+        return remainingToFlush;
+    }
 }
 
 

diff --git a/programs/zbufftest.c b/programs/zbufftest.c
index cdf3abef8726d676b6efba52ab80b7b7a7dc3438..3ff55fc225fdcd4e53096b82345b90c0426b29b6 100644 (file)
--- a/programs/zbufftest.c
+++ b/programs/zbufftest.c
@@ -427,10 +427,9 @@ static int fuzzerTests(U32 seed, U32 nbTests, unsigned startTest, double compres
                 size_t dstBuffSize = MIN(cBufferSize - cSize, randomDstSize);
                 remainingToFlush = ZBUFF_compressEnd(zc, cBuffer+cSize, &dstBuffSize);
                 CHECK (ZBUFF_isError(remainingToFlush), "flush error : %s", ZBUFF_getErrorName(remainingToFlush));
-                DISPLAY("flush %u bytes : still within context : %i \n", (U32)dstBuffSize, (int)remainingToFlush);
+                //DISPLAY("flush %u bytes : still within context : %i \n", (U32)dstBuffSize, (int)remainingToFlush);
                 cSize += dstBuffSize;
-            }
-        }
+        }   }
         crcOrig = XXH64_digest(&xxhState);
 
         /* multi - fragments decompression test */