zdict: check allocation sizes in ZDICT_trainBuffer_legacy()

lib/dictBuilder/zdict.c

@@ -471,11 +471,11 @@ static size_t ZDICT_trainBuffer_legacy(dictItem* dictList, U32 dictListSize,
                             const size_t* fileSizes, unsigned nbFiles,
                             unsigned minRatio, U32 notificationLevel)
 {
-    unsigned* const suffix0 = (unsigned*)malloc((bufferSize+2)*sizeof(*suffix0));
-    unsigned* const suffix = suffix0+1;
-    U32* reverseSuffix = (U32*)malloc((bufferSize)*sizeof(*reverseSuffix));
-    BYTE* doneMarks = (BYTE*)malloc((bufferSize+16)*sizeof(*doneMarks));   /* +16 for overflow security */
-    U32* filePos = (U32*)malloc(nbFiles * sizeof(*filePos));
+    unsigned* suffix0;
+    unsigned* suffix;
+    U32* reverseSuffix;
+    BYTE* doneMarks;
+    U32* filePos;
     size_t result = 0;
     clock_t displayClock = 0;
     clock_t const refreshRate = CLOCKS_PER_SEC * 3 / 10;
@@ -494,10 +494,37 @@ static size_t ZDICT_trainBuffer_legacy(dictItem* dictList, U32 dictListSize,
 
     /* init */
     DISPLAYLEVEL(2, "\r%70s\r", "");   /* clean display line */
+
+    /* Each working buffer is sized from bufferSize or nbFiles. Reject sizes
+     * whose byte counts cannot be represented instead of under-allocating:
+     * on 32-bit targets, (bufferSize+2)*sizeof(*suffix0) and
+     * bufferSize*sizeof(*reverseSuffix) wrap for bufferSize >= 0x3FFFFFFE
+     * (~1 GB of samples), so the suffix array under-allocates and
+     * divsufsort() writes out of bounds; nbFiles*sizeof(*filePos) wraps the
+     * same way once nbFiles > SIZE_MAX / sizeof(*filePos).
+     */
+    {   size_t suffix0Bytes, reverseSuffixBytes, doneMarksBytes, filePosBytes;
+        if (bufferSize > (size_t)(-1) - 2) return ERROR(memory_allocation);
+        suffix0Bytes = (bufferSize + 2) * sizeof(*suffix0);
+        if (suffix0Bytes / sizeof(*suffix0) != bufferSize + 2) return ERROR(memory_allocation);
+        reverseSuffixBytes = bufferSize * sizeof(*reverseSuffix);
+        if (reverseSuffixBytes / sizeof(*reverseSuffix) != bufferSize) return ERROR(memory_allocation);
+        if (bufferSize > (size_t)(-1) - 16) return ERROR(memory_allocation);
+        doneMarksBytes = (bufferSize + 16) * sizeof(*doneMarks);
+        if (doneMarksBytes / sizeof(*doneMarks) != bufferSize + 16) return ERROR(memory_allocation);
+        filePosBytes = (size_t)nbFiles * sizeof(*filePos);
+        if (filePosBytes / sizeof(*filePos) != (size_t)nbFiles) return ERROR(memory_allocation);
+
+        suffix0 = (unsigned*)malloc(suffix0Bytes);
+        reverseSuffix = (U32*)malloc(reverseSuffixBytes);
+        doneMarks = (BYTE*)malloc(doneMarksBytes);   /* +16 for overflow security */
+        filePos = (U32*)malloc(filePosBytes);
+    }
     if (!suffix0 || !reverseSuffix || !doneMarks || !filePos) {
         result = ERROR(memory_allocation);
         goto _cleanup;
     }
+    suffix = suffix0 + 1;
     if (minRatio < MINRATIO) minRatio = MINRATIO;
     memset(doneMarks, 0, bufferSize+16);
 

tests/fuzzer.c

@@ -26,6 +26,10 @@
 #include <stdio.h>        /* fgets, sscanf */
 #include <string.h>       /* strcmp */
 #include <time.h>         /* time(), time_t */
+#ifndef _WIN32
+#  include <unistd.h>     /* fork, _exit */
+#  include <sys/wait.h>   /* waitpid */
+#endif
 #undef NDEBUG             /* always enable assert() */
 #include <assert.h>
 #define ZSTD_STATIC_LINKING_ONLY  /* ZSTD_compressContinue, ZSTD_compressBlock */
@@ -3995,6 +3999,60 @@ static int basicUnitTests(U32 const seed, double compressibility)
         if (dictID==0) goto _output_error;
         DISPLAYLEVEL(3, "OK : %u \n", (unsigned)dictID);
 
+        DISPLAYLEVEL(3, "test%3i : ZDICT_trainFromBuffer_legacy : ", testNb++);
+        {   ZDICT_legacy_params_t legacyParams;
+            size_t const legacyCapacity = 16 KB;
+            memset(&legacyParams, 0, sizeof(legacyParams));
+            dictSize = ZDICT_trainFromBuffer_legacy(dictBuffer, legacyCapacity,
+                                                   CNBuffer, samplesSizes, nbSamples,
+                                                   legacyParams);
+            if (ZDICT_isError(dictSize)) goto _output_error;
+        }
+        DISPLAYLEVEL(3, "OK, created dictionary of size %u \n", (unsigned)dictSize);
+
+        DISPLAYLEVEL(3, "test%3i : ZDICT_trainFromBuffer_legacy rejects wrapping allocation sizes : ", testNb++);
+        {   /* (bufferSize+2)*sizeof(unsigned) wraps on 32-bit targets for
+             * bufferSize >= 0x3FFFFFFE : the suffix array under-allocates and
+             * divsufsort() writes out of bounds. The call must fail cleanly
+             * instead. Only reachable on 32-bit : on 64-bit the products are
+             * representable and a huge allocation just fails gracefully. */
+            if (!MEM_32bits()) {
+                DISPLAYLEVEL(3, "skipped (allocation sizes do not wrap on 64-bit) \n");
+            } else {
+                size_t const wrapSize = (size_t)0x40000000;   /* 1 GB */
+                void* const wrapBuffer = malloc(wrapSize);   /* untouched : reads as zeros */
+                size_t wrapSampleSize = wrapSize;
+                if (wrapBuffer == NULL) {
+                    DISPLAYLEVEL(3, "skipped (cannot allocate witness buffer) \n");
+                } else {
+#ifdef _WIN32
+                    (void)wrapSampleSize;
+                    DISPLAYLEVEL(3, "skipped (no fork on Windows) \n");
+#else
+                    {   /* crash-isolated : an unfixed tree dies inside this call */
+                        pid_t const child = fork();
+                        assert(child >= 0);
+                        if (child == 0) {
+                            ZDICT_legacy_params_t legacyParams;
+                            size_t wrapRC;
+                            memset(&legacyParams, 0, sizeof(legacyParams));
+                            wrapRC = ZDICT_trainFromBuffer_legacy(
+                                dictBuffer, 16 KB, wrapBuffer, &wrapSampleSize, 1,
+                                legacyParams);
+                            fflush(stdout);
+                            _exit(ZDICT_isError(wrapRC) ? 0 : 1);
+                        }
+                        {   int status = 0;
+                            if (waitpid(child, &status, 0) != child) goto _output_error;
+                            if (!WIFEXITED(status) || (WEXITSTATUS(status) != 0)) goto _output_error;
+                    }   }
+#endif
+                    free(wrapBuffer);
+                }
+            }
+        }
+        DISPLAYLEVEL(3, "OK \n");
+
         ZSTD_freeCCtx(cctx);
         free(dictBuffer);
         free(samplesSizes);