From 92f22379344c0ad7d44cd42e56332ef924264414 Mon Sep 17 00:00:00 2001
From: Toji <155128248+tojixlugar@users.noreply.github.com>
Date: Thu, 9 Jul 2026 23:45:30 +0530
Subject: [PATCH] dibio: fix unsigned underflow in DiB_findMaxMem, guard
 nbSamples overflow

DiB_findMaxMem()'s retry loop decremented requiredMem (unsigned long long)
with no floor check. If malloc() kept failing all the way down to `step`
(8MB), the next `requiredMem -= step` would underflow and wrap around to
a value near ULLONG_MAX instead of terminating, turning what should be a
clean out-of-memory failure into an unbounded loop of futile allocation
attempts.

Fix: stop the loop once requiredMem would drop to or below `step`, and
fail explicitly via EXM_THROW if no allocation ever succeeds, instead of
silently returning an unusable value.

Also added an explicit check that fs.nbSamples is positive before it's
used to size the sampleSizes allocation in DiB_trainFromFiles(). nbSamples
is accumulated as a plain int in DiB_fileStats() and could in principle
overflow with a large enough/chunked input set; previously this relied
on malloc() happening to fail and being caught by a later check further
down. Failing fast at the point of computation is clearer and doesn't
depend on that side effect.

Both changes are no-ops on the normal/success path verified by running
`zstd --train` against a sample corpus and confirming dictionary output
is unchanged.
---
 programs/dibio.c | 36 +++++++++++++++++++++++++++++++++---
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/programs/dibio.c b/programs/dibio.c
index dc629103ba2..4794d2232a2 100644
--- a/programs/dibio.c
+++ b/programs/dibio.c
@@ -10,6 +10,8 @@
 
 
 
+
+
 /* **************************************
 *  Compiler Warnings
 ****************************************/
@@ -211,6 +213,20 @@ static void DiB_shuffle(const char** fileNamesTable, unsigned nbFiles) {
 /*-********************************************************
 *  Dictionary training functions
 **********************************************************/
+/* DiB_findMaxMem() :
+ * Tries successively smaller allocation sizes (in `step` decrements) starting
+ * from a rounded-up estimate of `requiredMem`, to discover how much memory is
+ * actually available, capped at g_maxMemory.
+ *
+ * FIX: the original loop had no floor check on `requiredMem`. If malloc() kept
+ * failing all the way down to `step`, the next `requiredMem -= step` would
+ * underflow (requiredMem is unsigned long long), wrapping around to a value
+ * near ULLONG_MAX instead of terminating. That turned a legitimate
+ * out-of-memory condition into an unbounded loop of futile allocation
+ * attempts instead of a clean, reportable failure. The loop condition now
+ * stops once requiredMem would drop to/under `step`, and if no allocation
+ * ever succeeds, it fails loudly via EXM_THROW instead of returning garbage.
+ */
 static size_t DiB_findMaxMem(unsigned long long requiredMem)
 {
     size_t const step = 8 MB;
@@ -220,11 +236,14 @@ static size_t DiB_findMaxMem(unsigned long long requiredMem)
     requiredMem += step;
     if (requiredMem > g_maxMemory) requiredMem = g_maxMemory;
 
-    while (!testmem) {
+    while (!testmem && requiredMem > step) {
         testmem = malloc((size_t)requiredMem);
-        requiredMem -= step;
+        if (!testmem) requiredMem -= step;
     }
 
+    if (!testmem)
+        EXM_THROW(31, "not enough memory for training data buffer");
+
     free(testmem);
     return (size_t)requiredMem;
 }
@@ -357,7 +376,18 @@ int DiB_trainFromFiles(const char* dictFileName, size_t maxDictSize,
             loadedSize = (size_t)MIN(loadedSize, memLimit);
         }
         srcBuffer = malloc(loadedSize+NOISELENGTH);
-        sampleSizes = (size_t*)malloc(fs.nbSamples * sizeof(size_t));
+
+        /* FIX: fs.nbSamples is accumulated as a plain int in DiB_fileStats()
+         * (summing (fileSize + chunkSize - 1) / chunkSize per file). With a
+         * sufficiently large/chunked input set this sum can in principle
+         * overflow int and go negative. The subsequent malloc() would then be
+         * fed a negative value implicitly converted to a huge size_t, which
+         * (aside from being wasteful) is undefined/surprising behavior to
+         * rely on for safety. Fail explicitly and early instead of relying on
+         * the malloc-fails-so-the-later-check-catches-it side effect. */
+        if (fs.nbSamples <= 0)
+            EXM_THROW(13, "invalid sample count (possible overflow while counting samples)");
+        sampleSizes = (size_t*)malloc((size_t)fs.nbSamples * sizeof(size_t));
     }
 
     /* Checks */