From 2c1804e957fe90037890b1d6da99c745dce049e8 Mon Sep 17 00:00:00 2001
From: Krystle Salazar <krystle.salazar@automattic.com>
Date: Wed, 5 Oct 2022 17:45:49 -0400
Subject: [PATCH] Catch struct errors in `watermark/` endpoint (#932)

---
 api/catalog/api/utils/watermark.py    | 14 ++++++++------
 api/env.docker                        |  3 +++
 api/test/unit/utils/watermark_test.py | 22 +++++++++++++++++++++-
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/api/catalog/api/utils/watermark.py b/api/catalog/api/utils/watermark.py
index 9e2c56db2..d8a70e5f4 100644
--- a/api/catalog/api/utils/watermark.py
+++ b/api/catalog/api/utils/watermark.py
@@ -1,5 +1,6 @@
 import logging
 import os
+import struct
 from enum import Flag, auto
 from io import BytesIO
 from textwrap import wrap
@@ -158,17 +159,18 @@ def _open_image(url):
         response = requests.get(url, headers=HEADERS)
         img_bytes = BytesIO(response.content)
         img = Image.open(img_bytes)
-        # Preserve EXIF metadata
-        if "exif" in img.info:
-            exif = piexif.load(img.info["exif"])
-        else:
-            exif = None
-        return img, exif
     except requests.exceptions.RequestException as e:
         capture_exception(e)
         logger.error(f"Error loading image data: {e}")
         return None, None
 
+    try:
+        # Preserve EXIF metadata
+        exif = piexif.load(img.info["exif"]) if "exif" in img.info else None
+        return img, exif
+    except struct.error:
+        return img, None
+
 
 def _print_attribution_on_image(img, image_info):
     """
diff --git a/api/env.docker b/api/env.docker
index 65db473de..2fd6629d7 100644
--- a/api/env.docker
+++ b/api/env.docker
@@ -4,6 +4,7 @@ DJANGO_SETTINGS_MODULE=catalog.settings
 DJANGO_SECRET_KEY="ny#b__$$f6ry4wy8oxre97&-68u_0lk3gw(z=d40_dxey3zw0v1"
 DJANGO_DEBUG_ENABLED=True
 
+ENVIRONMENT=development
 ALLOWED_HOSTS=api.openverse.engineering,api-dev.openverse.engineering,host.docker.internal
 
 REDIS_HOST=cache
@@ -18,3 +19,5 @@ UPSTREAM_DATABASE_PORT=5432
 SEMANTIC_VERSION=1.0.0
 
 ELASTICSEARCH_URL=es
+
+WATERMARK_ENABLED=True
diff --git a/api/test/unit/utils/watermark_test.py b/api/test/unit/utils/watermark_test.py
index 3c0e8c50b..10ca2ff8c 100644
--- a/api/test/unit/utils/watermark_test.py
+++ b/api/test/unit/utils/watermark_test.py
@@ -1,12 +1,16 @@
 import json
+import struct
 from dataclasses import dataclass
+from io import BytesIO
 from pathlib import Path
 from typing import Callable
+from unittest import mock
 
 import pytest
+from PIL import Image
 from requests import Request, Response
 
-from catalog.api.utils.watermark import HEADERS, watermark
+from catalog.api.utils.watermark import HEADERS, _open_image, watermark
 
 
 _MOCK_IMAGE_PATH = Path(__file__).parent / ".." / ".." / "factory"
@@ -51,3 +55,19 @@ def test_sends_UA_header(requests):
     assert len(requests.requests) > 0
     for r in requests.requests:
         assert r.headers == HEADERS
+
+
+def test_catch_struct_errors_from_piexif(requests):
+    img_mock = Image.open(BytesIO(_MOCK_IMAGE_BYTES))
+    img_mock.info["exif"] = "bad_info"
+
+    with mock.patch("PIL.Image.open") as open_mock, mock.patch(
+        "piexif.load"
+    ) as load_mock:
+        open_mock.return_value = img_mock
+        load_mock.side_effect = struct.error("unpack requires a buffer of 2 bytes")
+
+        img, exif = _open_image("http://example.com/")
+
+    assert img is not None
+    assert exif is None