check pledged src size at end of multithreaded frame

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");

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

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

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;