Check size arithmetic where it is used in the quantizers

Tests/test_image_quantize.py

@@ -135,6 +135,46 @@ def test_colors() -> None:
     assert len(converted.palette.palette) == colors * len("RGB")
 
 
+@pytest.mark.parametrize("size", ((1, 1), (1, 300), (300, 1), (64, 64)))
+@pytest.mark.parametrize("colors", (1, 2, 256))
+@pytest.mark.parametrize(
+    "mode, method",
+    (
+        ("L", Image.Quantize.MEDIANCUT),
+        ("P", Image.Quantize.MEDIANCUT),
+        ("RGB", Image.Quantize.MEDIANCUT),
+        ("RGB", Image.Quantize.MAXCOVERAGE),
+        ("RGB", Image.Quantize.FASTOCTREE),
+        ("RGBA", Image.Quantize.FASTOCTREE),
+    ),
+)
+def test_quantize_sizes(
+    mode: str, method: Image.Quantize, colors: int, size: tuple[int, int]
+) -> None:
+    im = hopper(mode).resize(size)
+
+    converted = im.quantize(colors, method)
+    assert converted.mode == "P"
+    assert converted.size == size
+    assert converted.palette is not None
+    entries = len(converted.palette.palette) // len(converted.palette.mode)
+    assert 1 <= entries <= colors
+    assert max(converted.tobytes()) < entries
+
+
+@skip_unless_feature("libimagequant")
+@pytest.mark.parametrize("size", ((1, 1), (1, 300), (300, 1)))
+def test_libimagequant_quantize_sizes(size: tuple[int, int]) -> None:
+    im = hopper().resize(size)
+
+    converted = im.quantize(2, Image.Quantize.LIBIMAGEQUANT)
+    assert converted.mode == "P"
+    assert converted.size == size
+    colors = converted.getcolors()
+    assert colors is not None
+    assert len(colors) <= 2
+
+
 def test_transparent_colors_equal() -> None:
     im = Image.new("RGBA", (1, 2), (0, 0, 0, 0))
     px = im.load()

checks/check_large_memory.py

@@ -55,3 +55,9 @@ def test_size_greater_than_int() -> None:
     assert numpy is not None
     arr = numpy.ndarray(shape=(16394, 16394))
     Image.fromarray(arr)
+
+
+def test_quantize_more_pixels_than_int() -> None:
+    im = Image.new("L", (2**16, 2**15))
+    with pytest.raises(MemoryError):
+        im.quantize()

src/libImaging/CheckedArithmetic.h

@@ -0,0 +1,61 @@
+/*
+ * The Python Imaging Library
+ *
+ * checked arithmetic for sizes
+ *
+ * See the README file for information on usage and redistribution.
+ */
+
+#ifndef __CHECKED_ARITHMETIC_H__
+#define __CHECKED_ARITHMETIC_H__
+
+#include <stddef.h>
+#include <stdint.h>
+
+/* Helpers to compute sizes, such as the number of bytes to allocate, without
+   overflowing. Each one stores the result in *result and returns 1, or returns
+   0 and leaves *result unchanged if the result does not fit in a size_t. */
+
+#if defined(__has_builtin)
+#if __has_builtin(__builtin_add_overflow) && __has_builtin(__builtin_mul_overflow)
+#define IMAGING_HAVE_OVERFLOW_BUILTINS
+#endif
+#elif defined(__GNUC__) && __GNUC__ >= 5
+#define IMAGING_HAVE_OVERFLOW_BUILTINS
+#endif
+
+static inline int
+ImagingCheckedAdd(size_t a, size_t b, size_t *result) {
+#ifdef IMAGING_HAVE_OVERFLOW_BUILTINS
+    size_t r;
+    if (__builtin_add_overflow(a, b, &r)) {
+        return 0;
+    }
+    *result = r;
+#else
+    if (a > SIZE_MAX - b) {
+        return 0;
+    }
+    *result = a + b;
+#endif
+    return 1;
+}
+
+static inline int
+ImagingCheckedMul(size_t a, size_t b, size_t *result) {
+#ifdef IMAGING_HAVE_OVERFLOW_BUILTINS
+    size_t r;
+    if (__builtin_mul_overflow(a, b, &r)) {
+        return 0;
+    }
+    *result = r;
+#else
+    if (b != 0 && a > SIZE_MAX / b) {
+        return 0;
+    }
+    *result = a * b;
+#endif
+    return 1;
+}
+
+#endif

src/libImaging/Quant.c

@@ -28,6 +28,7 @@
 #include "QuantPngQuant.h"
 #include "QuantHash.h"
 #include "QuantHeap.h"
