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

Commit

Permalink
Catch struct errors in watermark/ endpoint (#932)
Browse files Browse the repository at this point in the history
  • Loading branch information
krysal authored Oct 5, 2022
1 parent f431b1b commit 2c1804e
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 7 deletions.
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

0 comments on commit 2c1804e

Please sign in to comment.