Merge pull request #110 from Hipo/handle-os-error-in-base64-image-field

Handle invalid Base64 inputs in Base64ImageField
Hipo · Feb 7, 2020 · f6d5d0a · f6d5d0a
2 parents 18dc207 + 24acce3
commit f6d5d0a
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 6 deletions.
diff --git a/drf_extra_fields/fields.py b/drf_extra_fields/fields.py
@@ -126,7 +126,11 @@ def get_file_extension(self, filename, decoded_file):
         # Try with PIL as fallback if format not detected due
         # to bug in imghdr https://bugs.python.org/issue16512
         if extension is None:
-            image = Image.open(io.BytesIO(decoded_file))
+            try:
+                image = Image.open(io.BytesIO(decoded_file))
+            except (OSError, IOError):
+                raise ValidationError(self.INVALID_FILE_MESSAGE)
+
             extension = image.format.lower()
 
         extension = "jpg" if extension == "jpeg" else extension

diff --git a/setup.py b/setup.py
@@ -9,7 +9,7 @@
 
 setup(
     name='django-extra-fields',
-    version='2.0.2',
+    version='2.0.3',
     packages=['drf_extra_fields',
               'drf_extra_fields.runtests'],
     include_package_data=True,

diff --git a/tests/test_fields.py b/tests/test_fields.py
@@ -107,16 +107,26 @@ def test_create_with_base64_prefix(self):
         self.assertEqual(serializer.validated_data['created'], uploaded_image.created)
         self.assertFalse(serializer.validated_data is uploaded_image)
 
+    def test_create_with_invalid_base64(self):
+        """
+        Test for creating Base64 image with an invalid Base64 string in the server side
+        """
+        now = datetime.datetime.now()
+        file = 'this_is_not_a_base64'
+        serializer = UploadedBase64ImageSerializer(data={'created': now, 'file': file})
+
+        self.assertFalse(serializer.is_valid())
+        self.assertEqual(serializer.errors, {'file': [Base64ImageField.INVALID_FILE_MESSAGE]})
+
     def test_validation_error_with_non_file(self):
         """
         Passing non-base64 should raise a validation error.
         """
         now = datetime.datetime.now()
-        errmsg = "Please upload a valid image."
         serializer = UploadedBase64ImageSerializer(data={'created': now,
                                                          'file': 'abc'})
         self.assertFalse(serializer.is_valid())
-        self.assertEqual(serializer.errors, {'file': [errmsg]})
+        self.assertEqual(serializer.errors, {'file': [Base64ImageField.INVALID_FILE_MESSAGE]})
 
     def test_remove_with_empty_string(self):
         """
@@ -228,10 +238,9 @@ def test_validation_error_with_non_file(self):
         Passing non-base64 should raise a validation error.
         """
         now = datetime.datetime.now()
-        errmsg = "Please upload a valid file."
         serializer = UploadedBase64FileSerializer(data={'created': now, 'file': 'abc'})
         self.assertFalse(serializer.is_valid())
-        self.assertEqual(serializer.errors, {'file': [errmsg]})
+        self.assertEqual(serializer.errors, {'file': [Base64FileField.INVALID_FILE_MESSAGE]})
 
     def test_remove_with_empty_string(self):
         """