From 2f723fb707737e07953d455ad9b62ff70f732a17 Mon Sep 17 00:00:00 2001 From: Yashvardhan Thanvi Date: Wed, 14 Oct 2026 10:39:00 +0530 Subject: [PATCH] Check size arithmetic where it is used in the quantizers Add ImagingCheckedMul() and ImagingCheckedAdd() in CheckedArithmetic.h for computing sizes without overflowing, and use them for the allocation sizes in Quant.c, QuantOctree.c and QuantPngQuant.c, so that each size is checked in the function that computes it. Some of these sizes were only safe because of checks made by the callers. Allocate the two distance tables in one function instead of three copies, allocate the hash table arrays with calloc() instead of malloc() and memset(), and use sizeof(uint32_t) for the uint32_t output of quantize_octree(). The output is unchanged for valid images. --- Tests/test_image_quantize.py | 40 +++++++++++++++ checks/check_large_memory.py | 6 +++ src/libImaging/CheckedArithmetic.h | 61 ++++++++++++++++++++++ src/libImaging/Quant.c | 82 +++++++++++++----------------- src/libImaging/QuantHash.c | 6 +-- src/libImaging/QuantOctree.c | 20 ++++---- src/libImaging/QuantPngQuant.c | 20 +++++--- 7 files changed, 168 insertions(+), 67 deletions(-) create mode 100644 src/libImaging/CheckedArithmetic.h diff --git a/Tests/test_image_quantize.py b/Tests/test_image_quantize.py index ec1c13977d1..1a1fc875adf 100644 --- a/Tests/test_image_quantize.py +++ b/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() diff --git a/checks/check_large_memory.py b/checks/check_large_memory.py index 012a94b47fc..f0116eff55f 100644 --- a/checks/check_large_memory.py +++ b/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() diff --git a/src/libImaging/CheckedArithmetic.h b/src/libImaging/CheckedArithmetic.h new file mode 100644 index 00000000000..d4e966c46d9 --- /dev/null +++ b/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 +#include + +/* 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 diff --git a/src/libImaging/Quant.c b/src/libImaging/Quant.c index 6d3dd2757c3..4684a4a9335 100644 --- a/src/libImaging/Quant.c +++ b/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(); } diff --git a/src/libImaging/QuantHash.c b/src/libImaging/QuantHash.c index bf2f29fde85..c32725bd70e 100644 --- a/src/libImaging/QuantHash.c +++ b/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; diff --git a/src/libImaging/QuantOctree.c b/src/libImaging/QuantOctree.c index ce62f0a3017..1ba65606f98 100644 --- a/src/libImaging/QuantOctree.c +++ b/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; } diff --git a/src/libImaging/QuantPngQuant.c b/src/libImaging/QuantPngQuant.c index 89e7bf35ce0..f43a4dc74de 100644 --- a/src/libImaging/QuantPngQuant.c +++ b/src/libImaging/QuantPngQuant.c @@ -13,6 +13,7 @@ #include #include "QuantPngQuant.h" +#include "CheckedArithmetic.h" #ifdef HAVE_LIBIMAGEQUANT #include @@ -34,11 +35,16 @@ quantize_pngquant( liq_result *remap = NULL; unsigned char *charMatrix = NULL; unsigned char **charMatrixRows = NULL; - unsigned int i, y; + unsigned int y; + size_t i, nPixels; *palette = NULL; *paletteLength = 0; *quantizedPixels = NULL; + if (!ImagingCheckedMul(width, height, &nPixels)) { + goto err; + } + /* configure pngquant */ attr = liq_attr_create(); if (!attr) { @@ -65,7 +71,7 @@ quantize_pngquant( /* write output palette */ const liq_palette *l_palette = liq_get_palette(remap); *paletteLength = l_palette->count; - *palette = malloc(sizeof(Pixel) * l_palette->count); + *palette = calloc(l_palette->count, sizeof(Pixel)); if (!*palette) { goto err; } @@ -77,27 +83,27 @@ quantize_pngquant( } /* write output pixels (pngquant uses char array) */ - charMatrix = malloc(width * height); + charMatrix = malloc(nPixels); if (!charMatrix) { goto err; } - charMatrixRows = malloc(height * sizeof(unsigned char *)); + charMatrixRows = calloc(height, sizeof(unsigned char *)); if (!charMatrixRows) { goto err; } for (y = 0; y < height; y++) { - charMatrixRows[y] = &charMatrix[y * width]; + charMatrixRows[y] = &charMatrix[(size_t)y * width]; } if (LIQ_OK != liq_write_remapped_image_rows(remap, image, charMatrixRows)) { goto err; } /* transcribe output pixels (pillow uses uint32_t array) */ - *quantizedPixels = malloc(sizeof(uint32_t) * width * height); + *quantizedPixels = calloc(nPixels, sizeof(uint32_t)); if (!*quantizedPixels) { goto err; } - for (i = 0; i < width * height; i++) { + for (i = 0; i < nPixels; i++) { (*quantizedPixels)[i] = charMatrix[i]; }