From: Yann Collet <cyan@fb.com>
Date: Thu, 12 Jan 2017 01:01:28 +0000 (+0100)
Subject: zstdmt : correctly check for cctx and buffer allocation
X-Git-Tag: v1.1.3^2~19^2~45
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=b05c4828eaf67fce7d2be9ef70a53591351b9ec8;p=thirdparty%2Fzstd.git

zstdmt : correctly check for cctx and buffer allocation

Result from getBuffer and getCCtx could be NULL when allocation fails.
Now correctly checks : job creation stop and last job reports an allocation error.
releaseBuffer and releaseCCtx are now also compatible with NULL input.

Identified a new potential issue :
when early job fails, later jobs are not collected for resource retrieval.
---

diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c
index c4dbb6ced..d4800dce1 100644
--- a/lib/compress/zstd_compress.c
+++ b/lib/compress/zstd_compress.c
@@ -2627,9 +2627,9 @@ size_t ZSTD_compressBegin_usingDict(ZSTD_CCtx* cctx, const void* dict, size_t di
 }
 
 
-size_t ZSTD_compressBegin(ZSTD_CCtx* zc, int compressionLevel)
+size_t ZSTD_compressBegin(ZSTD_CCtx* cctx, int compressionLevel)
 {
-    return ZSTD_compressBegin_usingDict(zc, NULL, 0, compressionLevel);
+    return ZSTD_compressBegin_usingDict(cctx, NULL, 0, compressionLevel);
 }
 
 
diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c
index 6fe37a6f4..7e8bb9f3e 100644
--- a/lib/compress/zstdmt_compress.c
+++ b/lib/compress/zstdmt_compress.c
@@ -95,6 +95,7 @@ static buffer_t ZSTDMT_getBuffer(ZSTDMT_bufferPool* pool, size_t bSize)
 /* store buffer for later re-use, up to pool capacity */
 static void ZSTDMT_releaseBuffer(ZSTDMT_bufferPool* pool, buffer_t buf)
 {
+    if (buf.start == NULL) return;   /* release on NULL */
     if (pool->nbBuffers < pool->totalBuffers) {
         pool->bTable[pool->nbBuffers++] = buf;   /* store for later re-use */
         return;
@@ -187,10 +188,10 @@ void ZSTDMT_compressChunk(void* jobDescription)
     if (ZSTD_isError(hSize)) { job->cSize = hSize; goto _endJob; }
     hSize = ZSTD_compressContinue(job->cctx, dstBuff.start, dstBuff.size, job->srcStart, 0);   /* flush frame header */
     if (ZSTD_isError(hSize)) { job->cSize = hSize; goto _endJob; }
-    if (job->firstChunk) {   /* preserve frame header when it is first chunk - otherwise, overwrite */
+    if (job->firstChunk) {   /* preserve frame header when it is first chunk */
         dstBuff.start = (char*)dstBuff.start + hSize;
         dstBuff.size -= hSize;
-    } else
+    } else                  /* otherwise, overwrite */
         hSize = 0;
 
     job->cSize = (job->lastChunk) ?   /* last chunk signal */
@@ -258,7 +259,7 @@ size_t ZSTDMT_compressCCtx(ZSTDMT_CCtx* mtctx,
     ZSTD_parameters params = ZSTD_getParams(compressionLevel, srcSize, 0);
     size_t const chunkTargetSize = (size_t)1 << (params.cParams.windowLog + 2);
     unsigned const nbChunksMax = (unsigned)(srcSize / chunkTargetSize) + (srcSize < chunkTargetSize) /* min 1 */;
-    unsigned const nbChunks = MIN(nbChunksMax, mtctx->nbThreads);
+    unsigned nbChunks = MIN(nbChunksMax, mtctx->nbThreads);
     size_t const proposedChunkSize = (srcSize + (nbChunks-1)) / nbChunks;
     size_t const avgChunkSize = ((proposedChunkSize & 0x1FFFF) < 0xFFFF) ? proposedChunkSize + 0xFFFF : proposedChunkSize;   /* avoid too small last block */
     size_t remainingSrcSize = srcSize;
@@ -274,7 +275,14 @@ size_t ZSTDMT_compressCCtx(ZSTDMT_CCtx* mtctx,
             size_t const chunkSize = MIN(remainingSrcSize, avgChunkSize);
             size_t const dstBufferCapacity = u ? ZSTD_compressBound(chunkSize) : dstCapacity;
             buffer_t const dstBuffer = u ? ZSTDMT_getBuffer(mtctx->buffPool, dstBufferCapacity) : (buffer_t){ dst, dstCapacity };
-            ZSTD_CCtx* const cctx = ZSTDMT_getCCtx(mtctx->cctxPool);   /* should check for NULL ! */
+            ZSTD_CCtx* const cctx = ZSTDMT_getCCtx(mtctx->cctxPool);
+
+            if ((cctx==NULL) || (dstBuffer.start==NULL)) {
+                mtctx->jobs[u].cSize = ERROR(memory_allocation);   /* job result */
+                mtctx->jobs[u].jobCompleted = 1;
+                nbChunks = u+1;
+                break;   /* let's wait for previous jobs to complete, but don't start new ones */
+            }
 
             mtctx->jobs[u].srcStart = srcStart + frameStartPos;
             mtctx->jobs[u].srcSize = chunkSize;
@@ -310,8 +318,8 @@ size_t ZSTDMT_compressCCtx(ZSTDMT_CCtx* mtctx,
 
             ZSTDMT_releaseCCtx(mtctx->cctxPool, mtctx->jobs[chunkID].cctx);
             {   size_t const cSize = mtctx->jobs[chunkID].cSize;
-                if (ZSTD_isError(cSize)) return cSize;
-                if (dstPos + cSize > dstCapacity) return ERROR(dstSize_tooSmall);
+                if (ZSTD_isError(cSize)) return cSize;  /* leaving here : later ressources won't be released */
+                if (dstPos + cSize > dstCapacity) return ERROR(dstSize_tooSmall);   /* leaving here : later ressources won't be released */
                 if (chunkID) {   /* note : chunk 0 is already written directly into dst */
                     memcpy((char*)dst + dstPos, mtctx->jobs[chunkID].dstBuff.start, cSize);
                     ZSTDMT_releaseBuffer(mtctx->buffPool, mtctx->jobs[chunkID].dstBuff);