move image resizing to a subprocess

cogs/emote.py

@@ -1,9 +1,8 @@
 #!/usr/bin/env python3
 # encoding: utf-8
 
-import io
-import imghdr
 import asyncio
+import cgi
 import logging
 import weakref
 import traceback
@@ -186,29 +185,31 @@ class Emotes(commands.Cog):
 		return emote
 
 	async def fetch_emote(self, url):
-		# credits to @Liara#0001 (ID 136900814408122368) for most of this part
+		def validate_headers(response):
-		# https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L217-228
-		async with self.http.head(url, timeout=5) as response:
 			if response.reason != 'OK':
 				raise errors.HTTPException(response.status)
-			if response.headers.get('Content-Type') not in ('image/png', 'image/jpeg', 'image/gif'):
+			# some dumb servers also send '; charset=UTF-8' which we should ignore
+			mimetype, options = cgi.parse_header(response.headers.get('Content-Type', ''))
+			if mimetype not in {'image/png', 'image/jpeg', 'image/gif'}:
 				raise errors.InvalidImageError
 
-		async with self.http.get(url) as response:
+		try:
-			if response.reason != 'OK':
+			async with self.http.head(url, timeout=5) as response:
-				raise errors.HTTPException(response.status)
+				validate_headers(response)
-			return io.BytesIO(await response.read())
+		except aiohttp.ServerDisconnectedError as exception:
+			validate_headers(exception.message)
 
-	async def create_emote_from_bytes(self, guild, name, author_id, image_data: io.BytesIO, *, reason=None):
+		async with self.http.get(url) as response:
-		# resize_until_small is normally blocking, because wand is.
+			validate_headers(response)
-		# run_in_executor is magic that makes it non blocking somehow.
+			return await response.read()
-		# also, None as the executor arg means "use the loop's default executor"
+
-		image_data = await self.bot.loop.run_in_executor(None, utils.image.resize_until_small, image_data)
+	async def create_emote_from_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None):
+		image_data = await utils.image.resize_in_subprocess(image_data)
 		if reason is None:
 			reason = f'Created by {utils.format_user(self.bot, author_id)}'
 		return await guild.create_custom_emoji(
 			name=name,
-			image=image_data.read(),
+			image=image_data,
 			reason=reason)
 
 	@commands.command(aliases=('delete', 'delet', 'rm'))

utils/image.py

@@ -1,18 +1,27 @@
+#!/usr/bin/env python3
+# encoding: utf-8
+
+import asyncio
+import base64
+import contextlib
 import io
 import logging
-import contextlib
+import sys
+import typing
 
 logger = logging.getLogger(__name__)
 
 try:
-	from wand.image import Image
+	import wand.image
-except ImportError:
+except (ImportError, OSError):
 	logger.warn('Failed to import wand.image. Image manipulation functions will be unavailable.')
+else:
+	import wand.exceptions
 
 from utils import errors
 
-def resize_until_small(image_data: io.BytesIO):
+def resize_until_small(image_data: io.BytesIO) -> None:
-	"""If the image_data is bigger than 256KB, resize it until it's not"""
+	"""If the image_data is bigger than 256KB, resize it until it's not."""
 	# It's important that we only attempt to resize the image when we have to,
 	# ie when it exceeds the Discord limit of 256KiB.
 	# Apparently some <256KiB images become larger when we attempt to resize them,
@@ -21,43 +30,108 @@ def resize_until_small(image_data: io.BytesIO):
 	image_size = size(image_data)
 	while image_size > 256 * 2**10 and max_resolution >= 32:  # don't resize past 32x32 or 256KiB
 		logger.debug('image size too big (%s bytes)', image_size)
-		logger.debug('attempting resize to %s*%s pixels', max_resolution, max_resolution)
+		logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution)
-		image_data = thumbnail(image_data, (max_resolution, max_resolution))
+
+		try:
+			thumbnail(image_data, (max_resolution, max_resolution))
+		except wand.exceptions.CoderError:
+			raise errors.InvalidImageError
+
 		image_size = size(image_data)
 		max_resolution //= 2
-	return image_data
 
-def size(data: io.BytesIO):
+def thumbnail(image_data: io.BytesIO, max_size=(128, 128)) -> None:
-	"""return the size, in bytes, of the data a BytesIO object represents"""
-	with preserve_position(data):
-		data.seek(0, io.SEEK_END)
-		return data.tell()
-
-def thumbnail(image_data: io.BytesIO, max_size=(128, 128)):
 	"""Resize an image in place to no more than max_size pixels, preserving aspect ratio."""
