- Test file type checks

- Test file size checks

Author: alexlambson

kirovy/constants/api_codes.py

@@ -31,3 +31,4 @@ class FileUploadApiCodes(enum.StrEnum):
 
     e.g. a temporary map does not support custom image uploads.
     """
+    TOO_LARGE = "file-too-large"

kirovy/views/base_views.py

@@ -28,7 +28,7 @@
 from kirovy.request import KirovyRequest
 from kirovy.response import KirovyResponse
 from kirovy.serializers import KirovySerializer, CncNetUserOwnedModelSerializer
-
+from kirovy.utils import file_utils
 
 _LOGGER = logging.get_logger(__name__)
 
@@ -176,6 +176,8 @@ class FileUploadBaseView(APIView, metaclass=ABCMeta):
 
     serializer_class: t.ClassVar[t.Type[CncNetUserOwnedModelSerializer]]
 
+    max_file_size: t.ClassVar[file_utils.ByteSized] = file_utils.ByteSized(mega=3)
+
     def get_parent_object(self, request: KirovyRequest) -> GameScopedUserOwnedModel:
 
         parent_object_id = request.data.get(self.file_parent_attr_name)
@@ -193,6 +195,14 @@ def get_parent_object(self, request: KirovyRequest) -> GameScopedUserOwnedModel:
 
     def post(self, request: KirovyRequest, format=None) -> KirovyResponse[ui_objects.ResultResponseData]:
         uploaded_file: UploadedFile = request.data["file"]
+
+        if file_utils.ByteSized(uploaded_file.size) > self.max_file_size:
+            raise KirovyValidationError(
+                "File too large",
+                code=api_codes.FileUploadApiCodes.TOO_LARGE,
+                additional={"max_size": str(self.max_file_size)},
+            )
+
         parent_object = self.get_parent_object(request)
         self.extra_verification(request, uploaded_file, parent_object)
 

kirovy/views/cnc_map_views.py

@@ -23,7 +23,6 @@
     CncMap,
     CncMapFile,
 )
-from kirovy.models.cnc_game import GameScopedUserOwnedModel
 from kirovy.models.cnc_map import CncMapImageFile
 from kirovy.objects import ui_objects
 from kirovy.request import KirovyRequest
@@ -341,6 +340,8 @@ def extra_serializer_data(
                 height = image.height
                 width = image.width
         except (DecompressionBombError, UnidentifiedImageError) as e:
+            # This should be in verification, but we need the width and height,
+            # and loading the image twice is inefficient.
             _LOGGER.warning(
                 "user-attempted-bad-image-upload",
                 {"user_id": request.user.id, "username": request.user.username, "e": str(e)},

tests/test_views/test_map_upload.py

@@ -222,3 +222,38 @@ def test_map_image_upload__map_is_banned(create_cnc_map, file_map_image, client_
     assert response.status_code == status.HTTP_403_FORBIDDEN
 
     assert cnc_map.cncmapimagefile_set.select_related().count() == original_image_count
+
+
+def test_map_image_upload__not_an_image(create_cnc_map, file_map_desert, client_user):
+    """Test that map image uploads fail for non-image files."""
+    cnc_map = create_cnc_map(user_id=client_user.kirovy_user.id)
+    original_image_count = cnc_map.cncmapimagefile_set.select_related().count()
+    response = client_user.post(
+        "/maps/img/",
+        {"file": file_map_desert, "cnc_map_id": str(cnc_map.id)},
+        format="multipart",
+        content_type=None,
+    )
+
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert response.data["code"] == UploadApiCodes.FILE_EXTENSION_NOT_SUPPORTED
+    assert cnc_map.cncmapimagefile_set.select_related().count() == original_image_count
+
+
+def test_map_image_upload__too_large(create_cnc_map, file_binary, client_user):
+    """Test that map image uploads fail for files that are too large.
+
+    This test works because the file size check happens before the extension check
+    """
+    cnc_map = create_cnc_map(user_id=client_user.kirovy_user.id)
+    original_image_count = cnc_map.cncmapimagefile_set.select_related().count()
+    response = client_user.post(
+        "/maps/img/",
+        {"file": file_binary, "cnc_map_id": str(cnc_map.id)},
+        format="multipart",
+        content_type=None,
+    )
+
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert response.data["code"] == FileUploadApiCodes.TOO_LARGE
+    assert cnc_map.cncmapimagefile_set.select_related().count() == original_image_count