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]; }