Merge pull request #638 from adamtheturtle/bad-image

Validate query image is not corrupted

Author: adamtheturtle

src/mock_vws/_mock_web_query_api.py

@@ -83,6 +83,60 @@ def validate_image_format(
     )
 
 
+@wrapt.decorator
+def validate_image_file_contents(
+    wrapped: Callable[..., str],
+    instance: Any,  # pylint: disable=unused-argument
+    args: Tuple[_RequestObjectProxy, _Context],
+    kwargs: Dict,
+) -> str:
+    """
+    Validate the format of the image given to the query endpoint.
+
+    Args:
+        wrapped: An endpoint function for `requests_mock`.
+        instance: The class that the endpoint function is in.
+        args: The arguments given to the endpoint function.
+        kwargs: The keyword arguments given to the endpoint function.
+
+    Returns:
+        The result of calling the endpoint.
+        An `UNPROCESSABLE_ENTITY` response if the image is given and is not
+        either a PNG or a JPEG.
+    """
+    request, context = args
+    body_file = io.BytesIO(request.body)
+
+    _, pdict = cgi.parse_header(request.headers['Content-Type'])
+    parsed = cgi.parse_multipart(
+        fp=body_file,
+        pdict={
+            'boundary': pdict['boundary'].encode(),
+        },
+    )
+
+    [image] = parsed['image']
+
+    image_file = io.BytesIO(image)
+    try:
+        Image.open(image_file).verify()
+    except SyntaxError:
+        context.status_code = codes.UNPROCESSABLE_ENTITY
+        transaction_id = uuid.uuid4().hex
+        result_code = ResultCodes.BAD_IMAGE.value
+
+        # The response has an unusual format of separators, so we construct it
+        # manually.
+        return (
+            '{"transaction_id": '
+            f'"{transaction_id}",'
+            f'"result_code":"{result_code}"'
+            '}'
+        )
+
+    return wrapped(*args, **kwargs)
+
+
 @wrapt.decorator
 def validate_date_header_given(
     wrapped: Callable[..., str],
@@ -476,6 +530,7 @@ def decorator(method: Callable[..., str]) -> Callable[..., str]:
             validate_date_header_given,
             validate_include_target_data,
             validate_max_num_results,
+            validate_image_file_contents,
             validate_image_format,
             validate_image_field_given,
             validate_extra_fields,

tests/mock_vws/test_query.py

@@ -1054,6 +1054,51 @@ def test_inactive(
         assert response.json()['results'] == []
 
 
+@pytest.mark.usefixtures('verify_mock_vuforia')
+class TestBadImage:
+    """
+    Tests for bad images.
+    """
+
+    def test_corrupted(
+        self,
+        vuforia_database_keys: VuforiaDatabaseKeys,
+        png_rgb: io.BytesIO,
+    ) -> None:
+        """
+        A "BadImage" result is returned when a corrupted image is given.
+        """
+        original_data = png_rgb.getvalue()
+        corrupted_data = original_data.replace(b'IEND', b'\x00' + b'IEND')
+
+        body = {'image': ('image.jpeg', corrupted_data, 'image/jpeg')}
+
+        response = query(
+            vuforia_database_keys=vuforia_database_keys,
+            body=body,
+        )
+
+        assert_vwq_failure(
+            response=response,
+            status_code=codes.UNPROCESSABLE_ENTITY,
+            content_type='application/json',
+        )
+        assert response.json().keys() == {'transaction_id', 'result_code'}
+        assert_valid_transaction_id(response=response)
+        assert_valid_date_header(response=response)
+        result_code = response.json()['result_code']
+        transaction_id = response.json()['transaction_id']
+        assert result_code == ResultCodes.BAD_IMAGE.value
+        # The separators are inconsistent and we test this.
+        expected_text = (
+            '{"transaction_id": '
+            f'"{transaction_id}",'
+            f'"result_code":"{result_code}"'
+            '}'
+        )
+        assert response.text == expected_text
+
+
 @pytest.mark.usefixtures('verify_mock_vuforia')
 class TestMaximumImageSize:
     """