From 233c7e339217b8719678026ffa86f10e3a089934 Mon Sep 17 00:00:00 2001 From: avionicharshit-byte Date: Tue, 8 Sep 2026 23:09:29 +0530 Subject: [PATCH] fix: check pledged src size at end of multithreaded frame --- lib/compress/zstd_compress.c | 13 ++++- tests/cli-tests/compression/mt-stream-size.sh | 34 ++++++++++++ .../mt-stream-size.sh.stdout.exact | 4 ++ tests/zstreamtest.c | 52 +++++++++++++++++++ 4 files changed, 102 insertions(+), 1 deletion(-) create mode 100755 tests/cli-tests/compression/mt-stream-size.sh create mode 100644 tests/cli-tests/compression/mt-stream-size.sh.stdout.exact diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 50c3e6f7279..b051254839f 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -6534,8 +6534,19 @@ size_t ZSTD_compressStream2( ZSTD_CCtx* cctx, cctx->producedCSize += (U64)(output->pos - opos); if ( ZSTD_isError(flushMin) || (endOp == ZSTD_e_end && flushMin == 0) ) { /* compression completed */ - if (flushMin == 0) + if (flushMin == 0) { + /* end of frame : control src size, same as the single-thread path */ + if (cctx->pledgedSrcSizePlusOne != 0) { + ZSTD_STATIC_ASSERT(ZSTD_CONTENTSIZE_UNKNOWN == (unsigned long long)-1); + RETURN_ERROR_IF( + cctx->pledgedSrcSizePlusOne != cctx->consumedSrcSize+1, + srcSize_wrong, + "error : pledgedSrcSize = %u, while realSrcSize = %u", + (unsigned)cctx->pledgedSrcSizePlusOne-1, + (unsigned)cctx->consumedSrcSize); + } ZSTD_CCtx_trace(cctx, 0); + } ZSTD_CCtx_reset(cctx, ZSTD_reset_session_only); } FORWARD_IF_ERROR(flushMin, "ZSTDMT_compressStream_generic failed"); diff --git a/tests/cli-tests/compression/mt-stream-size.sh b/tests/cli-tests/compression/mt-stream-size.sh new file mode 100755 index 00000000000..5aea86035e2 --- /dev/null +++ b/tests/cli-tests/compression/mt-stream-size.sh @@ -0,0 +1,34 @@ +#!/bin/sh + +. "$COMMON/platform.sh" + +set -e + +if [ -z "$hasMT" ]; then + cat "$CLI_TESTS/compression/mt-stream-size.sh.stdout.exact" + exit 0 +fi + +# The input must be large enough to span several jobs, otherwise zstd falls +# back to compressing it with a single thread. +datagen -g2M > file2M +SIZE=2097152 + +echo "over-declared --stream-size must fail" +zstd -T2 -B512K --stream-size=$((SIZE + 1)) < file2M > $INTOVOID 2> stderr && die "should have failed" +ret=$? +[ "$ret" -eq 11 ] || die "expected exit code 11, got $ret" +grep -q "Src size is incorrect" stderr || die "expected a 'Src size is incorrect' error" + +echo "under-declared --stream-size must fail" +zstd -T2 -B512K --stream-size=$((SIZE - 1)) < file2M > $INTOVOID 2> stderr && die "should have failed" +ret=$? +[ "$ret" -eq 11 ] || die "expected exit code 11, got $ret" +grep -q "Src size is incorrect" stderr || die "expected a 'Src size is incorrect' error" + +echo "exact --stream-size must succeed and round-trip" +zstd -T2 -B512K --stream-size=$SIZE < file2M > file2M.zst +zstd -d < file2M.zst | cmp - file2M + +echo "no --stream-size must still succeed and round-trip" +zstd -T2 -B512K < file2M | zstd -d | cmp - file2M diff --git a/tests/cli-tests/compression/mt-stream-size.sh.stdout.exact b/tests/cli-tests/compression/mt-stream-size.sh.stdout.exact new file mode 100644 index 00000000000..cf0920ae1e4 --- /dev/null +++ b/tests/cli-tests/compression/mt-stream-size.sh.stdout.exact @@ -0,0 +1,4 @@ +over-declared --stream-size must fail +under-declared --stream-size must fail +exact --stream-size must succeed and round-trip +no --stream-size must still succeed and round-trip diff --git a/tests/zstreamtest.c b/tests/zstreamtest.c index 88a20a3e270..8585e9be06f 100644 --- a/tests/zstreamtest.c +++ b/tests/zstreamtest.c @@ -260,6 +260,29 @@ static U32 badParameters(ZSTD_CCtx* zc, ZSTD_parameters const savedParams) return 0; } +/* Compress @src with @pledgedSrcSize declared up-front, streaming the input so + * that a first ZSTD_e_end call doesn't auto-determine pledgedSrcSize instead. + * @return : the first error encountered, or 0, with *cSizePtr set. */ +static size_t FUZ_compressWithPledgedSrcSize(ZSTD_CCtx* cctx, + void* dst, size_t dstCapacity, + const void* src, size_t srcSize, + unsigned long long pledgedSrcSize, + size_t* cSizePtr) +{ + ZSTD_outBuffer out = { dst, dstCapacity, 0 }; + ZSTD_inBuffer in = { src, srcSize, 0 }; + size_t r = ZSTD_CCtx_reset(cctx, ZSTD_reset_session_only); + if (!ZSTD_isError(r)) r = ZSTD_CCtx_setPledgedSrcSize(cctx, pledgedSrcSize); + while ((in.pos < in.size) && !ZSTD_isError(r)) + r = ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_continue); + do { + if (ZSTD_isError(r)) break; + r = ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_end); + } while (r != 0); + *cSizePtr = out.pos; + return r; +} + static int basicUnitTests(U32 seed, double compressibility, int bigTests) { size_t const CNBufferSize = COMPRESSIBLE_NOISE_LENGTH; @@ -1493,6 +1516,35 @@ static int basicUnitTests(U32 seed, double compressibility, int bigTests) } DISPLAYLEVEL(3, "OK \n"); + /* wrong pledgedSrcSize, multithreaded : the frame must span several jobs */ + { size_t const jobSize = 1 MB; + size_t const srcSize = 4 * jobSize; + assert(srcSize <= CNBufferSize); + CHECK_Z( ZSTD_CCtx_setParameter(mtctx, ZSTD_c_jobSize, (int)jobSize) ); + + DISPLAYLEVEL(3, "test%3i : too large srcSize with multiple threads : %u bytes : ", testNb++, (unsigned)srcSize); + { size_t const r = FUZ_compressWithPledgedSrcSize(mtctx, compressedBuffer, compressedBufferSize, CNBuffer, srcSize, srcSize+1, &cSize); + if (ZSTD_getErrorCode(r) != ZSTD_error_srcSize_wrong) goto _output_error; /* must fail : wrong srcSize */ + DISPLAYLEVEL(3, "OK (error detected : %s) \n", ZSTD_getErrorName(r)); + } + + DISPLAYLEVEL(3, "test%3i : too small srcSize with multiple threads : %u bytes : ", testNb++, (unsigned)srcSize); + { size_t const r = FUZ_compressWithPledgedSrcSize(mtctx, compressedBuffer, compressedBufferSize, CNBuffer, srcSize, srcSize-1, &cSize); + if (ZSTD_getErrorCode(r) != ZSTD_error_srcSize_wrong) goto _output_error; /* must fail : wrong srcSize */ + DISPLAYLEVEL(3, "OK (error detected : %s) \n", ZSTD_getErrorName(r)); + } + + DISPLAYLEVEL(3, "test%3i : exact srcSize with multiple threads : %u bytes : ", testNb++, (unsigned)srcSize); + CHECK_Z( FUZ_compressWithPledgedSrcSize(mtctx, compressedBuffer, compressedBufferSize, CNBuffer, srcSize, srcSize, &cSize) ); + if (ZSTD_findDecompressedSize(compressedBuffer, cSize) != srcSize) goto _output_error; + CHECK_Z( ZSTD_decompress(decodedBuffer, decodedBufferSize, compressedBuffer, cSize) ); + if (memcmp(decodedBuffer, CNBuffer, srcSize) != 0) goto _output_error; + DISPLAYLEVEL(3, "OK \n"); + + CHECK_Z( ZSTD_CCtx_reset(mtctx, ZSTD_reset_session_only) ); + CHECK_Z( ZSTD_CCtx_setParameter(mtctx, ZSTD_c_jobSize, 0) ); + } + /* Complex multithreading + dictionary test */ { U32 const nbWorkers = 2; size_t const jobSize = 4 * 1 MB;