From d2ca3ed6fe8ca58ca31a441a3ebecdc943f4431f Mon Sep 17 00:00:00 2001
From: Jordan <dordsor21@gmail.com>
Date: Sun, 17 Mar 2024 20:25:09 +0100
Subject: [PATCH] fix: better image reading errors (#2632)

- closes #2593
---
 .../core/command/tool/brush/HeightBrush.java  |  2 +-
 .../core/util/MainUtil.java                   | 21 +++++++++++++++++--
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/worldedit-core/src/main/java/com/fastasyncworldedit/core/command/tool/brush/HeightBrush.java b/worldedit-core/src/main/java/com/fastasyncworldedit/core/command/tool/brush/HeightBrush.java
index fea4884f1..2c6ea5ec0 100644
--- a/worldedit-core/src/main/java/com/fastasyncworldedit/core/command/tool/brush/HeightBrush.java
+++ b/worldedit-core/src/main/java/com/fastasyncworldedit/core/command/tool/brush/HeightBrush.java
@@ -53,7 +53,7 @@ public class HeightBrush implements Brush {
             try {
                 heightMap = ScalableHeightMap.fromPNG(stream);
             } catch (IOException e) {
-                throw new FaweException(Caption.of("fawe.worldedit.brush.brush.height.invalid"));
+                throw new FaweException(Caption.of("fawe.worldedit.brush.brush.height.invalid", e.getMessage()));
             }
         } else if (clipboard != null) {
             heightMap = ScalableHeightMap.fromClipboard(clipboard, minY, maxY);
diff --git a/worldedit-core/src/main/java/com/fastasyncworldedit/core/util/MainUtil.java b/worldedit-core/src/main/java/com/fastasyncworldedit/core/util/MainUtil.java
index 8122a840c..cbe9c6897 100644
--- a/worldedit-core/src/main/java/com/fastasyncworldedit/core/util/MainUtil.java
+++ b/worldedit-core/src/main/java/com/fastasyncworldedit/core/util/MainUtil.java
@@ -38,6 +38,8 @@ import org.apache.logging.log4j.Logger;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import javax.imageio.ImageIO;
+import javax.imageio.ImageReadParam;
+import javax.imageio.ImageReader;
 import java.awt.Graphics2D;
 import java.awt.image.BufferedImage;
 import java.io.BufferedOutputStream;
@@ -70,6 +72,7 @@ import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -519,8 +522,22 @@ public class MainUtil {
         return destFile;
     }
 
-    public static BufferedImage readImage(InputStream in) throws IOException {
-        return MainUtil.toRGB(ImageIO.read(in));
+    public static BufferedImage readImage(InputStream stream) throws IOException {
+        Iterator<ImageReader> iter = ImageIO.getImageReaders(stream);
+        if (!iter.hasNext()) {
+            throw new IOException("Could not get image reader from stream.");
+        }
+        ImageReader reader = iter.next();
+        ImageReadParam param = reader.getDefaultReadParam();
+        reader.setInput(stream, true, true);
+        BufferedImage bi;
+        try {
+            bi = reader.read(0, param);
+        } finally {
+            reader.dispose();
+            stream.close();
+        }
+        return MainUtil.toRGB(bi);
     }
 
     public static BufferedImage readImage(URL url) throws IOException {