From 12ff813412fb2cfd2eb7f1e8b7f15e0c7e1cc30a Mon Sep 17 00:00:00 2001
From: Aron Carroll <aron@replicate.com>
Date: Wed, 16 Oct 2024 11:37:15 +0100
Subject: [PATCH 1/2] Fix type annotations in Input

---
 python/cog/types.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/python/cog/types.py b/python/cog/types.py
index 143a89513..0bcb0afa3 100644
--- a/python/cog/types.py
+++ b/python/cog/types.py
@@ -22,13 +22,13 @@
 
 def Input(
     default: Any = ...,
-    description: str = None,
-    ge: float = None,
-    le: float = None,
-    min_length: int = None,
-    max_length: int = None,
-    regex: str = None,
-    choices: List[Union[str, int]] = None,
+    description: Optional[str] = None,
+    ge: Optional[float] = None,
+    le: Optional[float] = None,
+    min_length: Optional[int] = None,
+    max_length: Optional[int] = None,
+    regex: Optional[str] = None,
+    choices: Optional[List[Union[str, int]]] = None,
 ) -> Any:
     """Input is similar to pydantic.Field, but doesn't require a default value to be the first argument."""
     return Field(

From f6d3bafcb8b6ee952e1b26bd78f7ab5fc5ccfa88 Mon Sep 17 00:00:00 2001
From: Aron Carroll <aron@replicate.com>
Date: Wed, 16 Oct 2024 11:37:38 +0100
Subject: [PATCH 2/2] 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")