Validate url passed to URLFile conforms to HTTP protocol

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)
 

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():

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")