Add a default User-Agent header when downloading images using pycurl (#…

…932)

* Add a default 'User-Agent: Marqobot/1.0' header to all request to download images
* Follow redirections when downloading images

src/marqo/s2_inference/clip_utils.py

@@ -32,6 +32,7 @@
 OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
 OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)
 BICUBIC = InterpolationMode.BICUBIC
+DEFAULT_HEADERS = {'User-Agent': 'Marqobot/1.0'}
 
 
 def get_allowed_image_types():
@@ -153,16 +154,20 @@ def download_image_from_url(image_path: str, image_download_headers: dict, timeo
                                  f"The url could not be encoded properly. Original error: {e}")
     buffer = BytesIO()
     c = pycurl.Curl()
-    c.setopt(c.CAINFO, certifi.where())
-    c.setopt(c.URL, encoded_url)
-    c.setopt(c.WRITEDATA, buffer)
-    c.setopt(c.TIMEOUT_MS, timeout_ms)
-    c.setopt(c.HTTPHEADER, [f"{k}: {v}" for k, v in image_download_headers.items()])
+    c.setopt(pycurl.CAINFO, certifi.where())
+    c.setopt(pycurl.URL, encoded_url)
+    c.setopt(pycurl.WRITEDATA, buffer)
+    c.setopt(pycurl.TIMEOUT_MS, timeout_ms)
+    c.setopt(pycurl.FOLLOWLOCATION, 1)
+
+    headers = DEFAULT_HEADERS.copy()
+    headers.update(image_download_headers)
+    c.setopt(pycurl.HTTPHEADER, [f"{k}: {v}" for k, v in headers.items()])
 
     try:
         c.perform()
-        if c.getinfo(c.RESPONSE_CODE) != 200:
-            raise ImageDownloadError(f"image url `{image_path}` returned {c.getinfo(c.RESPONSE_CODE)}")
+        if c.getinfo(pycurl.RESPONSE_CODE) != 200:
+            raise ImageDownloadError(f"image url `{image_path}` returned {c.getinfo(pycurl.RESPONSE_CODE)}")
     except pycurl.error as e:
         raise ImageDownloadError(f"Marqo encountered an error when downloading the image url {image_path}. "
                                  f"The original error is: {e}")

src/marqo/version.py

@@ -1,4 +1,4 @@
-__version__ = "2.11.0"
+__version__ = "2.11.1"
 
 def get_version() -> str:
     return f"{__version__}"

tests/marqo_test.py

@@ -1,9 +1,15 @@
+import contextlib
+import socket
+import threading
 import time
 import unittest
 import uuid
+from typing import Generator
 from unittest.mock import patch, Mock
 
+import uvicorn
 import vespa.application as pyvespa
+from starlette.applications import Starlette
 
 from marqo import config, version
 from marqo.vespa.zookeeper_client import ZookeeperClient
@@ -307,3 +313,38 @@ def assertRaisesStrict(self, expected_exception):
 
 class AsyncMarqoTestCase(unittest.IsolatedAsyncioTestCase, MarqoTestCase):
     pass
+
+
+class MockHttpServer:
+    """
+    A MockHttpServer that takes a Starlette app as input, start the uvicorn server
+    in a thread, and yield the server url (with random port binding). After the test,
+    it automatically shuts down the server.
+
+    This can be used in individual tests, or as a test fixture in class or module scope.
+    Example usage:
+
+    app = Starlette(routes=[
+        Route('/path1', lambda _: Response({"a":"b"}, status_code=200)),
+        Route('/image.jpg', lambda _: Response(b'\x00\x00\x00\xff', media_type='image/png')),
+    ])
+
+    with MockHttpServer(app).run_in_thread() as base_url:
+        run_some_tests
+    """
+    def __init__(self, app: Starlette):
+        self.server = uvicorn.Server(config=uvicorn.Config(app=app))
+
+    @contextlib.contextmanager
+    def run_in_thread(self) -> Generator[str, None, None]:
+        (sock := socket.socket()).bind(("127.0.0.1", 0))
+        thread = threading.Thread(target=self.server.run, kwargs={"sockets": [sock]})
+        thread.start()
+        try:
+            while not self.server.started:
+                time.sleep(1)
+            address, port = sock.getsockname()
+            yield f'http://{address}:{port}'
+        finally:
+            self.server.should_exit = True
+            thread.join()

tests/s2_inference/test_image_downloading.py

@@ -1,9 +1,15 @@
 from unittest import TestCase
+from unittest.mock import patch, MagicMock
 
+import pycurl
 import pytest
+from starlette.applications import Starlette
+from starlette.responses import Response
+from starlette.routing import Route
 
 from marqo.s2_inference.clip_utils import encode_url, download_image_from_url
 from marqo.s2_inference.errors import ImageDownloadError
+from tests.marqo_test import MockHttpServer
 
 
 @pytest.mark.unittest
@@ -47,4 +53,40 @@ def test_download_image_from_url_handleDifferentUrlsCorrectly(self):
         for url, expected, msg in self.test_cases:
             with self.subTest(url=url, expected=expected, msg=msg):
                 with self.assertRaises(ImageDownloadError) as cm:
-                    download_image_from_url(image_path=url + ".jpg", image_download_headers={})
+                    download_image_from_url(image_path=url + ".jpg", image_download_headers={})
+
+    def test_download_image_from_url_handlesUrlRequiringUserAgentHeader(self):
+        url_requiring_user_agent_header = "https://docs.marqo.ai/2.0.0/Examples/marqo.jpg"
+        try:
+            download_image_from_url(image_path=url_requiring_user_agent_header, image_download_headers={})
+        except Exception as e:
+            self.fail(f"Exception was raised when downloading {url_requiring_user_agent_header}: {e}")
+
+    @patch('pycurl.Curl')
+    def test_download_image_from_url_mergesDefaultHeadersWithCustomHeaders(self, mock_curl):
+        mock_curl_instance = mock_curl.return_value
+        mock_curl_instance.setopt = MagicMock()
+        mock_curl_instance.perform = MagicMock()
+        mock_curl_instance.getinfo = MagicMock(return_value=200)
+
+        test_cases = [
+            ({}, ['User-Agent: Marqobot/1.0'], "Empty header"),
+            ({'a': 'b'}, ['User-Agent: Marqobot/1.0', 'a: b'], "Basic headers"),
+            ({'User-Agent': 'Marqobot-Image/1.0'}, ['User-Agent: Marqobot-Image/1.0'], "Headers with override"),
+        ]
+
+        for (headers, expected_headers, msg) in test_cases:
+            with self.subTest(headers=headers, expected_headers=expected_headers, msg=msg):
+                download_image_from_url('http://example.com/image.jpg', image_download_headers=headers)
+                mock_curl_instance.setopt.assert_called_with(pycurl.HTTPHEADER, expected_headers)
+
+    def test_download_image_from_url_handlesRedirection(self):
+        image_content = b'\x00\x00\x00\xff'
+        app = Starlette(routes=[
+            Route('/missing_image.jpg', lambda _: Response(status_code=301, headers={'Location': '/image.jpg'})),
+            Route('/image.jpg', lambda _: Response(image_content, media_type='image/png')),
+        ])
+
+        with MockHttpServer(app).run_in_thread() as base_url:
+            result = download_image_from_url(f'{base_url}/missing_image.jpg', image_download_headers={})
+            self.assertEqual(result.getvalue(), image_content)