Skip to content
This repository has been archived by the owner on Feb 22, 2023. It is now read-only.

Catch struct errors in watermark/ endpoint #932

Merged
merged 3 commits into from
Oct 5, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions api/catalog/api/utils/watermark.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
import os
import struct
from enum import Flag, auto
from io import BytesIO
from textwrap import wrap
Expand Down Expand Up @@ -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):
"""
Expand Down
3 changes: 3 additions & 0 deletions api/env.docker
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -18,3 +19,5 @@ UPSTREAM_DATABASE_PORT=5432
SEMANTIC_VERSION=1.0.0

ELASTICSEARCH_URL=es

WATERMARK_ENABLED=True
22 changes: 21 additions & 1 deletion api/test/unit/utils/watermark_test.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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