diff --git a/lib/common/pool.c b/lib/common/pool.c index dd5fb0c4d4b..0177e3ccb22 100644 --- a/lib/common/pool.c +++ b/lib/common/pool.c @@ -117,7 +117,12 @@ POOL_ctx* POOL_create_advanced(size_t numThreads, size_t queueSize, { POOL_ctx* ctx; /* Check parameters */ - if (!numThreads) { return NULL; } + if (!numThreads || numThreads > POOL_MAX_THREADS) { + return NULL; + } + if (queueSize >= (size_t)-1 / sizeof(POOL_job)) { + return NULL; + } /* Allocate the context and zero initialize */ ctx = (POOL_ctx*)ZSTD_customCalloc(sizeof(POOL_ctx), customMem); if (!ctx) { return NULL; } @@ -190,6 +195,7 @@ void POOL_free(POOL_ctx *ctx) { * Waits for all queued jobs to finish executing. */ void POOL_joinJobs(POOL_ctx* ctx) { + if (!ctx) { return; } ZSTD_pthread_mutex_lock(&ctx->queueMutex); while(!ctx->queueEmpty || ctx->numThreadsBusy > 0) { ZSTD_pthread_cond_wait(&ctx->queuePushCond, &ctx->queueMutex); @@ -218,6 +224,7 @@ static int POOL_resize_internal(POOL_ctx* ctx, size_t numThreads) return 0; } /* numThreads > threadCapacity */ + if (numThreads > POOL_MAX_THREADS) return 1; ctx->threadLimit = numThreads; { ZSTD_pthread_t* const threadPool = (ZSTD_pthread_t*)ZSTD_customCalloc(numThreads * sizeof(ZSTD_pthread_t), ctx->customMem); if (!threadPool) return 1; @@ -283,6 +290,7 @@ POOL_add_internal(POOL_ctx* ctx, POOL_function function, void *opaque) void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque) { assert(ctx != NULL); + assert(function != NULL); ZSTD_pthread_mutex_lock(&ctx->queueMutex); /* Wait until there is space in the queue for the new job */ while (isQueueFull(ctx) && (!ctx->shutdown)) { @@ -296,6 +304,7 @@ void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque) int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque) { assert(ctx != NULL); + assert(function != NULL); ZSTD_pthread_mutex_lock(&ctx->queueMutex); if (isQueueFull(ctx)) { ZSTD_pthread_mutex_unlock(&ctx->queueMutex); @@ -350,11 +359,13 @@ int POOL_resize(POOL_ctx* ctx, size_t numThreads) { void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque) { (void)ctx; + assert(function != NULL); function(opaque); } int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque) { (void)ctx; + assert(function != NULL); function(opaque); return 1; } diff --git a/lib/common/pool.h b/lib/common/pool.h index f39b7f1eb99..d78da3e84d7 100644 --- a/lib/common/pool.h +++ b/lib/common/pool.h @@ -18,6 +18,8 @@ typedef struct POOL_ctx_s POOL_ctx; +#define POOL_MAX_THREADS 10000 + /*! POOL_create() : * Create a thread pool with at most `numThreads` threads. * `numThreads` must be at least 1. diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index cb710475c2a..7b9a9be553e 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -143,10 +143,14 @@ static ZSTDMT_bufferPool* ZSTDMT_createBufferPool(unsigned maxNbBuffers, ZSTD_cu /* only works at initialization, not during compression */ static size_t ZSTDMT_sizeof_bufferPool(ZSTDMT_bufferPool* bufPool) { - size_t const poolSize = sizeof(*bufPool); - size_t const arraySize = bufPool->totalBuffers * sizeof(Buffer); + size_t poolSize; + size_t arraySize; unsigned u; size_t totalBufferSize = 0; + + if (bufPool == NULL) return 0; + poolSize = sizeof(*bufPool); + arraySize = bufPool->totalBuffers * sizeof(Buffer); ZSTD_pthread_mutex_lock(&bufPool->poolMutex); for (u=0; utotalBuffers; u++) totalBufferSize += bufPool->buffers[u].capacity; @@ -161,6 +165,7 @@ static size_t ZSTDMT_sizeof_bufferPool(ZSTDMT_bufferPool* bufPool) * as they become freely interchangeable, reducing malloc/free usages and memory fragmentation */ static void ZSTDMT_setBufferSize(ZSTDMT_bufferPool* const bufPool, size_t const bSize) { + if (bufPool == NULL) return; ZSTD_pthread_mutex_lock(&bufPool->poolMutex); DEBUGLOG(4, "ZSTDMT_setBufferSize: bSize = %u", (U32)bSize); bufPool->bufferSize = bSize; @@ -417,6 +422,7 @@ static ZSTDMT_CCtxPool* ZSTDMT_expandCCtxPool(ZSTDMT_CCtxPool* srcPool, /* only works during initialization phase, not during compression */ static size_t ZSTDMT_sizeof_CCtxPool(ZSTDMT_CCtxPool* cctxPool) { + if (cctxPool == NULL) return 0; ZSTD_pthread_mutex_lock(&cctxPool->poolMutex); { unsigned const nbWorkers = cctxPool->totalCCtx; size_t const poolSize = sizeof(*cctxPool); @@ -434,6 +440,7 @@ static size_t ZSTDMT_sizeof_CCtxPool(ZSTDMT_CCtxPool* cctxPool) static ZSTD_CCtx* ZSTDMT_getCCtx(ZSTDMT_CCtxPool* cctxPool) { + if (cctxPool == NULL) return NULL; DEBUGLOG(5, "ZSTDMT_getCCtx"); ZSTD_pthread_mutex_lock(&cctxPool->poolMutex); if (cctxPool->availCCtx) { @@ -449,7 +456,7 @@ static ZSTD_CCtx* ZSTDMT_getCCtx(ZSTDMT_CCtxPool* cctxPool) static void ZSTDMT_releaseCCtx(ZSTDMT_CCtxPool* pool, ZSTD_CCtx* cctx) { - if (cctx==NULL) return; /* compatibility with release on NULL */ + if (cctx==NULL || pool==NULL) return; ZSTD_pthread_mutex_lock(&pool->poolMutex); if (pool->availCCtx < pool->totalCCtx) pool->cctxs[pool->availCCtx++] = cctx; diff --git a/tests/poolTests.c b/tests/poolTests.c index 9e62722bfa0..50dc2de6591 100644 --- a/tests/poolTests.c +++ b/tests/poolTests.c @@ -233,6 +233,42 @@ int main(int argc, const char **argv) { return 1; } + if (POOL_create(1, (size_t)-1)) { + printf("FAIL: should not create POOL with overflowing queueSize\n"); + return 1; + } + + if (POOL_create((size_t)-1, 1)) { + printf("FAIL: should not create POOL with overflowing numThreads\n"); + return 1; + } + + if (POOL_create(POOL_MAX_THREADS + 1, 1)) { + printf("FAIL: should not create POOL exceeding POOL_MAX_THREADS\n"); + return 1; + } + + POOL_joinJobs(NULL); + + { + POOL_ctx* const testCtx = POOL_create(1, 1); + if (!testCtx) { + printf("FAIL: could not create POOL for resize test\n"); + return 1; + } + if (POOL_resize(testCtx, (size_t)-1) == 0) { + printf("FAIL: POOL_resize with overflowing numThreads should fail\n"); + POOL_free(testCtx); + return 1; + } + if (POOL_resize(testCtx, POOL_MAX_THREADS + 1) == 0) { + printf("FAIL: POOL_resize exceeding POOL_MAX_THREADS should fail\n"); + POOL_free(testCtx); + return 1; + } + POOL_free(testCtx); + } + for (numThreads = 1; numThreads <= 4; ++numThreads) { size_t queueSize; for (queueSize = 0; queueSize <= 2; ++queueSize) {