Validate image size in pixels at upload time (#3003)

src/tests/api/test_upload.py

@@ -19,8 +19,14 @@
 # You should have received a copy of the GNU Affero General Public License along with this program.  If not, see
 # <https://www.gnu.org/licenses/>.
 #
+import io
+import math
+
 import pytest
 from django.core.files.base import ContentFile
+from PIL import Image
+from PIL.Image import MAX_IMAGE_PIXELS
+from tests.const import SAMPLE_PNG
 
 
 @pytest.mark.django_db
@@ -29,7 +35,7 @@ def test_upload_file(token_client):
         '/api/v1/upload',
         data={
             'media_type': 'application/pdf',
-            'file': ContentFile('file.pdf', 'invalid pdf content')
+            'file': ContentFile('invalid pdf content')
         },
         format='upload',
         HTTP_CONTENT_DISPOSITION='attachment; filename="file.pdf"',
@@ -44,9 +50,74 @@ def test_upload_file_extension_mismatch(token_client):
         '/api/v1/upload',
         data={
             'media_type': 'application/pdf',
-            'file': ContentFile('file.png', 'invalid pdf content')
+            'file': ContentFile('invalid pdf content')
         },
         format='upload',
         HTTP_CONTENT_DISPOSITION='attachment; filename="file.png"',
     )
     assert r.status_code == 400
+    assert r.data == ['File name "file.png" has an invalid extension for type "application/pdf"']
+
+
+@pytest.mark.django_db
+def test_upload_file_extension_not_allowed(token_client):
+    r = token_client.post(
+        '/api/v1/upload',
+        data={
+            'media_type': 'application/octet-stream',
+            'file': ContentFile('invalid pdf content')
+        },
+        format='upload',
+        HTTP_CONTENT_DISPOSITION='attachment; filename="file.bin"',
+    )
+    assert r.status_code == 400
+    assert r.data == ['Content type "application/octet-stream" is not allowed']
+
+
+@pytest.mark.django_db
+def test_upload_invalid_image(token_client):
+    r = token_client.post(
+        '/api/v1/upload',
+        data={
+            'media_type': 'image/png',
+            'file': ContentFile('invalid png content')
+        },
+        format='upload',
+        HTTP_CONTENT_DISPOSITION='attachment; filename="file.png"',
+    )
+    assert r.status_code == 400
+    assert r.data == ['Upload a valid image. The file you uploaded was either not an image or a corrupted image.']
+
+
+@pytest.mark.django_db
+def test_upload_valid_image(token_client):
+    r = token_client.post(
+        '/api/v1/upload',
+        data={
+            'media_type': 'image/png',
+            'file': ContentFile(SAMPLE_PNG)
+        },
+        format='upload',
+        HTTP_CONTENT_DISPOSITION='attachment; filename="file.png"',
+    )
+    assert r.status_code == 201
+
+
+@pytest.mark.django_db
+@pytest.mark.filterwarnings("ignore")
+def test_upload_image_with_invalid_dimensions(token_client):
+    d = int(math.sqrt(MAX_IMAGE_PIXELS)) + 100
+    img = Image.new('RGB', (d, d), color='red')
+    output = io.BytesIO()
+    img.save(output, format='PNG')
+    r = token_client.post(
+        '/api/v1/upload',
+        data={
+            'media_type': 'image/png',
+            'file': ContentFile(output.getvalue())
+        },
+        format='upload',
+        HTTP_CONTENT_DISPOSITION='attachment; filename="file.png"',
+    )
+    assert r.status_code == 400
+    assert r.data == ['The file you uploaded has a very large number of pixels, please upload a picture with smaller dimensions.']