From 2f2639438a8809c1698128ea7aa427af4e782613 Mon Sep 17 00:00:00 2001
From: Yann Collet <yann.collet.73@gmail.com>
Date: Mon, 26 Sep 2016 14:06:08 +0200
Subject: [PATCH] zstreamtest can fuzztest pledgedSrcSize

---
 lib/decompress/zstd_decompress.c |  1 +
 tests/zstreamtest.c              | 18 ++++++++++--------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/lib/decompress/zstd_decompress.c b/lib/decompress/zstd_decompress.c
index 3410bbc0a..47b5f42c7 100644
--- a/lib/decompress/zstd_decompress.c
+++ b/lib/decompress/zstd_decompress.c
@@ -1554,6 +1554,7 @@ size_t ZSTD_initDStream(ZSTD_DStream* zds)
 
 size_t ZSTD_resetDStream(ZSTD_DStream* zds)
 {
+    if (zds->ddict == NULL) return ERROR(stage_wrong);  /* must be init at least once */
     zds->stage = zdss_loadHeader;
     zds->lhSize = zds->inPos = zds->outStart = zds->outEnd = 0;
     zds->legacyVersion = 0;
diff --git a/tests/zstreamtest.c b/tests/zstreamtest.c
index 085de8139..7dcd8ea07 100644
--- a/tests/zstreamtest.c
+++ b/tests/zstreamtest.c
@@ -436,7 +436,8 @@ static int fuzzerTests(U32 seed, U32 nbTests, unsigned startTest, double compres
 
         /* compression init */
         if (maxTestSize /* at least one test happened */ && resetAllowed && (FUZ_rand(&lseed)&1)) {
-            ZSTD_resetCStream(zc, 0);
+            U64 const pledgedSrcSize = (FUZ_rand(&lseed) & 3) ? 0 : maxTestSize;
+            ZSTD_resetCStream(zc, pledgedSrcSize);
         } else {
             U32 const testLog = FUZ_rand(&lseed) % maxSrcLog;
             U32 const cLevel = (FUZ_rand(&lseed) % (ZSTD_maxCLevel() - (testLog/3))) + 1;
@@ -449,22 +450,23 @@ static int fuzzerTests(U32 seed, U32 nbTests, unsigned startTest, double compres
             {   ZSTD_parameters params = ZSTD_getParams(cLevel, 0, dictSize);
                 params.fParams.checksumFlag = FUZ_rand(&lseed) & 1;
                 params.fParams.noDictIDFlag = FUZ_rand(&lseed) & 1;
-                {   size_t const initError = ZSTD_initCStream_advanced(zc, dict, dictSize, params, 0);
+                {   U64 const pledgedSrcSize = (FUZ_rand(&lseed) & 3) ? 0 : maxTestSize;
+                    size_t const initError = ZSTD_initCStream_advanced(zc, dict, dictSize, params, pledgedSrcSize);
                     CHECK (ZSTD_isError(initError),"ZSTD_initCStream_advanced error : %s", ZSTD_getErrorName(initError));
         }   }   }
 
         /* multi-segments compression test */
         XXH64_reset(&xxhState, 0);
-        {   U32 const maxNbChunks = (FUZ_rand(&lseed) & 127) + 2;
-            ZSTD_outBuffer outBuff = { cBuffer, cBufferSize, 0 } ;
+        {   ZSTD_outBuffer outBuff = { cBuffer, cBufferSize, 0 } ;
             U32 n;
-            for (n=0, cSize=0, totalTestSize=0 ; (n<maxNbChunks) && (totalTestSize < maxTestSize) ; n++) {
+            for (n=0, cSize=0, totalTestSize=0 ; totalTestSize < maxTestSize ; n++) {
                 /* compress random chunk into random size dst buffer */
-                {   size_t const readChunkSize = FUZ_randomLength(&lseed, maxSampleLog);
-                    size_t const srcStart = FUZ_rand(&lseed) % (srcBufferSize - readChunkSize);
+                {   size_t const randomSrcSize = FUZ_randomLength(&lseed, maxSampleLog);
+                    size_t const srcSize = MIN (maxTestSize-totalTestSize, randomSrcSize);
+                    size_t const srcStart = FUZ_rand(&lseed) % (srcBufferSize - srcSize);
                     size_t const randomDstSize = FUZ_randomLength(&lseed, maxSampleLog);
                     size_t const dstBuffSize = MIN(cBufferSize - cSize, randomDstSize);
-                    ZSTD_inBuffer inBuff = { srcBuffer+srcStart, readChunkSize, 0 };
+                    ZSTD_inBuffer inBuff = { srcBuffer+srcStart, srcSize, 0 };
                     outBuff.size = outBuff.pos + dstBuffSize;
 
                     { size_t const compressionError = ZSTD_compressStream(zc, &outBuff, &inBuff);