Repository: facebook/zstd

Problem statement:
ZSTD_compressStream2  can't stop with wrong ZSTD_EndDirective argument
`ZSTD_EndDirective` enum's value is 0 (continue), 1 (flush), 2 (end).
If pass `-1` or `3` to `ZSTD_compressStream2()` function, the function can't stop.

Seems an infinite loop in `ZSTD_compressStream_generic()`:
https://github.com/facebook/zstd/blob/dev/lib/compress/zstd_compress.c#L3808.

Test code:
```c
#include <stdio.h>     // printf
#include <zstd.h>      // presumes zstd library is installed

void compress()
{
    char buffer[256];
    ZSTD_inBuffer in;
    ZSTD_outBuffer out;
    ZSTD_CCtx *cctx;
    
    in.src = &in;
    in.size = 0;
    in.pos = 0;
    
    out.dst = buffer;
    out.size = sizeof(buffer);
    out.pos = 0;
    
    cctx = ZSTD_createCCtx();
    
    printf("ZSTD_compressStream2 started.\n");
    ZSTD_compressStream2(cctx, &out, &in, 3);  // can't stop other than 0,1,2
    printf("ZSTD_compressStream2 finished.\n");
}

int main(int argc, const char** argv)
{
    compress();
    return 0;
}
```


Hints:
Passing a value other than a valid `ZSTD_EndDirective` is undefined behavior. It would be nice to report an error instead of infinite looping, but it isn't a requirement of the API.
> 
> 
> Passing a value other than a valid `ZSTD_EndDirective` is undefined behavior. It would be nice to report an error instead of infinite looping, but it isn't a requirement of the API.

Yes, downstream code can easily handle this problem.
Considering any value beyond the valid range of `ZSTD_EndDirective` as UB (Undefined Behavior) on reaching `ZSTD_compressStream_generic()` is the correct thing to do to, since it's an internal function, and everything should be clean on reaching this point.

In contrast, `ZSTD_compressStream2()` is a public function, and while it's reasonable to require `endDirective` to be valid, it's also a place where we have no direct control; the user is entirely in charge, and could mess up.

