(svn r22873) -Fix [FS#4747]: Validate image dimensions before loading. (Based on patch by monoid)

3 files changed, 35 insertions, 3 deletions

diff --git a/src/heightmap.cpp b/src/heightmap.cpp
index 07b1e5e2a..b85bcd9c8 100644
--- a/src/heightmap.cpp
+++ b/src/heightmap.cpp
@@ -142,13 +142,24 @@ static bool ReadHeightmapPNG(char *filename, uint *x, uint *y, byte **map)
 		return false;
 	}
 
+	uint width = png_get_image_width(png_ptr, info_ptr);
+	uint height = png_get_image_height(png_ptr, info_ptr);
+
+	/* Check if image dimensions don't overflow a size_t to avoid memory corruption. */
+	if ((uint64)width * height >= (size_t)-1) {
+		ShowErrorMessage(STR_ERROR_PNGMAP, STR_ERROR_HEIGHTMAP_TOO_LARGE, WL_ERROR);
+		fclose(fp);
+		png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
+		return false;
+	}
+
 	if (map != NULL) {
-		*map = MallocT<byte>(png_get_image_width(png_ptr, info_ptr) * png_get_image_height(png_ptr, info_ptr));
+		*map = MallocT<byte>(width * height);
 		ReadHeightmapPNGImageData(*map, png_ptr, info_ptr);
 	}
 
-	*x = png_get_image_width(png_ptr, info_ptr);
-	*y = png_get_image_height(png_ptr, info_ptr);
+	*x = width;
+	*y = height;
 
 	fclose(fp);
 	png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
@@ -243,6 +254,14 @@ static bool ReadHeightmapBMP(char *filename, uint *x, uint *y, byte **map)
 		return false;
 	}
 
+	/* Check if image dimensions don't overflow a size_t to avoid memory corruption. */
+	if ((uint64)info.width * info.height >= (size_t)-1 / (info.bpp == 24 ? 3 : 1)) {
+		ShowErrorMessage(STR_ERROR_BMPMAP, STR_ERROR_HEIGHTMAP_TOO_LARGE, WL_ERROR);
+		fclose(f);
+		BmpDestroyData(&data);
+		return false;
+	}
+
 	if (map != NULL) {
 		if (!BmpReadBitmap(&buffer, &info, &data)) {
 			ShowErrorMessage(STR_ERROR_BMPMAP, STR_ERROR_BMPMAP_IMAGE_TYPE, WL_ERROR);
diff --git a/src/lang/english.txt b/src/lang/english.txt
index f8b341ad0..e2b778c3a 100644
--- a/src/lang/english.txt
+++ b/src/lang/english.txt
@@ -3452,6 +3452,8 @@ STR_ERROR_PNGMAP_MISC                                           :{WHITE}... some
 STR_ERROR_BMPMAP                                                :{WHITE}Can't load landscape from BMP...
 STR_ERROR_BMPMAP_IMAGE_TYPE                                     :{WHITE}... could not convert image type
 
+STR_ERROR_HEIGHTMAP_TOO_LARGE                                   :{WHITE}... image is too large
+
 STR_WARNING_HEIGHTMAP_SCALE_CAPTION                             :{WHITE}Scale warning
 STR_WARNING_HEIGHTMAP_SCALE_MESSAGE                             :{YELLOW}Resizing source map too much is not recommended. Continue with the generation?
 
diff --git a/src/spriteloader/png.cpp b/src/spriteloader/png.cpp
index d31184385..27ff3cc9b 100644
--- a/src/spriteloader/png.cpp
+++ b/src/spriteloader/png.cpp
@@ -108,7 +108,17 @@ static bool LoadPNG(SpriteLoader::Sprite *sprite, const char *filename, uint32 i
 
 		sprite->height = png_get_image_height(png_ptr, info_ptr);
 		sprite->width  = png_get_image_width(png_ptr, info_ptr);
+		/* Check if sprite dimensions aren't larger than what is allowed in GRF-files. */
+		if (sprite->height > UINT8_MAX || sprite->width > UINT16_MAX) {
+			png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
+			return false;
+		}
 		sprite->AllocateData(sprite->width * sprite->height);
+	} else if (sprite->height != png_get_image_height(png_ptr, info_ptr) || sprite->width != png_get_image_width(png_ptr, info_ptr)) {
+		/* Make sure the mask image isn't larger than the sprite image. */
+		DEBUG(misc, 0, "Ignoring mask for SpriteID %d as it isn't the same dimension as the masked sprite", id);
+		png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
+		return true;
 	}
 
 	bit_depth  = png_get_bit_depth(png_ptr, info_ptr);
@@ -116,6 +126,7 @@ static bool LoadPNG(SpriteLoader::Sprite *sprite, const char *filename, uint32 i
 
 	if (mask && (bit_depth != 8 || colour_type != PNG_COLOR_TYPE_PALETTE)) {
 		DEBUG(misc, 0, "Ignoring mask for SpriteID %d as it isn't a 8 bit palette image", id);
+		png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
 		return true;
 	}