From 523e0a7214a86942feabcbee8a0e0e195080b0f2 Mon Sep 17 00:00:00 2001
From: Danielle Rozenblit <drozenblit@fb.com>
Date: Thu, 12 Jan 2023 08:17:08 -0800
Subject: [PATCH] add simple test for maxBlockSize expected functionality

---
 lib/compress/zstd_compress.c |  2 +-
 lib/zstd.h                   |  2 +-
 tests/zstreamtest.c          | 43 ++++++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c
index e27cbd6133d..6872d3364b1 100644
--- a/lib/compress/zstd_compress.c
+++ b/lib/compress/zstd_compress.c
@@ -288,7 +288,7 @@ static int ZSTD_resolveExternalSequenceValidation(int mode) {
 #endif
 }
 
-/* Sets the default maximum block size. */
+/* Resolves maxBlockSize to the default if no value is present. */
 static size_t ZSTD_resolveMaxBlockSize(size_t maxBlockSize) {
     if (maxBlockSize == 0) {
         return ZSTD_BLOCKSIZE_MAX;
diff --git a/lib/zstd.h b/lib/zstd.h
index 2287886858c..0c0fc1afc91 100644
--- a/lib/zstd.h
+++ b/lib/zstd.h
@@ -141,7 +141,7 @@ ZSTDLIB_API const char* ZSTD_versionString(void);
 
 #define ZSTD_BLOCKSIZELOG_MAX  17
 #define ZSTD_BLOCKSIZE_MAX     (1<<ZSTD_BLOCKSIZELOG_MAX)
-#define ZSTD_BLOCKSIZE_MAX_MIN 1 << 10 /* The minimum valid max blocksize. Maximum blocksizes smaller than this make compressBound() inaccurate. */
+#define ZSTD_BLOCKSIZE_MAX_MIN (1 << 10) /* The minimum valid max blocksize. Maximum blocksizes smaller than this make compressBound() inaccurate. */
 
 
 /***************************************
diff --git a/tests/zstreamtest.c b/tests/zstreamtest.c
index 664ff632ac1..5a57360dc8f 100644
--- a/tests/zstreamtest.c
+++ b/tests/zstreamtest.c
@@ -1925,6 +1925,49 @@ static int basicUnitTests(U32 seed, double compressibility, int bigTests)
     }
     DISPLAYLEVEL(3, "OK \n");
 
+
+    /* Test maxBlockSize cctx param functionality */
+    DISPLAYLEVEL(3, "test%3i : Testing maxBlockSize PR#3418: ", testNb++);
+    {
+        ZSTD_CCtx* cctx = ZSTD_createCCtx();
+        size_t const srcSize = 2 << 10;
+        void* const src = CNBuffer;
+        size_t const dstSize = ZSTD_compressBound(srcSize);
+        void* const dst1 = compressedBuffer;
+        void* const dst2 = (BYTE*)compressedBuffer + dstSize;
+        size_t size1, size2;
+        void* const checkBuf = malloc(srcSize);
+
+        memset(src, 'x', srcSize);
+
+        /* Quick test to make sure maxBlockSize bounds are enforced */
+        assert(ZSTD_isError(ZSTD_CCtx_setParameter(cctx, ZSTD_c_maxBlockSize, ZSTD_BLOCKSIZE_MAX_MIN - 1)));
+        assert(ZSTD_isError(ZSTD_CCtx_setParameter(cctx, ZSTD_c_maxBlockSize, ZSTD_BLOCKSIZE_MAX + 1)));
+
+        /* maxBlockSize = 1KB */
+        CHECK_Z(ZSTD_CCtx_setParameter(cctx, ZSTD_c_maxBlockSize, 1 << 10));
+        size1 = ZSTD_compress2(cctx, dst1, dstSize, src, srcSize);
+
+        if (ZSTD_isError(size1)) goto _output_error;
+        CHECK_Z(ZSTD_decompress(checkBuf, srcSize, dst1, size1));
+        CHECK(memcmp(src, checkBuf, srcSize) != 0, "Corruption!");
+
+        /* maxBlockSize = 3KB */
+        CHECK_Z(ZSTD_CCtx_setParameter(cctx, ZSTD_c_maxBlockSize, 3 << 10));
+        size2 = ZSTD_compress2(cctx, dst2, dstSize, src, srcSize);
+
+        if (ZSTD_isError(size2)) goto _output_error;
+        CHECK_Z(ZSTD_decompress(checkBuf, srcSize, dst2, size2));
+        CHECK(memcmp(src, checkBuf, srcSize) != 0, "Corruption!");
+
+        /* Compressed output should not be equal */
+        assert(memcmp(&dst1, &dst2, dstSize) != 0);
+        assert(size1 - size2 == 4); /* We add another RLE block with header + character */
+        ZSTD_freeCCtx(cctx);
+        free(checkBuf);
+    }
+    DISPLAYLEVEL(3, "OK \n");
+
 _end:
     FUZ_freeDictionary(dictionary);
     ZSTD_freeCStream(zc);