+#include "CheckedArithmetic.h"
 
 typedef struct {
     uint32_t scale;
@@ -564,6 +565,32 @@ resort_distance_tables(
     return 1;
 }
 
+/* The distance tables have nEntries * nEntries elements and are indexed with
+   uint32_t arithmetic, so that product has to fit in a uint32_t as well. */
+static int
+alloc_distance_tables(
+    uint32_t **avgDist, uint32_t ***avgDistSortKey, uint32_t nEntries
+) {
+    size_t n;
+
+    *avgDist = NULL;
+    *avgDistSortKey = NULL;
+    if (!ImagingCheckedMul(nEntries, nEntries, &n) || n > UINT32_MAX) {
+        return 0;
+    }
+    /* malloc check ok, using calloc for overflow */
+    *avgDist = calloc(n, sizeof(uint32_t));
+    *avgDistSortKey = calloc(n, sizeof(uint32_t *));
+    if (!*avgDist || !*avgDistSortKey) {
+        free(*avgDist);
+        free(*avgDistSortKey);
+        *avgDist = NULL;
+        *avgDistSortKey = NULL;
+        return 0;
+    }
+    return 1;
+}
+
 static int
 build_distance_tables(
     uint32_t *avgDist, uint32_t **avgDistSortKey, Pixel *p, uint32_t nEntries
@@ -899,9 +926,6 @@ k_means(
     int changes;
     int built = 0;
 
-    if (nPaletteEntries > UINT32_MAX / (sizeof(uint32_t))) {
-        return 0;
-    }
     /* malloc check ok, using calloc */
     if (!(count = calloc(nPaletteEntries, sizeof(uint32_t)))) {
         return 0;
@@ -916,22 +940,10 @@ k_means(
         }
     }
 
-    /* this is enough of a check, since the multiplication n*size is done above */
-    if (nPaletteEntries > UINT32_MAX / nPaletteEntries) {
-        goto error_1;
-    }
-    /* malloc check ok, using calloc, checking n*n above */
-    avgDist = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t));
-    if (!avgDist) {
+    if (!alloc_distance_tables(&avgDist, &avgDistSortKey, nPaletteEntries)) {
         goto error_1;
     }
 
-    /* malloc check ok, using calloc, checking n*n above */
-    avgDistSortKey = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t *));
-    if (!avgDistSortKey) {
-        goto error_2;
-    }
-
     while (1) {
         if (!built) {
             compute_palette_from_quantized_pixels(
@@ -987,7 +999,6 @@ k_means(
     if (avgDistSortKey) {
         free(avgDistSortKey);
     }
-error_2:
     if (avgDist) {
         free(avgDist);
     }
@@ -1062,21 +1073,10 @@ quantize(
         goto error_4;
     }
 
-    if (nPaletteEntries > UINT32_MAX / nPaletteEntries) {
-        goto error_5;
-    }
-    /* malloc check ok, using calloc for overflow, check of n*n above */
-    avgDist = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t));
-    if (!avgDist) {
+    if (!alloc_distance_tables(&avgDist, &avgDistSortKey, nPaletteEntries)) {
         goto error_5;
     }
 
-    /* malloc check ok, using calloc for overflow, check of n*n above */
-    avgDistSortKey = calloc(nPaletteEntries * nPaletteEntries, sizeof(uint32_t *));
-    if (!avgDistSortKey) {
-        goto error_6;
-    }
-
     if (!build_distance_tables(avgDist, avgDistSortKey, p, nPaletteEntries)) {
         goto error_7;
     }
@@ -1108,7 +1108,6 @@ quantize(
     if (avgDistSortKey) {
         free(avgDistSortKey);
     }
-error_6:
     if (avgDist) {
         free(avgDist);
     }
@@ -1208,22 +1207,10 @@ quantize2(
         goto error_1;
     }
 
-    if (nQuantPixels > UINT32_MAX / nQuantPixels) {
+    if (!alloc_distance_tables(&avgDist, &avgDistSortKey, nQuantPixels)) {
         goto error_2;
     }
 
-    /* malloc check ok, using calloc for overflow, check of n*n above */
-    avgDist = calloc(nQuantPixels * nQuantPixels, sizeof(uint32_t));
-    if (!avgDist) {
-        goto error_2;
-    }
-
-    /* malloc check ok, using calloc for overflow, check of n*n above */
-    avgDistSortKey = calloc(nQuantPixels * nQuantPixels, sizeof(uint32_t *));
-    if (!avgDistSortKey) {
-        goto error_3;
-    }
-
     if (!build_distance_tables(avgDist, avgDistSortKey, p, nQuantPixels)) {
         goto error_4;
     }
@@ -1246,7 +1233,6 @@ quantize2(
 
 error_4:
     free(avgDistSortKey);
-error_3:
     free(avgDist);
 error_2:
     free(qp);
@@ -1297,11 +1283,13 @@ ImagingQuantize(Imaging im, int colors, int mode, int kmeans) {
     // Hoisted here as these are invariant over the loops below.
     int xsize = im->xsize, ysize = im->ysize;
 
-    if (xsize > INT_MAX / ysize) {
+    /* the pixel index below is an int, so the pixel count has to fit in one */
+    size_t nPixels;
+    if (!ImagingCheckedMul(xsize, ysize, &nPixels) || nPixels > INT_MAX) {
         return ImagingError_MemoryError();
     }
-    /* malloc check ok, using calloc for final overflow, x*y above */
-    p = calloc((size_t)xsize * ysize, sizeof(Pixel));
+    /* malloc check ok, using calloc for overflow */
+    p = calloc(nPixels, sizeof(Pixel));
     if (!p) {
         return ImagingError_MemoryError();
     }

src/libImaging/QuantHash.c

@@ -56,12 +56,11 @@ hashtable_new(HashFunc hf, HashCmpFunc cf) {
     h->length = MIN_LENGTH;
     h->count = 0;
     h->userData = NULL;
-    h->table = malloc(sizeof(HashNode *) * h->length);
+    h->table = calloc(h->length, sizeof(HashNode *));
     if (!h->table) {
         free(h);
         return NULL;
     }
-    memset(h->table, 0, sizeof(HashNode *) * h->length);
     return h;
 }
 
@@ -94,14 +93,13 @@ _hashtable_rehash(HashTable *h, CollisionFunc cf, uint32_t newSize) {
     HashNode *n, *nn;
     uint32_t oldSize;
     oldSize = h->length;
-    h->table = malloc(sizeof(HashNode *) * newSize);
+    h->table = calloc(newSize, sizeof(HashNode *));
     if (!h->table) {
         h->table = oldTable;
         return;
     }
     h->length = newSize;
     h->count = 0;
-    memset(h->table, 0, sizeof(HashNode *) * h->length);
     for (i = 0; i < oldSize; i++) {
         for (n = oldTable[i]; n; n = nn) {
             nn = n->next;

src/libImaging/QuantOctree.c

@@ -30,6 +30,7 @@
 
 #include "ImagingUtils.h"
 #include "QuantOctree.h"
+#include "CheckedArithmetic.h"
 
 typedef struct _ColorBucket {
     /* contains palette index when used for look up cube */
@@ -173,15 +174,16 @@ compare_bucket_count(const ColorBucket a, const ColorBucket b) {
 static ColorBucket
 create_sorted_color_palette(const ColorCube cube) {
     ColorBucket buckets;
-    if (cube->size > LONG_MAX / sizeof(struct _ColorBucket)) {
+    size_t size;
+    if (!ImagingCheckedMul(cube->size, sizeof(struct _ColorBucket), &size)) {
         return NULL;
     }
-    /* malloc check ok, calloc + overflow check above for memcpy */
-    buckets = calloc(cube->size, sizeof(struct _ColorBucket));
+    /* malloc check ok, overflow checked above */
+    buckets = malloc(size);
     if (!buckets) {
         return NULL;
     }
-    memcpy(buckets, cube->buckets, sizeof(struct _ColorBucket) * cube->size);
+    memcpy(buckets, cube->buckets, size);
 
     qsort(
         buckets,
@@ -332,12 +334,12 @@ combined_palette(
     unsigned long nBucketsB
 ) {
     ColorBucket result;
-    if (nBucketsA > LONG_MAX - nBucketsB ||
-        (nBucketsA + nBucketsB) > LONG_MAX / sizeof(struct _ColorBucket)) {
+    size_t nBuckets;
+    if (!ImagingCheckedAdd(nBucketsA, nBucketsB, &nBuckets)) {
         return NULL;
     }
-    /* malloc check ok, overflow check above */
-    result = calloc(nBucketsA + nBucketsB, sizeof(struct _ColorBucket));
+    /* malloc check ok, using calloc for overflow, sum checked above */
+    result = calloc(nBuckets, sizeof(struct _ColorBucket));
     if (!result) {
         return NULL;
     }
@@ -508,7 +510,7 @@ quantize_octree(
 
     /* create result pixels and map palette indices */
     /* malloc check ok, calloc for overflow */
-    qp = calloc(nPixels, sizeof(Pixel));
+    qp = calloc(nPixels, sizeof(uint32_t));
     if (!qp) {
         goto error;
     }