dictBuilder: prevent training size underflow and div-by-zero in cover/fastcover

lib/dictBuilder/cover.c

@@ -644,6 +644,13 @@ static size_t COVER_ctx_init(COVER_ctx_t *ctx, const void *samplesBuffer,
                  (unsigned)(totalSamplesSize>>20), (COVER_MAX_SAMPLES_SIZE >> 20));
     return ERROR(srcSize_wrong);
   }
+  if (trainingSamplesSize < MAX(d, sizeof(U64))) {
+    DISPLAYLEVEL(1,
+                 "Total training samples size is too small (%u bytes), "
+                 "minimum is %u bytes\n",
+                 (unsigned)trainingSamplesSize, (unsigned)MAX(d, sizeof(U64)));
+    return ERROR(srcSize_wrong);
+  }
   /* Check if there are at least 5 training samples */
   if (nbTrainSamples < 5) {
     DISPLAYLEVEL(1, "Total number of training samples is %u and is invalid.", nbTrainSamples);
@@ -736,14 +743,15 @@ COVER_epoch_info_t COVER_computeEpochs(U32 maxDictSize,
 {
   const U32 minEpochSize = k * 10;
   COVER_epoch_info_t epochs;
-  epochs.num = MAX(1, maxDictSize / k / passes);
-  epochs.size = nbDmers / epochs.num;
+  epochs.num = (k == 0 || passes == 0) ? 1 : MAX(1, maxDictSize / k / passes);
+  epochs.size = epochs.num > 0 ? nbDmers / epochs.num : 0;
   if (epochs.size >= minEpochSize) {
       assert(epochs.size * epochs.num <= nbDmers);
       return epochs;
   }
   epochs.size = MIN(minEpochSize, nbDmers);
-  epochs.num = nbDmers / epochs.size;
+  epochs.num = epochs.size > 0 ? nbDmers / epochs.size : 1;
+  epochs.num = MAX(1, epochs.num);
   assert(epochs.size * epochs.num <= nbDmers);
   return epochs;
 }

lib/dictBuilder/fastcover.c

@@ -326,6 +326,14 @@ FASTCOVER_ctx_init(FASTCOVER_ctx_t* ctx,
                     (unsigned)(totalSamplesSize >> 20), (FASTCOVER_MAX_SAMPLES_SIZE >> 20));
         return ERROR(srcSize_wrong);
     }
+    if (trainingSamplesSize < MAX(d, sizeof(U64))) {
+        DISPLAYLEVEL(1,
+                     "Total training samples size is too small (%u bytes), "
+                     "minimum is %u bytes\n",
+                     (unsigned)trainingSamplesSize,
+                     (unsigned)MAX(d, sizeof(U64)));
+        return ERROR(srcSize_wrong);
+    }
 
     /* Check if there are at least 5 training samples */
     if (nbTrainSamples < 5) {

tests/fuzzer.c

@@ -4000,6 +4000,50 @@ static int basicUnitTests(U32 const seed, double compressibility)
         free(samplesSizes);
     }
 
+    /* COVER / FASTCOVER training size validation tests */
+    {   size_t const dictSize = 16 KB;
+        void* const dictBuffer = malloc(dictSize);
+        size_t const smallSamplesSizes[8] = { 1, 1, 1, 1, 1, 1, 100, 100 };
+        ZDICT_cover_params_t coverParams;
+        ZDICT_fastCover_params_t fastCoverParams;
+        size_t res;
+
+        if (dictBuffer == NULL) goto _output_error;
+
+        DISPLAYLEVEL(3, "test%3i : COVER training size underflow : ", testNb++);
+        memset(&coverParams, 0, sizeof(coverParams));
+        coverParams.splitPoint = 0.75;
+        coverParams.d = 8;
+        coverParams.k = 16;
+        res = ZDICT_optimizeTrainFromBuffer_cover(dictBuffer, dictSize,
+                                                 CNBuffer, smallSamplesSizes, 8,
+                                                 &coverParams);
+        if (!ZDICT_isError(res) ||
+            ZSTD_getErrorCode(res) != ZSTD_error_srcSize_wrong) {
+            free(dictBuffer);
+            goto _output_error;
+        }
+        DISPLAYLEVEL(3, "OK \n");
+
+        DISPLAYLEVEL(3, "test%3i : FASTCOVER training underflow : ", testNb++);
+        memset(&fastCoverParams, 0, sizeof(fastCoverParams));
+        fastCoverParams.d = 8;
+        fastCoverParams.k = 16;
+        fastCoverParams.f = 14;
+        res = ZDICT_optimizeTrainFromBuffer_fastCover(dictBuffer, dictSize,
+                                                     CNBuffer,
+                                                     smallSamplesSizes, 8,
+                                                     &fastCoverParams);
+        if (!ZDICT_isError(res) ||
+            ZSTD_getErrorCode(res) != ZSTD_error_srcSize_wrong) {
+            free(dictBuffer);
+            goto _output_error;
+        }
+        DISPLAYLEVEL(3, "OK \n");
+
+        free(dictBuffer);
+    }
+
     /* Decompression defense tests */
     DISPLAYLEVEL(3, "test%3i : Check input length for magic number : ", testNb++);
     { size_t const r = ZSTD_decompress(decodedBuffer, CNBuffSize, CNBuffer, 3);   /* too small input */