improve error handling, limit upload size

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 {
+    public void uploadImage(Context ctx) {
-        InputStream body = ctx.bodyInputStream();
+        try {
-        BufferedImage image = Imaging.getBufferedImage(body);
+            InputStream body = ctx.bodyInputStream();
+            BufferedImage image = Imaging.getBufferedImage(body);
 
-        byte[] png = Imaging.writeImageToBytes(image, ImageFormats.PNG);
+            byte[] png = Imaging.writeImageToBytes(image, ImageFormats.PNG);
-        Instant expiresAt = Instant.now().plus(this.expireTime);
+            if (png.length > maxImageSize) {
+                ctx
+                        .status(HttpStatus.CONTENT_TOO_LARGE)
+                        .result("Image too large");
 
-        String id = storage.put(png, expiresAt);
+                return;
-        logger.info("Stored {} bytes as {}", png.length, id);
+            }
 
-        ctx.
+            Instant expiresAt = Instant.now().plus(this.expireTime);
-                status(HttpStatus.CREATED).
-                contentType(ContentType.TEXT_PLAIN).
-                result(id);
-    }
 
-    public void downloadImage(Context ctx) throws IOException {
+            String id = storage.put(png, expiresAt);
-        String id = ctx.pathParam("id");
+            logger.info("Uploaded {} image bytes as {}", png.length, id);
-        byte[] png = storage.get(id);
+
-        if (png == null) {
+            ctx.
-            ctx.status(HttpStatus.NOT_FOUND);
+                    status(HttpStatus.CREATED).
-            return;
+                    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.
+    public void downloadImage(Context ctx) {
-                status(HttpStatus.OK).
+        try {
-                contentType(ContentType.IMAGE_PNG).
+            String id = ctx.pathParam("id");
-                result(png);
+            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);
+        }
     }
 }

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(() -> {