From 0153d6dbcdc14738de102d062ec9ca2b7cc5bcc3 Mon Sep 17 00:00:00 2001 From: Alex Vandiver Date: Wed, 12 Jun 2024 14:56:41 +0000 Subject: [PATCH] thumbnailing: Move resizing functions into zerver.lib.thumbnail. --- .../generate_integration_bots_avatars.py | 2 +- zerver/data_import/slack.py | 3 +- zerver/lib/avatar.py | 2 +- zerver/lib/import_realm.py | 3 +- zerver/lib/thumbnail.py | 141 +++++++++++++++++ zerver/lib/upload/base.py | 142 ------------------ zerver/lib/upload/local.py | 11 +- zerver/lib/upload/s3.py | 5 +- zerver/tests/test_auth_backends.py | 2 +- zerver/tests/test_realm_emoji.py | 2 +- zerver/tests/test_transfer.py | 2 +- zerver/tests/test_upload.py | 9 +- zerver/tests/test_upload_local.py | 2 +- zerver/tests/test_upload_s3.py | 12 +- 14 files changed, 165 insertions(+), 173 deletions(-) diff --git a/tools/setup/generate_integration_bots_avatars.py b/tools/setup/generate_integration_bots_avatars.py index 39d6fb6af5..9f1600e6c1 100755 --- a/tools/setup/generate_integration_bots_avatars.py +++ b/tools/setup/generate_integration_bots_avatars.py @@ -21,7 +21,7 @@ from PIL import Image from zerver.lib.integrations import INTEGRATIONS from zerver.lib.storage import static_path -from zerver.lib.upload.base import DEFAULT_AVATAR_SIZE, resize_avatar +from zerver.lib.thumbnail import DEFAULT_AVATAR_SIZE, resize_avatar def create_square_image(png: bytes) -> bytes: diff --git a/zerver/data_import/slack.py b/zerver/data_import/slack.py index b6545f8a56..ab9d314745 100644 --- a/zerver/data_import/slack.py +++ b/zerver/data_import/slack.py @@ -47,7 +47,8 @@ from zerver.data_import.slack_message_conversion import ( from zerver.lib.emoji import codepoint_to_name from zerver.lib.export import MESSAGE_BATCH_CHUNK_SIZE from zerver.lib.storage import static_path -from zerver.lib.upload.base import resize_logo, sanitize_name +from zerver.lib.thumbnail import resize_logo +from zerver.lib.upload.base import sanitize_name from zerver.models import ( CustomProfileField, CustomProfileFieldValue, diff --git a/zerver/lib/avatar.py b/zerver/lib/avatar.py index d4a50e8539..4e7882253b 100644 --- a/zerver/lib/avatar.py +++ b/zerver/lib/avatar.py @@ -9,8 +9,8 @@ from zerver.lib.avatar_hash import ( user_avatar_content_hash, user_avatar_path_from_ids, ) +from zerver.lib.thumbnail import MEDIUM_AVATAR_SIZE from zerver.lib.upload import get_avatar_url -from zerver.lib.upload.base import MEDIUM_AVATAR_SIZE from zerver.lib.url_encoding import append_url_query_string from zerver.models import UserProfile diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index 155c8c9e63..8481c87a07 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -32,9 +32,10 @@ from zerver.lib.push_notifications import sends_notifications_directly from zerver.lib.remote_server import maybe_enqueue_audit_log_upload from zerver.lib.server_initialization import create_internal_realm, server_initialized from zerver.lib.streams import render_stream_description +from zerver.lib.thumbnail import BadImageError from zerver.lib.timestamp import datetime_to_timestamp from zerver.lib.upload import upload_backend -from zerver.lib.upload.base import BadImageError, sanitize_name +from zerver.lib.upload.base import sanitize_name from zerver.lib.upload.s3 import get_bucket from zerver.lib.user_counts import realm_user_count_by_role from zerver.lib.user_groups import create_system_user_groups_for_realm diff --git a/zerver/lib/thumbnail.py b/zerver/lib/thumbnail.py index 2822d24c20..ba2afba8ff 100644 --- a/zerver/lib/thumbnail.py +++ b/zerver/lib/thumbnail.py @@ -1,8 +1,28 @@ +import io +from typing import Optional, Tuple from urllib.parse import urljoin from django.utils.http import url_has_allowed_host_and_scheme +from django.utils.translation import gettext as _ +from PIL import GifImagePlugin, Image, ImageOps, PngImagePlugin +from PIL.Image import DecompressionBombError from zerver.lib.camo import get_camo_url +from zerver.lib.exceptions import ErrorCode, JsonableError + +DEFAULT_AVATAR_SIZE = 100 +MEDIUM_AVATAR_SIZE = 500 +DEFAULT_EMOJI_SIZE = 64 + +# These sizes were selected based on looking at the maximum common +# sizes in a library of animated custom emoji, balanced against the +# network cost of very large emoji images. +MAX_EMOJI_GIF_SIZE = 128 +MAX_EMOJI_GIF_FILE_SIZE_BYTES = 128 * 1024 * 1024 # 128 kb + + +class BadImageError(JsonableError): + code = ErrorCode.BAD_IMAGE def user_uploads_or_external(url: str) -> bool: @@ -17,3 +37,124 @@ def generate_thumbnail_url(path: str, size: str = "0x0") -> str: if url_has_allowed_host_and_scheme(path, allowed_hosts=None): return path return get_camo_url(path) + + +def resize_avatar(image_data: bytes, size: int = DEFAULT_AVATAR_SIZE) -> bytes: + try: + im = Image.open(io.BytesIO(image_data)) + im = ImageOps.exif_transpose(im) + im = ImageOps.fit(im, (size, size), Image.Resampling.LANCZOS) + except OSError: + raise BadImageError(_("Could not decode image; did you upload an image file?")) + except DecompressionBombError: + raise BadImageError(_("Image size exceeds limit.")) + out = io.BytesIO() + if im.mode == "CMYK": + im = im.convert("RGB") + im.save(out, format="png") + return out.getvalue() + + +def resize_logo(image_data: bytes) -> bytes: + try: + im = Image.open(io.BytesIO(image_data)) + im = ImageOps.exif_transpose(im) + im.thumbnail((8 * DEFAULT_AVATAR_SIZE, DEFAULT_AVATAR_SIZE), Image.Resampling.LANCZOS) + except OSError: + raise BadImageError(_("Could not decode image; did you upload an image file?")) + except DecompressionBombError: + raise BadImageError(_("Image size exceeds limit.")) + out = io.BytesIO() + if im.mode == "CMYK": + im = im.convert("RGB") + im.save(out, format="png") + return out.getvalue() + + +def resize_animated(im: Image.Image, size: int = DEFAULT_EMOJI_SIZE) -> bytes: + assert im.n_frames > 1 + frames = [] + duration_info = [] + disposals = [] + # If 'loop' info is not set then loop for infinite number of times. + loop = im.info.get("loop", 0) + for frame_num in range(im.n_frames): + im.seek(frame_num) + new_frame = im.copy() + new_frame.paste(im, (0, 0), im.convert("RGBA")) + new_frame = ImageOps.pad(new_frame, (size, size), Image.Resampling.LANCZOS) + frames.append(new_frame) + if im.info.get("duration") is None: # nocoverage + raise BadImageError(_("Corrupt animated image.")) + duration_info.append(im.info["duration"]) + if isinstance(im, GifImagePlugin.GifImageFile): + disposals.append( + im.disposal_method # type: ignore[attr-defined] # private member missing from stubs + ) + elif isinstance(im, PngImagePlugin.PngImageFile): + disposals.append(im.info.get("disposal", PngImagePlugin.Disposal.OP_NONE)) + else: # nocoverage + raise BadImageError(_("Unknown animated image format.")) + out = io.BytesIO() + frames[0].save( + out, + save_all=True, + optimize=False, + format=im.format, + append_images=frames[1:], + duration=duration_info, + disposal=disposals, + loop=loop, + ) + + return out.getvalue() + + +def resize_emoji( + image_data: bytes, size: int = DEFAULT_EMOJI_SIZE +) -> Tuple[bytes, bool, Optional[bytes]]: + # This function returns three values: + # 1) Emoji image data. + # 2) If emoji is gif i.e. animated. + # 3) If is animated then return still image data i.e. first frame of gif. + + try: + im = Image.open(io.BytesIO(image_data)) + image_format = im.format + if getattr(im, "n_frames", 1) > 1: + # There are a number of bugs in Pillow which cause results + # in resized images being broken. To work around this we + # only resize under certain conditions to minimize the + # chance of creating ugly images. + should_resize = ( + im.size[0] != im.size[1] # not square + or im.size[0] > MAX_EMOJI_GIF_SIZE # dimensions too large + or len(image_data) > MAX_EMOJI_GIF_FILE_SIZE_BYTES # filesize too large + ) + + # Generate a still image from the first frame. Since + # we're converting the format to PNG anyway, we resize unconditionally. + still_image = im.copy() + still_image.seek(0) + still_image = ImageOps.exif_transpose(still_image) + still_image = ImageOps.fit(still_image, (size, size), Image.Resampling.LANCZOS) + out = io.BytesIO() + still_image.save(out, format="PNG") + still_image_data = out.getvalue() + + if should_resize: + image_data = resize_animated(im, size) + + return image_data, True, still_image_data + else: + # Note that this is essentially duplicated in the + # still_image code path, above. + im = ImageOps.exif_transpose(im) + im = ImageOps.fit(im, (size, size), Image.Resampling.LANCZOS) + out = io.BytesIO() + im.save(out, format=image_format) + return out.getvalue(), False, None + except OSError: + raise BadImageError(_("Could not decode image; did you upload an image file?")) + except DecompressionBombError: + raise BadImageError(_("Image size exceeds limit.")) diff --git a/zerver/lib/upload/base.py b/zerver/lib/upload/base.py index 84592d430c..d5e76c1599 100644 --- a/zerver/lib/upload/base.py +++ b/zerver/lib/upload/base.py @@ -1,29 +1,12 @@ -import io import os import re import unicodedata from datetime import datetime from typing import IO, Any, BinaryIO, Callable, Iterator, List, Optional, Tuple -from django.utils.translation import gettext as _ -from PIL import GifImagePlugin, Image, ImageOps, PngImagePlugin -from PIL.Image import DecompressionBombError - -from zerver.lib.exceptions import ErrorCode, JsonableError from zerver.models import Attachment, Realm, UserProfile from zerver.models.users import is_cross_realm_bot_email -DEFAULT_AVATAR_SIZE = 100 -MEDIUM_AVATAR_SIZE = 500 -DEFAULT_EMOJI_SIZE = 64 - -# These sizes were selected based on looking at the maximum common -# sizes in a library of animated custom emoji, balanced against the -# network cost of very large emoji images. -MAX_EMOJI_GIF_SIZE = 128 -MAX_EMOJI_GIF_FILE_SIZE_BYTES = 128 * 1024 * 1024 # 128 kb - - INLINE_MIME_TYPES = [ "application/pdf", "audio/aac", @@ -66,131 +49,6 @@ def sanitize_name(value: str) -> str: return value -class BadImageError(JsonableError): - code = ErrorCode.BAD_IMAGE - - -def resize_avatar(image_data: bytes, size: int = DEFAULT_AVATAR_SIZE) -> bytes: - try: - im = Image.open(io.BytesIO(image_data)) - im = ImageOps.exif_transpose(im) - im = ImageOps.fit(im, (size, size), Image.Resampling.LANCZOS) - except OSError: - raise BadImageError(_("Could not decode image; did you upload an image file?")) - except DecompressionBombError: - raise BadImageError(_("Image size exceeds limit.")) - out = io.BytesIO() - if im.mode == "CMYK": - im = im.convert("RGB") - im.save(out, format="png") - return out.getvalue() - - -def resize_logo(image_data: bytes) -> bytes: - try: - im = Image.open(io.BytesIO(image_data)) - im = ImageOps.exif_transpose(im) - im.thumbnail((8 * DEFAULT_AVATAR_SIZE, DEFAULT_AVATAR_SIZE), Image.Resampling.LANCZOS) - except OSError: - raise BadImageError(_("Could not decode image; did you upload an image file?")) - except DecompressionBombError: - raise BadImageError(_("Image size exceeds limit.")) - out = io.BytesIO() - if im.mode == "CMYK": - im = im.convert("RGB") - im.save(out, format="png") - return out.getvalue() - - -def resize_animated(im: Image.Image, size: int = DEFAULT_EMOJI_SIZE) -> bytes: - assert im.n_frames > 1 - frames = [] - duration_info = [] - disposals = [] - # If 'loop' info is not set then loop for infinite number of times. - loop = im.info.get("loop", 0) - for frame_num in range(im.n_frames): - im.seek(frame_num) - new_frame = im.copy() - new_frame.paste(im, (0, 0), im.convert("RGBA")) - new_frame = ImageOps.pad(new_frame, (size, size), Image.Resampling.LANCZOS) - frames.append(new_frame) - if im.info.get("duration") is None: # nocoverage - raise BadImageError(_("Corrupt animated image.")) - duration_info.append(im.info["duration"]) - if isinstance(im, GifImagePlugin.GifImageFile): - disposals.append( - im.disposal_method # type: ignore[attr-defined] # private member missing from stubs - ) - elif isinstance(im, PngImagePlugin.PngImageFile): - disposals.append(im.info.get("disposal", PngImagePlugin.Disposal.OP_NONE)) - else: # nocoverage - raise BadImageError(_("Unknown animated image format.")) - out = io.BytesIO() - frames[0].save( - out, - save_all=True, - optimize=False, - format=im.format, - append_images=frames[1:], - duration=duration_info, - disposal=disposals, - loop=loop, - ) - - return out.getvalue() - - -def resize_emoji( - image_data: bytes, size: int = DEFAULT_EMOJI_SIZE -) -> Tuple[bytes, bool, Optional[bytes]]: - # This function returns three values: - # 1) Emoji image data. - # 2) If emoji is gif i.e. animated. - # 3) If is animated then return still image data i.e. first frame of gif. - - try: - im = Image.open(io.BytesIO(image_data)) - image_format = im.format - if getattr(im, "n_frames", 1) > 1: - # There are a number of bugs in Pillow which cause results - # in resized images being broken. To work around this we - # only resize under certain conditions to minimize the - # chance of creating ugly images. - should_resize = ( - im.size[0] != im.size[1] # not square - or im.size[0] > MAX_EMOJI_GIF_SIZE # dimensions too large - or len(image_data) > MAX_EMOJI_GIF_FILE_SIZE_BYTES # filesize too large - ) - - # Generate a still image from the first frame. Since - # we're converting the format to PNG anyway, we resize unconditionally. - still_image = im.copy() - still_image.seek(0) - still_image = ImageOps.exif_transpose(still_image) - still_image = ImageOps.fit(still_image, (size, size), Image.Resampling.LANCZOS) - out = io.BytesIO() - still_image.save(out, format="PNG") - still_image_data = out.getvalue() - - if should_resize: - image_data = resize_animated(im, size) - - return image_data, True, still_image_data - else: - # Note that this is essentially duplicated in the - # still_image code path, above. - im = ImageOps.exif_transpose(im) - im = ImageOps.fit(im, (size, size), Image.Resampling.LANCZOS) - out = io.BytesIO() - im.save(out, format=image_format) - return out.getvalue(), False, None - except OSError: - raise BadImageError(_("Could not decode image; did you upload an image file?")) - except DecompressionBombError: - raise BadImageError(_("Image size exceeds limit.")) - - class ZulipUploadBackend: # Message attachment uploads def get_public_upload_root_url(self) -> str: diff --git a/zerver/lib/upload/local.py b/zerver/lib/upload/local.py index 0a18ef912a..9e8d0c59f5 100644 --- a/zerver/lib/upload/local.py +++ b/zerver/lib/upload/local.py @@ -10,16 +10,9 @@ from django.conf import settings from typing_extensions import override from zerver.lib.avatar_hash import user_avatar_path +from zerver.lib.thumbnail import MEDIUM_AVATAR_SIZE, resize_avatar, resize_emoji, resize_logo from zerver.lib.timestamp import timestamp_to_datetime -from zerver.lib.upload.base import ( - MEDIUM_AVATAR_SIZE, - ZulipUploadBackend, - create_attachment, - resize_avatar, - resize_emoji, - resize_logo, - sanitize_name, -) +from zerver.lib.upload.base import ZulipUploadBackend, create_attachment, sanitize_name from zerver.lib.utils import assert_is_not_none from zerver.models import Realm, RealmEmoji, UserProfile diff --git a/zerver/lib/upload/s3.py b/zerver/lib/upload/s3.py index 98ed7de301..185dc4a270 100644 --- a/zerver/lib/upload/s3.py +++ b/zerver/lib/upload/s3.py @@ -14,14 +14,11 @@ from typing_extensions import override from zerver.lib.avatar_hash import user_avatar_path from zerver.lib.mime_types import guess_type +from zerver.lib.thumbnail import MEDIUM_AVATAR_SIZE, resize_avatar, resize_emoji, resize_logo from zerver.lib.upload.base import ( INLINE_MIME_TYPES, - MEDIUM_AVATAR_SIZE, ZulipUploadBackend, create_attachment, - resize_avatar, - resize_emoji, - resize_logo, sanitize_name, ) from zerver.models import Realm, RealmEmoji, UserProfile diff --git a/zerver/tests/test_auth_backends.py b/zerver/tests/test_auth_backends.py index 050d72edcd..ffd753c4b0 100644 --- a/zerver/tests/test_auth_backends.py +++ b/zerver/tests/test_auth_backends.py @@ -92,8 +92,8 @@ from zerver.lib.test_helpers import ( read_test_image_file, use_s3_backend, ) +from zerver.lib.thumbnail import DEFAULT_AVATAR_SIZE, MEDIUM_AVATAR_SIZE, resize_avatar from zerver.lib.types import Validator -from zerver.lib.upload.base import DEFAULT_AVATAR_SIZE, MEDIUM_AVATAR_SIZE, resize_avatar from zerver.lib.user_groups import is_user_in_group from zerver.lib.users import get_all_api_keys, get_api_key, get_users_for_api from zerver.lib.utils import assert_is_not_none diff --git a/zerver/tests/test_realm_emoji.py b/zerver/tests/test_realm_emoji.py index f266284fc1..01a7aab44b 100644 --- a/zerver/tests/test_realm_emoji.py +++ b/zerver/tests/test_realm_emoji.py @@ -8,7 +8,7 @@ from zerver.actions.users import do_change_user_role from zerver.lib.exceptions import JsonableError from zerver.lib.test_classes import ZulipTestCase from zerver.lib.test_helpers import get_test_image_file -from zerver.lib.upload.base import BadImageError +from zerver.lib.thumbnail import BadImageError from zerver.models import Realm, RealmEmoji, UserProfile from zerver.models.realms import CommonPolicyEnum, get_realm diff --git a/zerver/tests/test_transfer.py b/zerver/tests/test_transfer.py index e1bc8cf709..bdf10e12be 100644 --- a/zerver/tests/test_transfer.py +++ b/zerver/tests/test_transfer.py @@ -13,6 +13,7 @@ from zerver.lib.test_helpers import ( get_test_image_file, read_test_image_file, ) +from zerver.lib.thumbnail import resize_emoji from zerver.lib.transfer import ( transfer_avatars_to_s3, transfer_emoji_to_s3, @@ -20,7 +21,6 @@ from zerver.lib.transfer import ( transfer_uploads_to_s3, ) from zerver.lib.upload import upload_message_attachment -from zerver.lib.upload.base import resize_emoji from zerver.models import Attachment, RealmEmoji diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index df874d26f6..4b10ec341d 100644 --- a/zerver/tests/test_upload.py +++ b/zerver/tests/test_upload.py @@ -37,8 +37,9 @@ from zerver.lib.test_helpers import ( ratelimit_rule, read_test_image_file, ) +from zerver.lib.thumbnail import BadImageError, resize_emoji from zerver.lib.upload import upload_message_attachment -from zerver.lib.upload.base import BadImageError, ZulipUploadBackend, resize_emoji, sanitize_name +from zerver.lib.upload.base import ZulipUploadBackend, sanitize_name from zerver.lib.upload.local import LocalUploadBackend from zerver.lib.upload.s3 import S3UploadBackend from zerver.lib.users import get_api_key @@ -1422,15 +1423,15 @@ class EmojiTest(UploadSerializeMixin, ZulipTestCase): animated_large_img_data = read_test_image_file(f"animated_large_img.{img_format}") # Test an image larger than max is resized - with patch("zerver.lib.upload.base.MAX_EMOJI_GIF_SIZE", 128): + with patch("zerver.lib.thumbnail.MAX_EMOJI_GIF_SIZE", 128): test_resize() # Test an image file larger than max is resized - with patch("zerver.lib.upload.base.MAX_EMOJI_GIF_FILE_SIZE_BYTES", 3 * 1024 * 1024): + with patch("zerver.lib.thumbnail.MAX_EMOJI_GIF_FILE_SIZE_BYTES", 3 * 1024 * 1024): test_resize() # Test an image smaller than max and smaller than file size max is not resized - with patch("zerver.lib.upload.base.MAX_EMOJI_GIF_SIZE", 512): + with patch("zerver.lib.thumbnail.MAX_EMOJI_GIF_SIZE", 512): test_resize(size=256) # Test a non-animated GIF image which does need to be resized diff --git a/zerver/tests/test_upload_local.py b/zerver/tests/test_upload_local.py index 156524ddf5..39d027550c 100644 --- a/zerver/tests/test_upload_local.py +++ b/zerver/tests/test_upload_local.py @@ -10,6 +10,7 @@ import zerver.lib.upload from zerver.lib.avatar_hash import user_avatar_path from zerver.lib.test_classes import UploadSerializeMixin, ZulipTestCase from zerver.lib.test_helpers import get_test_image_file, read_test_image_file +from zerver.lib.thumbnail import DEFAULT_EMOJI_SIZE, MEDIUM_AVATAR_SIZE, resize_avatar from zerver.lib.upload import ( all_message_attachments, delete_export_tarball, @@ -20,7 +21,6 @@ from zerver.lib.upload import ( upload_export_tarball, upload_message_attachment, ) -from zerver.lib.upload.base import DEFAULT_EMOJI_SIZE, MEDIUM_AVATAR_SIZE, resize_avatar from zerver.lib.upload.local import write_local_file from zerver.models import Attachment, RealmEmoji from zerver.models.realms import get_realm diff --git a/zerver/tests/test_upload_s3.py b/zerver/tests/test_upload_s3.py index 76e032a0a2..7474072e71 100644 --- a/zerver/tests/test_upload_s3.py +++ b/zerver/tests/test_upload_s3.py @@ -20,6 +20,12 @@ from zerver.lib.test_helpers import ( read_test_image_file, use_s3_backend, ) +from zerver.lib.thumbnail import ( + DEFAULT_AVATAR_SIZE, + DEFAULT_EMOJI_SIZE, + MEDIUM_AVATAR_SIZE, + resize_avatar, +) from zerver.lib.upload import ( all_message_attachments, delete_export_tarball, @@ -29,12 +35,6 @@ from zerver.lib.upload import ( upload_export_tarball, upload_message_attachment, ) -from zerver.lib.upload.base import ( - DEFAULT_AVATAR_SIZE, - DEFAULT_EMOJI_SIZE, - MEDIUM_AVATAR_SIZE, - resize_avatar, -) from zerver.lib.upload.s3 import S3UploadBackend from zerver.models import Attachment, RealmEmoji, UserProfile from zerver.models.realms import get_realm