A similar scenario is providing an input where `in.pos > in.size`, which is obviously invalid. 
We could also invoke `UB`, but have also decided to [actively check that the structure is valid](https://github.com/facebook/zstd/blob/v1.4.5/lib/compress/zstd_compress.c#L3970), and return an error code when it's not.

This is in part because this check is trivial and costs (proportionally) almost nothing. 
Checking the validity of `endDirective` seems almost as simple. Therefore, we could add an equivalent check.

As a general principle, relying solely on `UB`, hence transferring the full responsibility of respecting interface's requirements to the user, is the right thing to do when checking these requirements is too difficult or too costly. When when such a cost is minimal, it's generally better to reduce `UB` territory.


It would be nice to not add any overhead.
Since the user of zstd are almost developer, maybe can verify the parameter with `assert()`.

Basically, it is not easy to make mistakes when using `endDirective` argument.

Patch:
diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c
index eedd9b67f5c..8c8c33a859d 100644
--- a/lib/compress/zstd_compress.c
+++ b/lib/compress/zstd_compress.c
@@ -3937,6 +3937,7 @@ static size_t ZSTD_compressStream_generic(ZSTD_CStream* zcs,
     assert(zcs->outBuffSize > 0);
     assert(output->pos <= output->size);
     assert(input->pos <= input->size);
+    assert((U32)flushMode <= (U32)ZSTD_e_end);
 
     while (someMoreWork) {
         switch(zcs->streamStage)
@@ -4083,9 +4084,10 @@ size_t ZSTD_compressStream2( ZSTD_CCtx* cctx,
 {
     DEBUGLOG(5, "ZSTD_compressStream2, endOp=%u ", (unsigned)endOp);
     /* check conditions */
-    RETURN_ERROR_IF(output->pos > output->size, GENERIC, "invalid buffer");
-    RETURN_ERROR_IF(input->pos  > input->size, GENERIC, "invalid buffer");
-    assert(cctx!=NULL);
+    RETURN_ERROR_IF(output->pos > output->size, dstSize_tooSmall, "invalid output buffer");
+    RETURN_ERROR_IF(input->pos  > input->size, srcSize_wrong, "invalid input buffer");
+    RETURN_ERROR_IF((U32)endOp > (U32)ZSTD_e_end, parameter_outOfBound, "invalid endDirective");
+    assert(cctx != NULL);
 
     /* transparent initialization stage */
     if (cctx->streamStage == zcss_init) {
diff --git a/lib/zstd.h b/lib/zstd.h
index f8d5e84da1c..55246462711 100644
--- a/lib/zstd.h
+++ b/lib/zstd.h
@@ -677,8 +677,9 @@ typedef enum {
  *  - Compression parameters cannot be changed once compression is started (save a list of exceptions in multi-threading mode)
  *  - output->pos must be <= dstCapacity, input->pos must be <= srcSize
  *  - output->pos and input->pos will be updated. They are guaranteed to remain below their respective limit.
+ *  - endOp must be a valid directive
  *  - When nbWorkers==0 (default), function is blocking : it completes its job before returning to caller.
- *  - When nbWorkers>=1, function is non-blocking : it just acquires a copy of input, and distributes jobs to internal worker threads, flush whatever is available,
+ *  - When nbWorkers>=1, function is non-blocking : it copies a portion of input, distributes jobs to internal worker threads, flush to output whatever is available,
  *                                                  and then immediately returns, just indicating that there is some data remaining to be flushed.
  *                                                  The function nonetheless guarantees forward progress : it will return only after it reads or write at least 1+ byte.
  *  - Exception : if the first call requests a ZSTD_e_end directive and provides enough dstCapacity, the function delegates to ZSTD_compress2() which is always blocking.


Test patch:
diff --git a/tests/fuzzer.c b/tests/fuzzer.c
index 52086c7f389..8929c8fb396 100644
--- a/tests/fuzzer.c
+++ b/tests/fuzzer.c
@@ -439,6 +439,17 @@ static int basicUnitTests(U32 const seed, double compressibility)
     }   }
     DISPLAYLEVEL(3, "OK \n");
 
+    DISPLAYLEVEL(3, "test%3u : invalid endDirective : ", testNb++);
+    {   ZSTD_CCtx* const cctx = ZSTD_createCCtx();
+        ZSTD_inBuffer inb = { CNBuffer, CNBuffSize, 0 };
+        ZSTD_outBuffer outb = { compressedBuffer, compressedBufferSize, 0 };
+        if (cctx==NULL) goto _output_error;
+        CHECK( ZSTD_isError( ZSTD_compressStream2(cctx, &outb, &inb, (ZSTD_EndDirective) 3) ) );  /* must fail */
+        CHECK( ZSTD_isError( ZSTD_compressStream2(cctx, &outb, &inb, (ZSTD_EndDirective)-1) ) );  /* must fail */
+        ZSTD_freeCCtx(cctx);
+    }
+    DISPLAYLEVEL(3, "OK \n");
+
     DISPLAYLEVEL(3, "test%3i : ZSTD_checkCParams : ", testNb++);
     {
         ZSTD_parameters params = ZSTD_getParams(3, 0, 0);
@@ -866,36 +877,35 @@ static int basicUnitTests(U32 const seed, double compressibility)
         assert(cctx != NULL);
         ZSTD_CCtx_setParameter(cctx, ZSTD_c_targetCBlockSize, 81);
 
-        { size_t read;
-          for (read = 0; read < streamCompressThreshold; read += streamCompressDelta) {
-            ZSTD_inBuffer in = {src, streamCompressDelta, 0};
-            ZSTD_outBuffer out = {dst, dstCapacity, 0};
-            CHECK_Z(ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_continue));
-            CHECK_Z(ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_end));
-            src += streamCompressDelta; srcSize -= streamCompressDelta;
-            dst += out.pos; dstCapacity -= out.pos;}}
+        {   size_t read;
+            for (read = 0; read < streamCompressThreshold; read += streamCompressDelta) {
+                ZSTD_inBuffer in = {src, streamCompressDelta, 0};
+                ZSTD_outBuffer out = {dst, dstCapacity, 0};
+                CHECK_Z(ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_continue));
+                CHECK_Z(ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_end));
+                src += streamCompressDelta; srcSize -= streamCompressDelta;
+                dst += out.pos; dstCapacity -= out.pos;
+        }   }
 
         /* This is trying to catch a dstSize_tooSmall error */
 
-        { ZSTD_inBuffer in = {src, srcSize, 0};
-          ZSTD_outBuffer out = {dst, dstCapacity, 0};
-          CHECK_Z(ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_end));}
+        {   ZSTD_inBuffer in = {src, srcSize, 0};
+            ZSTD_outBuffer out = {dst, dstCapacity, 0};
+            CHECK_Z(ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_end));
+        }
         ZSTD_freeCCtx(cctx);
     }
     DISPLAYLEVEL(3, "OK \n");
 
     DISPLAYLEVEL(3, "test%3d: superblock with no literals : ", testNb++);
     /* Generate the same data 20 times over */
