diff --git a/routes/gallery/gallery_routes.py b/routes/gallery/gallery_routes.py index f3aa63e5e..fd1244b9a 100644 --- a/routes/gallery/gallery_routes.py +++ b/routes/gallery/gallery_routes.py @@ -2415,11 +2415,10 @@ def setup_gallery_routes() -> APIRouter: image_b64 = body.get("image") amount = body.get("amount", 50) / 100.0 - from PIL import Image, ImageFilter + from PIL import ImageFilter import base64, io - img_bytes = _b64_bytes(image_b64) - img = Image.open(io.BytesIO(img_bytes)).convert("RGB") + img = _b64_to_pil_image(image_b64, mode="RGB") # Unsharp mask: radius=2, percent=amount*200, threshold=3 sharpened = img.filter(ImageFilter.UnsharpMask(radius=2, percent=int(amount * 200), threshold=3)) @@ -2454,8 +2453,7 @@ def setup_gallery_routes() -> APIRouter: except ImportError: raise HTTPException(500, "Server missing a required dependency") # Decode source image (RGB; Real-ESRGAN doesn't preserve alpha). - img_bytes = _b64_bytes(image_b64) - src = Image.open(io.BytesIO(img_bytes)).convert("RGB") + src = _b64_to_pil_image(image_b64, mode="RGB") try: patch_realesrgan_torchvision_compat() from realesrgan import RealESRGANer @@ -2507,8 +2505,7 @@ def setup_gallery_routes() -> APIRouter: import numpy as np except ImportError: raise HTTPException(500, "Server missing a required dependency") - img_bytes = _b64_bytes(image_b64) - src = Image.open(io.BytesIO(img_bytes)).convert("RGB") + src = _b64_to_pil_image(image_b64, mode="RGB") try: patch_realesrgan_torchvision_compat() from basicsr.archs.rrdbnet_arch import RRDBNet @@ -2760,8 +2757,7 @@ def setup_gallery_routes() -> APIRouter: from PIL import Image import base64, io - img_bytes = _b64_bytes(image_b64) - img = Image.open(io.BytesIO(img_bytes)).convert("RGBA") + img = _b64_to_pil_image(image_b64) W, H = img.size hint = None @@ -2850,8 +2846,7 @@ def setup_gallery_routes() -> APIRouter: from PIL import Image, ImageFilter, ImageEnhance import numpy as np - img_bytes = _b64_bytes(image_b64) - img = Image.open(io.BytesIO(img_bytes)).convert("RGB") + img = _b64_to_pil_image(image_b64, mode="RGB") # Try GFPGAN first (AI face restoration) try: diff --git a/tests/test_gallery_image_bad_input.py b/tests/test_gallery_image_bad_input.py new file mode 100644 index 000000000..7762ac64e --- /dev/null +++ b/tests/test_gallery_image_bad_input.py @@ -0,0 +1,44 @@ +import asyncio +import base64 +import io + +import pytest +from fastapi import HTTPException +from PIL import Image + +import routes.gallery_routes as gallery + +ROUTES = ['sharpen', 'remove-bg', 'denoise', 'upscale-local', 'enhance-face'] + + +def route(name): + return next(r for r in gallery.setup_gallery_routes().routes if r.path == f'/api/image/{name}') + + +def call(name, monkeypatch, body): + monkeypatch.setattr(gallery, 'require_privilege', lambda *args: 'alice') + + class Request: + async def json(self): + return body + + return asyncio.run(route(name).endpoint(Request())) + + +@pytest.mark.parametrize('name', ROUTES) +@pytest.mark.parametrize('body', [ + {}, + {'image': '!!notb64'}, + {'image': base64.b64encode(b'hello').decode()}, +]) +def test_bad_image_is_a_client_error(name, body, monkeypatch): + with pytest.raises(HTTPException) as exc: + call(name, monkeypatch, body) + assert exc.value.status_code == 400 + + +def test_sharpen_still_accepts_a_valid_image(monkeypatch): + output = io.BytesIO() + Image.new('RGB', (4, 4), 'red').save(output, format='PNG') + result = call('sharpen', monkeypatch, {'image': base64.b64encode(output.getvalue()).decode()}) + assert Image.open(io.BytesIO(base64.b64decode(result['image']))).size == (4, 4)