From 4ceec30e0c49d065771915d0c68a7051d626e49e Mon Sep 17 00:00:00 2001
From: 1e99 <noemail@noemail>
Date: Mon, 16 Dec 2024 18:10:18 +0100
Subject: [PATCH] improve error handling, limit upload size

---
 .../eu/e99/pixelchat/server/ImageHandler.java | 85 ++++++++++++++-----
 server/src/eu/e99/pixelchat/server/Main.java  |  6 +-
 2 files changed, 66 insertions(+), 25 deletions(-)

diff --git a/server/src/eu/e99/pixelchat/server/ImageHandler.java b/server/src/eu/e99/pixelchat/server/ImageHandler.java
index be7c659..d8a6891 100644
--- a/server/src/eu/e99/pixelchat/server/ImageHandler.java
+++ b/server/src/eu/e99/pixelchat/server/ImageHandler.java
@@ -19,40 +19,79 @@ public class ImageHandler {
     private final Logger logger;
     private final Storage storage;
     private final TemporalAmount expireTime;
+    private final int maxImageSize;
 
-    public ImageHandler(Logger logger, Storage storage, TemporalAmount expireTime) {
+    public ImageHandler(Logger logger, Storage storage, TemporalAmount expireTime, int maxImageSize) {
         this.logger = logger;
         this.storage = storage;
         this.expireTime = expireTime;
+        this.maxImageSize = maxImageSize;
     }
 
-    public void uploadImage(Context ctx) throws IOException {
-        InputStream body = ctx.bodyInputStream();
-        BufferedImage image = Imaging.getBufferedImage(body);
+    public void uploadImage(Context ctx) {
+        try {
+            InputStream body = ctx.bodyInputStream();
+            BufferedImage image = Imaging.getBufferedImage(body);
 
-        byte[] png = Imaging.writeImageToBytes(image, ImageFormats.PNG);
-        Instant expiresAt = Instant.now().plus(this.expireTime);
+            byte[] png = Imaging.writeImageToBytes(image, ImageFormats.PNG);
+            if (png.length > maxImageSize) {
+                ctx
+                        .status(HttpStatus.CONTENT_TOO_LARGE)
+                        .result("Image too large");
 
-        String id = storage.put(png, expiresAt);
-        logger.info("Stored {} bytes as {}", png.length, id);
+                return;
+            }
 
-        ctx.
-                status(HttpStatus.CREATED).
-                contentType(ContentType.TEXT_PLAIN).
-                result(id);
-    }
+            Instant expiresAt = Instant.now().plus(this.expireTime);
 
-    public void downloadImage(Context ctx) throws IOException {
-        String id = ctx.pathParam("id");
-        byte[] png = storage.get(id);
-        if (png == null) {
-            ctx.status(HttpStatus.NOT_FOUND);
-            return;
+            String id = storage.put(png, expiresAt);
+            logger.info("Uploaded {} image bytes as {}", png.length, id);
+
+            ctx.
+                    status(HttpStatus.CREATED).
+                    contentType(ContentType.TEXT_PLAIN).
+                    result(id);
+        } catch (IllegalArgumentException ignored) {
+            ctx
+                    .status(HttpStatus.BAD_REQUEST)
+                    .result("Unrecognized image format");
+        } catch (IOException e) {
+            ctx
+                    .status(HttpStatus.INTERNAL_SERVER_ERROR)
+                    .result("Internal Server Error");
+
+            this.logger.error("Failed to upload image", e);
+        } catch (OutOfMemoryError e) {
+            ctx
+                    .status(HttpStatus.INSUFFICIENT_STORAGE)
+                    .result("Insufficient storage on server");
+
+            this.logger.error("Out of memory", e);
         }
+    }
 
-        ctx.
-                status(HttpStatus.OK).
-                contentType(ContentType.IMAGE_PNG).
-                result(png);
+    public void downloadImage(Context ctx) {
+        try {
+            String id = ctx.pathParam("id");
+            byte[] png = storage.get(id);
+            if (png == null) {
+                ctx
+                        .status(HttpStatus.NOT_FOUND)
+                        .result("Image not found");
+
+                return;
+            }
+
+            ctx.
+                    status(HttpStatus.OK).
+                    contentType(ContentType.IMAGE_PNG).
+                    result(png);
+        } catch (IOException e) {
+            ctx
+                    .status(HttpStatus.INTERNAL_SERVER_ERROR)
+                    .result("Internal Server Error");
+
+            this.logger.error("Failed to download image", e);
+        }
     }
 }
diff --git a/server/src/eu/e99/pixelchat/server/Main.java b/server/src/eu/e99/pixelchat/server/Main.java
index 3b17717..89a24bb 100644
--- a/server/src/eu/e99/pixelchat/server/Main.java
+++ b/server/src/eu/e99/pixelchat/server/Main.java
@@ -17,15 +17,17 @@ public class Main {
                 new MemoryStorage(),
                 Duration.ofSeconds(5),
                 Duration.ofMinutes(2),
+                16 * 1024 * 1024, // 16 MiB
                 3000
         );
     }
 
-    private static void run(Logger logger, Storage storage, Duration clearTime, TemporalAmount expireTime, int port) {
+    private static void run(Logger logger, Storage storage, Duration clearTime, TemporalAmount expireTime, int maxImageSize, int port) {
         ImageHandler handler = new ImageHandler(
                 logger,
                 storage,
-                expireTime
+                expireTime,
+                maxImageSize
         );
 
         Thread.ofVirtual().start(() -> {