-    {
-        size_t const avgChunkSize = CNBuffSize / 20;
+    {   size_t const avgChunkSize = CNBuffSize / 20;
         size_t b;
         for (b = 0; b < CNBuffSize; b += avgChunkSize) {
             size_t const chunkSize = MIN(CNBuffSize - b, avgChunkSize);
             RDG_genBuffer((char*)CNBuffer + b, chunkSize, compressibility, 0. /* auto */, seed);
-        }
-    }
-    {
-        ZSTD_CCtx* const cctx = ZSTD_createCCtx();
+    }   }
+    {   ZSTD_CCtx* const cctx = ZSTD_createCCtx();
         size_t const normalCSize = ZSTD_compress2(cctx, compressedBuffer, compressedBufferSize, CNBuffer, CNBuffSize);
         size_t const allowedExpansion = (CNBuffSize * 3 / 1000);
         size_t superCSize;
@@ -914,12 +924,11 @@ static int basicUnitTests(U32 const seed, double compressibility)
 
     RDG_genBuffer(CNBuffer, CNBuffSize, compressibility, 0. /*auto*/, seed);
     DISPLAYLEVEL(3, "test%3d: superblock enough room for checksum : ", testNb++)
-    {
-        /* This tests whether or not we leave enough room for the checksum at the end
-         * of the dst buffer. The bug that motivated this test was found by the
-         * stream_round_trip fuzzer but this crashes for the same reason and is
-         * far more compact than re-creating the stream_round_trip fuzzer's code path */
-        ZSTD_CCtx *cctx = ZSTD_createCCtx();
+    /* This tests whether or not we leave enough room for the checksum at the end
+     * of the dst buffer. The bug that motivated this test was found by the
+     * stream_round_trip fuzzer but this crashes for the same reason and is
+     * far more compact than re-creating the stream_round_trip fuzzer's code path */
+    {   ZSTD_CCtx* const cctx = ZSTD_createCCtx();
         ZSTD_CCtx_setParameter(cctx, ZSTD_c_targetCBlockSize, 64);
         assert(!ZSTD_isError(ZSTD_compress2(cctx, compressedBuffer, 1339, CNBuffer, 1278)));
         ZSTD_freeCCtx(cctx);
@@ -2788,10 +2797,8 @@ static int basicUnitTests(U32 const seed, double compressibility)
     DISPLAYLEVEL(3, "OK \n");
 #ifdef ZSTD_MULTITHREAD
     DISPLAYLEVEL(3, "test%3i : passing wrong full dict should fail on compressStream2 refPrefix ", testNb++);
-    {
-        ZSTD_CCtx* cctx = ZSTD_createCCtx();
-         /* A little more than ZSTDMT_JOBSIZE_MIN */
-        size_t const srcSize = 1 MB + 5;
+    {   ZSTD_CCtx* cctx = ZSTD_createCCtx();
+        size_t const srcSize = 1 MB + 5;   /* A little more than ZSTDMT_JOBSIZE_MIN */
         size_t const dstSize = ZSTD_compressBound(srcSize);
         void* const src = CNBuffer;
         void* const dst = compressedBuffer;
@@ -2808,10 +2815,8 @@ static int basicUnitTests(U32 const seed, double compressibility)
         /* lie and claim this is a full dict */
         CHECK_Z(ZSTD_CCtx_refPrefix_advanced(cctx, dict, srcSize, ZSTD_dct_fullDict));
 
-        {
-            ZSTD_outBuffer out = {dst, dstSize, 0};
+        {   ZSTD_outBuffer out = {dst, dstSize, 0};
             ZSTD_inBuffer in = {src, srcSize, 0};
-
             /* should fail because its not a full dict like we said it was */
             assert(ZSTD_isError(ZSTD_compressStream2(cctx, &out, &in, ZSTD_e_flush)));
         }
@@ -2822,10 +2827,8 @@ static int basicUnitTests(U32 const seed, double compressibility)
     DISPLAYLEVEL(3, "OK \n");
 
     DISPLAYLEVEL(3, "test%3i : small dictionary with multithreading and LDM ", testNb++);
-    {
-        ZSTD_CCtx* cctx = ZSTD_createCCtx();
-        /* A little more than ZSTDMT_JOBSIZE_MIN */
-        size_t const srcSize = 1 MB + 5;
+    {   ZSTD_CCtx* cctx = ZSTD_createCCtx();
+        size_t const srcSize = 1 MB + 5;   /* A little more than ZSTDMT_JOBSIZE_MIN */
         size_t const dictSize = 10;
         size_t const dstSize = ZSTD_compressBound(srcSize);
         void* const src = CNBuffer;
@@ -2853,8 +2856,7 @@ static int basicUnitTests(U32 const seed, double compressibility)
 
     /* note : this test is rather long, it would be great to find a way to speed up its execution */
     DISPLAYLEVEL(3, "test%3i : table cleanliness through index reduction : ", testNb++);
-    {
-        int cLevel;
+    {   int cLevel;
         size_t approxIndex = 0;
         size_t maxIndex = ((3U << 29) + (1U << ZSTD_WINDOWLOG_MAX)); /* ZSTD_CURRENT_MAX from zstd_compress_internal.h */
 
@@ -2923,8 +2925,7 @@ static int basicUnitTests(U32 const seed, double compressibility)
     DISPLAYLEVEL(3, "OK \n");
 
     DISPLAYLEVEL(3, "test%3i : testing cdict compression with different attachment strategies : ", testNb++);
-    {
-        ZSTD_CCtx* const cctx = ZSTD_createCCtx();
+    {   ZSTD_CCtx* const cctx = ZSTD_createCCtx();
         ZSTD_DCtx* const dctx = ZSTD_createDCtx();
         size_t dictSize = CNBuffSize;
         void* dict = (void*)malloc(dictSize);
@@ -2975,9 +2976,7 @@ static int basicUnitTests(U32 const seed, double compressibility)
 
                     CHECK_Z(ZSTD_CCtx_reset(cctx, ZSTD_reset_session_and_parameters));
                     ZSTD_freeCDict(cdict);
-                }
-            }
-        }
+        }   }   }
 
         ZSTD_freeCCtx(cctx);
         ZSTD_freeDCtx(dctx);