From 0639e8f90133c79a222616cc84876b8b330bd223 Mon Sep 17 00:00:00 2001 From: Aron Carroll Date: Wed, 16 Oct 2024 11:37:38 +0100 Subject: [PATCH] Validate url passed to URLFile conforms to HTTP protocol --- python/cog/types.py | 7 +++++++ python/tests/server/test_clients.py | 2 -- python/tests/test_types.py | 8 ++++++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/python/cog/types.py b/python/cog/types.py index 0bcb0afa3..d982439a8 100644 --- a/python/cog/types.py +++ b/python/cog/types.py @@ -205,6 +205,13 @@ class URLFile(io.IOBase): def __init__(self, url: str) -> None: parsed = urllib.parse.urlparse(url) + if parsed.scheme not in { + "http", + "https", + }: + raise ValueError( + "URLFile requires URL to conform to HTTP or HTTPS protocol" + ) object.__setattr__(self, "name", os.path.basename(parsed.path)) object.__setattr__(self, "__url__", url) diff --git a/python/tests/server/test_clients.py b/python/tests/server/test_clients.py index 506f58b89..e474d5178 100644 --- a/python/tests/server/test_clients.py +++ b/python/tests/server/test_clients.py @@ -10,8 +10,6 @@ import pytest from cog.server.clients import ClientManager -pytest.mark.asyncio - @pytest.mark.asyncio async def test_upload_files_without_url(): diff --git a/python/tests/test_types.py b/python/tests/test_types.py index f9563fa42..f6b5a85fa 100644 --- a/python/tests/test_types.py +++ b/python/tests/test_types.py @@ -19,6 +19,14 @@ def file_fixture(body: str): ) +def test_urlfile_protocol_validation(): + with pytest.raises(ValueError): + URLFile("file:///etc/shadow") + + with pytest.raises(ValueError): + URLFile("data:text/plain,hello") + + @mock.patch("urllib.request.urlopen", return_value=file_fixture("hello world")) def test_urlfile_acts_like_response(mock_urlopen): u = URLFile("https://example.com/some/url")