improve error handling, limit upload size

server/src/eu/e99/pixelchat/server/ImageHandler.java

@@ -19,34 +19,66 @@ 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) {
+        try {
             InputStream body = ctx.bodyInputStream();
             BufferedImage image = Imaging.getBufferedImage(body);
 
             byte[] png = Imaging.writeImageToBytes(image, ImageFormats.PNG);
+            if (png.length > maxImageSize) {
+                ctx
+                        .status(HttpStatus.CONTENT_TOO_LARGE)
+                        .result("Image too large");
+
+                return;
+            }
+
             Instant expiresAt = Instant.now().plus(this.expireTime);
 
             String id = storage.put(png, expiresAt);
-        logger.info("Stored {} bytes as {}", png.length, id);
+            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);
+        }
     }
 
-    public void downloadImage(Context ctx) throws IOException {
+    public void downloadImage(Context ctx) {
+        try {
             String id = ctx.pathParam("id");
             byte[] png = storage.get(id);
             if (png == null) {
-            ctx.status(HttpStatus.NOT_FOUND);
+                ctx
+                        .status(HttpStatus.NOT_FOUND)
+                        .result("Image not found");
+
                 return;
             }
 
@@ -54,5 +86,12 @@ public class ImageHandler {
                     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(() -> {