From 22e27300817b672f580f98da875fefdc965af1bb Mon Sep 17 00:00:00 2001
From: inikep <inikep@gmail.com>
Date: Tue, 27 Sep 2016 18:21:17 +0200
Subject: [PATCH] ZSTD_resetDStream moved to inflate()

---
 zlibWrapper/zstd_zlibwrapper.c | 18 ++++++++++++------
 zlibWrapper/zstd_zlibwrapper.h |  3 ++-
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/zlibWrapper/zstd_zlibwrapper.c b/zlibWrapper/zstd_zlibwrapper.c
index 97f9269a4..31e784a80 100644
--- a/zlibWrapper/zstd_zlibwrapper.c
+++ b/zlibWrapper/zstd_zlibwrapper.c
@@ -480,6 +480,7 @@ ZEXTERN int ZEXPORT z_inflateInit_ OF((z_streamp strm,
                                      const char *version, int stream_size))
 {
     if (g_ZWRAPdecompressionType == ZWRAP_FORCE_ZLIB) {
+        strm->reserved = ZWRAP_ZLIB_STREAM; /* mark as zlib stream */
         return inflateInit(strm);
     }
 
@@ -530,10 +531,6 @@ int ZWRAP_inflateReset_keepDict(z_streamp strm)
 
     {   ZWRAP_DCtx* zwd = (ZWRAP_DCtx*) strm->state;
         if (zwd == NULL) return Z_STREAM_ERROR;
-        if (zwd->zbd) { 
-            size_t const errorCode = ZSTD_resetDStream(zwd->zbd);
-            if (ZSTD_isError(errorCode)) return ZWRAPD_finishWithError(zwd, strm, 0); 
-        }
         ZWRAP_initDCtx(zwd);
         zwd->decompState = ZWRAP_useReset;
     }
@@ -707,6 +704,7 @@ ZEXTERN int ZEXPORT z_inflate OF((z_streamp strm, int flush))
         if (!zwd->zbd) {
             zwd->zbd = ZSTD_createDStream_advanced(zwd->customMem);
             if (zwd->zbd == NULL) { LOG_WRAPPERD("ERROR: ZSTD_createDStream_advanced\n"); goto error; }
+            zwd->decompState = ZWRAP_useInit;
         }
 
         if (strm->total_in < ZSTD_HEADERSIZE)
@@ -715,6 +713,9 @@ ZEXTERN int ZEXPORT z_inflate OF((z_streamp strm, int flush))
                 if (zwd->decompState == ZWRAP_useInit) {
                     errorCode = ZSTD_initDStream(zwd->zbd);
                     if (ZSTD_isError(errorCode)) { LOG_WRAPPERD("ERROR: ZSTD_initDStream errorCode=%s\n", ZSTD_getErrorName(errorCode)); goto error; }
+                } else {
+                    errorCode = ZSTD_resetDStream(zwd->zbd);
+                    if (ZSTD_isError(errorCode)) goto error;
                 }
             } else {
                 srcSize = MIN(strm->avail_in, ZSTD_HEADERSIZE - strm->total_in);
@@ -724,8 +725,13 @@ ZEXTERN int ZEXPORT z_inflate OF((z_streamp strm, int flush))
                 strm->avail_in -= srcSize;
                 if (strm->total_in < ZSTD_HEADERSIZE) return Z_OK;
 
-                errorCode = ZSTD_initDStream(zwd->zbd);
-                if (ZSTD_isError(errorCode)) { LOG_WRAPPERD("ERROR: ZSTD_initDStream errorCode=%s\n", ZSTD_getErrorName(errorCode)); goto error; }
+                if (zwd->decompState == ZWRAP_useInit) {
+                    errorCode = ZSTD_initDStream(zwd->zbd);
+                    if (ZSTD_isError(errorCode)) { LOG_WRAPPERD("ERROR: ZSTD_initDStream errorCode=%s\n", ZSTD_getErrorName(errorCode)); goto error; }
+                } else {
+                    errorCode = ZSTD_resetDStream(zwd->zbd);
+                    if (ZSTD_isError(errorCode)) goto error;
+                }
 
                 zwd->inBuffer.src = zwd->headerBuf;
                 zwd->inBuffer.size = ZSTD_HEADERSIZE;
diff --git a/zlibWrapper/zstd_zlibwrapper.h b/zlibWrapper/zstd_zlibwrapper.h
index 9abbb7aa2..873413907 100644
--- a/zlibWrapper/zstd_zlibwrapper.h
+++ b/zlibWrapper/zstd_zlibwrapper.h
@@ -61,7 +61,8 @@ void ZWRAP_setDecompressionType(ZWRAP_decompress_type type);
 /* checks zstd decompression type */
 ZWRAP_decompress_type ZWRAP_getDecompressionType(void);
 
-/* checks if zstd decompression is used for a given stream */
+/* Checks if zstd decompression is used for a given stream.
+   If will return 1 only when inflate() was called and zstd header was detected. */
 int ZWRAP_isUsingZSTDdecompression(z_streamp strm);
 
 /* Similar to inflateReset but preserves dictionary set using inflateSetDictionary.
-- 
2.47.2