implement file size upload limit

owrx/controllers/__init__.py

@@ -1,6 +1,10 @@
 from datetime import datetime, timezone
 
 
+class BodySizeError(Exception):
+    pass
+
+
 class Controller(object):
     def __init__(self, handler, request, options):
         self.handler = handler
@@ -33,10 +37,12 @@ class Controller(object):
         self.handler.send_header("Location", location)
         self.handler.end_headers()
 
-    def get_body(self):
+    def get_body(self, max_size=None):
         if "Content-Length" not in self.handler.headers:
             return None
         length = int(self.handler.headers["Content-Length"])
+        if max_size is not None and length > max_size:
+            raise BodySizeError("HTTP body exceeds maximum allowed size")
         return self.handler.rfile.read(length)
 
     def handle_request(self):

owrx/controllers/imageupload.py

@@ -1,11 +1,20 @@
+from owrx.controllers import BodySizeError
 from owrx.controllers.assets import AssetsController
 from owrx.controllers.admin import AuthorizationMixin
 from owrx.config.core import CoreConfig
+from owrx.form.input.gfx import AvatarInput, TopPhotoInput
 import uuid
 import json
 
 
 class ImageUploadController(AuthorizationMixin, AssetsController):
+    # max upload filesizes
+    max_sizes = {
+        # not the best idea to instantiate inputs, but i didn't want to duplicate the sizes here
+        "receiver_avatar": AvatarInput("id", "label").getMaxSize(),
+        "receiver_top_photo": TopPhotoInput("id", "label").getMaxSize(),
+    }
+
     def __init__(self, handler, request, options):
         super().__init__(handler, request, options)
         self.file = request.query["file"][0] if "file" in request.query else None
@@ -29,22 +38,37 @@ class ImageUploadController(AuthorizationMixin, AssetsController):
 
     def processImage(self):
         if "id" not in self.request.query:
-            self.send_response("{}", content_type="application/json", code=400)
-        # TODO: limit file size
-        contents = self.get_body()
+            self.send_json_response({"error": "missing id"}, code=400)
+            return
+        file_id = self.request.query["id"][0]
+
+        if file_id not in ImageUploadController.max_sizes:
+            self.send_json_response({"error": "unexpected image id"}, code=400)
+            return
+
+        try:
+            contents = self.get_body(ImageUploadController.max_sizes[file_id])
+        except BodySizeError:
+            self.send_json_response({"error": "file size too large"}, code=400)
+            return
+
         filetype = None
         if self._is_png(contents):
             filetype = "png"
         if self._is_jpg(contents):
             filetype = "jpg"
         if filetype is None:
-            self.send_response("{}", content_type="application/json", code=400)
-        else:
-            self.file = "{id}-{uuid}.{ext}".format(
-                id=self.request.query["id"][0],
-                uuid=uuid.uuid4().hex,
-                ext=filetype,
-            )
-            with open(self.getFilePath(), "wb") as f:
-                f.write(contents)
-            self.send_response(json.dumps({"file": self.file}), content_type="application/json")
+            self.send_json_response({"error": "unsupported file type"}, code=400)
+            return
+
+        self.file = "{id}-{uuid}.{ext}".format(
+            id=file_id,
+            uuid=uuid.uuid4().hex,
+            ext=filetype,
+        )
+        with open(self.getFilePath(), "wb") as f:
+            f.write(contents)
+        self.send_json_response({"file": self.file}, code=200)
+
+    def send_json_response(self, obj, code):
+        self.send_response(json.dumps(obj), code=code, content_type="application/json")