-	# Credit to @Liara#0001 (ID 136900814408122368)
+	with wand.image.Image(blob=image_data) as image:
-	# https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L243-247
+		new_resolution = scale_resolution((image.width, image.height), max_size)
-	image = Image(blob=image_data)
+		image.resize(*new_resolution)
-	image.resize(*scale_resolution((image.width, image.height), max_size))
+		image_data.truncate(0)
-	# we create a new buffer here because there's wand errors otherwise.
+		image_data.seek(0)
-	# specific error:
+		image.save(file=image_data)
-	# MissingDelegateError: no decode delegate for this image format `' @ error/blob.c/BlobToImage/353
+
-	out = io.BytesIO()
+	# allow resizing the original image more than once for memory profiling
-	image.save(file=out)
+	image_data.seek(0)
-	out.seek(0)
-	return out
 
 def scale_resolution(old_res, new_res):
-	# https://stackoverflow.com/a/6565988
 	"""Resize a resolution, preserving aspect ratio. Returned w,h will be <= new_res"""
+	# https://stackoverflow.com/a/6565988
+
 	old_width, old_height = old_res
 	new_width, new_height = new_res
+
 	old_ratio = old_width / old_height
 	new_ratio = new_width / new_height
 	if new_ratio > old_ratio:
 		return (old_width * new_height//old_height, new_height)
 	return new_width, old_height * new_width//old_width
 
+def mime_type_for_image(data):
+	if data.startswith(b'\x89PNG\r\n\x1a\n'):
+		return 'image/png'
+	elif data.startswith(b'\xFF\xD8') and data.rstrip(b'\0').endswith(b'\xFF\xD9'):
+		return 'image/jpeg'
+	elif data.startswith((b'GIF87a', b'GIF89a')):
+		return 'image/gif'
+	elif data.startswith(b'RIFF') and data[8:12] == b'WEBP':
+		return 'image/webp'
+	else:
+		raise errors.InvalidImageError
+
+def image_to_base64_url(data):
+	fmt = 'data:{mime};base64,{data}'
+	mime = mime_type_for_image(data)
+	b64 = base64.b64encode(data).decode('ascii')
+	return fmt.format(mime=mime, data=b64)
+
+def main() -> typing.NoReturn:
+	"""resize an image from stdin and write the resized version to stdout."""
+	import sys
+
+	data = io.BytesIO(sys.stdin.buffer.read())
+	try:
+		resize_until_small(data)
+	except errors.InvalidImageError:
+		# 2 is used because 1 is already used by python's default error handler
+		sys.exit(2)
+
+	stdout_write = sys.stdout.buffer.write  # getattr optimization
+
+	while True:
+		buf = data.read(16 * 1024)
+		if not buf:
+			break
+
+		stdout_write(buf)
+
+	sys.exit(0)
+
+async def resize_in_subprocess(image_data: bytes):
+	proc = await asyncio.create_subprocess_exec(
+		sys.executable, '-m', __name__,
+
+		stdin=asyncio.subprocess.PIPE,
+		stdout=asyncio.subprocess.PIPE,
+		stderr=asyncio.subprocess.PIPE)
+
+	try:
+		image_data, err = await asyncio.wait_for(proc.communicate(image_data), timeout=30)
+	except asyncio.TimeoutError:
+		proc.kill()
+		raise errors.ImageResizeTimeoutError
+	else:
+		if proc.returncode == 2:
+			raise errors.InvalidImageError
+		if proc.returncode != 0:
+			raise RuntimeError(err.decode('utf-8') + f'Return code: {proc.returncode}')
+
+	return image_data
+
+def size(fp):
+	"""return the size, in bytes, of the data a file-like object represents"""
+	with preserve_position(fp):
+		fp.seek(0, io.SEEK_END)
+		return fp.tell()
+
 class preserve_position(contextlib.AbstractContextManager):
 	def __init__(self, fp):
 		self.fp = fp
@@ -65,3 +139,6 @@ class preserve_position(contextlib.AbstractContextManager):
 
 	def __exit__(self, *excinfo):
 		self.fp.seek(self.old_pos)
+
+if __name__ == '__main__':
+	main()