Address code review: correctness fixes and tests for clean() re-encoding
- Move extension allowlist check before PillowImage.open() so SVG/unknown formats get a clear error rather than a generic 'unable to process' message. - Replace hasattr(self.image, 'file') guard with not self.image._committed so already-stored images are not unnecessarily re-encoded on every save. - Replace bare BytesIO assignment with ContentFile so Django's pre_save() reliably writes the sanitised content to storage. - Tighten except clause to (OSError, UnidentifiedImageError, DecompressionBombError). - Use ImageSequence.Iterator for animated GIF re-encoding to preserve all frames; pass duration/loop metadata through. - Add four model tests covering polyglot stripping, animated GIF frame preservation, disallowed extension rejection, and committed-image skip. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Brian Tiemann
parent
a0c91099d8
commit
0227be7ebe
|
|
@ -6,6 +6,7 @@ from pathlib import Path
|
|||
from django.conf import settings
|
||||
from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
|
||||
from django.contrib.postgres.fields import ArrayField
|
||||
from django.core.files.base import ContentFile
|
||||
from django.core.validators import ValidationError
|
||||
from django.db import models
|
||||
from django.urls import reverse
|
||||
|
|
@ -743,30 +744,47 @@ class ImageAttachment(ChangeLoggedModel):
|
|||
_("Image attachments cannot be assigned to this object type ({type}).").format(type=self.object_type)
|
||||
)
|
||||
|
||||
# Re-encode the image through Pillow to strip any embedded non-image
|
||||
# payloads (e.g. polyglot HTML/PNG files). This must be done after the
|
||||
# ImageField's own Pillow dimension check has already run.
|
||||
if self.image and hasattr(self.image, 'file'):
|
||||
# Re-encode new image uploads through Pillow to strip any embedded
|
||||
# non-image payloads (e.g. polyglot HTML/PNG files).
|
||||
# Only process uncommitted (new) uploads; skip already-stored images.
|
||||
if self.image and not self.image._committed:
|
||||
try:
|
||||
self.image.file.seek(0)
|
||||
img = PillowImage.open(self.image.file)
|
||||
img_format = img.format or 'PNG'
|
||||
# Validate extension against allowed formats
|
||||
# Validate extension before attempting to open with Pillow so
|
||||
# that disallowed types (e.g. SVG) get a clear error message
|
||||
# rather than a generic "unable to process" one.
|
||||
ext = Path(self.image.name).suffix.lstrip('.').lower()
|
||||
if ext not in IMAGE_ATTACHMENT_IMAGE_FORMATS:
|
||||
raise ValidationError(
|
||||
_("Unsupported image format: {ext}").format(ext=ext)
|
||||
)
|
||||
# Re-encode to a new buffer; this strips any non-image trailer data
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format=img_format)
|
||||
buf.seek(0)
|
||||
self.image.file = buf
|
||||
self.image.file.seek(0)
|
||||
img = PillowImage.open(self.image.file)
|
||||
img_format = img.format or 'PNG'
|
||||
# Re-encode into a fresh buffer; any non-image trailer bytes
|
||||
# are discarded. For animated GIFs, iterate through all frames
|
||||
# via ImageSequence so animation is preserved.
|
||||
buf = io.BytesIO()
|
||||
if img_format == 'GIF':
|
||||
from PIL import ImageSequence # noqa: PLC0415
|
||||
frames = [frame.copy() for frame in ImageSequence.Iterator(img)]
|
||||
if frames:
|
||||
frames[0].save(
|
||||
buf, format='GIF', save_all=True,
|
||||
append_images=frames[1:] if len(frames) > 1 else [],
|
||||
duration=img.info.get('duration', 100),
|
||||
loop=img.info.get('loop', 0),
|
||||
)
|
||||
else:
|
||||
img.save(buf, format='GIF')
|
||||
else:
|
||||
img.save(buf, format=img_format)
|
||||
# Replace the in-memory upload with the sanitised content so
|
||||
# Django writes the clean version to storage on model save.
|
||||
self.image.file = ContentFile(buf.getvalue(), name=self.image.name)
|
||||
except ValidationError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ValidationError(_("Unable to process the uploaded image."))
|
||||
except (OSError, PillowImage.UnidentifiedImageError, PillowImage.DecompressionBombError) as e:
|
||||
raise ValidationError(_("Unable to process the uploaded image.")) from e
|
||||
|
||||
def delete(self, *args, **kwargs):
|
||||
|
||||
|
|
|
|||
|
|
@ -189,7 +189,91 @@ class ImageAttachmentTestCase(TestCase):
|
|||
self.assertEqual(first.filename, 'action-buttons.png')
|
||||
self.assertEqual(second.filename, 'action-buttons_sdmmer4.png')
|
||||
|
||||
self.assertCountEqual(storage.files.keys(), {base_name, suffixed_name})
|
||||
# ------------------------------------------------------------------
|
||||
# Security: polyglot upload sanitisation (SR-001 / VM-326)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _make_polyglot_png(self, filename, payload=b'<script>alert(1)</script>'):
|
||||
"""Return a SimpleUploadedFile that is a valid PNG with extra bytes appended."""
|
||||
buf = io.BytesIO()
|
||||
Image.new('RGB', (1, 1), color='red').save(buf, format='PNG')
|
||||
# Append a non-image payload after the IEND chunk
|
||||
polyglot = buf.getvalue() + payload
|
||||
return SimpleUploadedFile(name=filename, content=polyglot, content_type='image/png')
|
||||
|
||||
def _make_animated_gif(self, filename, frames=2):
|
||||
"""Return a SimpleUploadedFile containing a minimal animated GIF."""
|
||||
buf = io.BytesIO()
|
||||
# Use distinct RGB colours and a size > 1×1 so GIF optimisation
|
||||
# does not collapse identical frames into a single frame.
|
||||
images = [Image.new('RGB', (2, 2), color=(i * 80, 0, 0)) for i in range(frames)]
|
||||
images[0].save(buf, format='GIF', save_all=True, append_images=images[1:], loop=0, duration=100)
|
||||
return SimpleUploadedFile(name=filename, content=buf.getvalue(), content_type='image/gif')
|
||||
|
||||
def _make_attachment(self, uploaded_file):
|
||||
return ImageAttachment(
|
||||
object_type=self.ct_site,
|
||||
object_id=self.site.pk,
|
||||
image=uploaded_file,
|
||||
)
|
||||
|
||||
def test_clean_strips_polyglot_payload_from_png(self):
|
||||
"""clean() re-encodes PNG uploads, stripping any non-image trailer data."""
|
||||
payload = b'<script>alert(1)</script>'
|
||||
ia = self._make_attachment(self._make_polyglot_png('test.png', payload=payload))
|
||||
# The raw upload contains the payload
|
||||
self.assertIn(payload, ia.image.file.read())
|
||||
ia.image.file.seek(0)
|
||||
|
||||
ia.full_clean()
|
||||
|
||||
# After clean(), the sanitised ContentFile must not contain the payload
|
||||
sanitised = ia.image.file.read()
|
||||
self.assertNotIn(payload, sanitised)
|
||||
# But it must still be a valid PNG
|
||||
buf = io.BytesIO(sanitised)
|
||||
img = Image.open(buf)
|
||||
self.assertEqual(img.format, 'PNG')
|
||||
|
||||
def test_clean_preserves_animated_gif_frames(self):
|
||||
"""clean() re-encodes GIF uploads while keeping all animation frames."""
|
||||
n_frames = 3
|
||||
ia = self._make_attachment(self._make_animated_gif('anim.gif', frames=n_frames))
|
||||
ia.full_clean()
|
||||
|
||||
sanitised = ia.image.file.read()
|
||||
buf = io.BytesIO(sanitised)
|
||||
img = Image.open(buf)
|
||||
self.assertEqual(img.format, 'GIF')
|
||||
# All frames must be preserved
|
||||
frame_count = 0
|
||||
try:
|
||||
while True:
|
||||
frame_count += 1
|
||||
img.seek(img.tell() + 1)
|
||||
except EOFError:
|
||||
pass
|
||||
self.assertEqual(frame_count, n_frames)
|
||||
|
||||
def test_clean_rejects_disallowed_extension(self):
|
||||
"""clean() raises ValidationError for file extensions not in IMAGE_ATTACHMENT_IMAGE_FORMATS."""
|
||||
svg_content = b'<svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>'
|
||||
ia = self._make_attachment(
|
||||
SimpleUploadedFile(name='evil.svg', content=svg_content, content_type='image/svg+xml')
|
||||
)
|
||||
with self.assertRaises(ValidationError) as cm:
|
||||
ia.clean()
|
||||
self.assertIn('svg', str(cm.exception).lower())
|
||||
|
||||
def test_clean_skips_reencode_for_existing_stored_image(self):
|
||||
"""clean() does not re-encode already-committed (stored) images."""
|
||||
ia = self._make_attachment(self._uploaded_png('existing.png'))
|
||||
# Mark as committed (simulates an already-saved instance fetched from DB)
|
||||
ia.image._committed = True
|
||||
original_file = ia.image.file
|
||||
# Call model clean() directly; the re-encoding block must be skipped
|
||||
ia.clean()
|
||||
self.assertIs(ia.image.file, original_file)
|
||||
|
||||
|
||||
class TableConfigTestCase(TestCase):
|
||||
|
|
|
|||
Loading…
Reference in New Issue