From 81197d4676b6b3f302a6b5849264c6931aefeea9 Mon Sep 17 00:00:00 2001 From: Sai Teja Erukude Date: Sun, 8 Nov 2026 23:08:54 -0600 Subject: [PATCH] fix: CVE-2026-56712 ssrf vulnerability --- pyproject.toml | 1 + python/cog/_download.py | 203 +++++++++++++++++++++++++++ python/cog/types.py | 52 +++++-- python/tests/test_download.py | 255 ++++++++++++++++++++++++++++++++++ 4 files changed, 499 insertions(+), 12 deletions(-) create mode 100644 python/cog/_download.py create mode 100644 python/tests/test_download.py diff --git a/pyproject.toml b/pyproject.toml index 2684826c03..81d47906b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "pyyaml>=6.0", "structlog>=21.0.0", "requests>=2.25.0", + "urllib3>=1.26.0,<3", "coglet>=0.1.0,<1.0", ] dynamic = ["version"] diff --git a/python/cog/_download.py b/python/cog/_download.py new file mode 100644 index 0000000000..fbbf8983a9 --- /dev/null +++ b/python/cog/_download.py @@ -0,0 +1,203 @@ +"""Fetch untrusted input URLs without allowing access to internal services.""" + +import ipaddress +import socket +from typing import Any, Iterator, Union, cast +from urllib.parse import urlsplit + +import requests +from requests.adapters import HTTPAdapter +from urllib3.connection import HTTPConnection, HTTPSConnection +from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool +from urllib3.exceptions import NewConnectionError + +IPAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address] + +# Azure's platform virtual IP is globally classified, but exposes internal +# services. Other common metadata endpoints are already non-global addresses. +# https://learn.microsoft.com/azure/virtual-network/what-is-ip-address-168-63-129-16 +_INTERNAL_NETWORKS = ( + ipaddress.ip_network("168.63.129.16/32"), + ipaddress.ip_network("64:ff9b::/96"), # NAT64 may translate to internal IPv4. + ipaddress.ip_network("64:ff9b:1::/48"), + ipaddress.ip_network("192.88.99.0/24"), # Deprecated relay addresses. + ipaddress.ip_network("3fff::/20"), # Documentation prefix added in 2024. + ipaddress.ip_network("fec0::/10"), # Deprecated IPv6 site-local addresses. +) +_PROTOCOL_V4 = ipaddress.ip_network("192.0.0.0/24") +_PROTOCOL_V6 = ipaddress.ip_network("2001::/23") +_GLOBAL_PROTOCOL_V6 = tuple( + ipaddress.ip_network(cidr) + for cidr in ( + "2001:1::1/128", + "2001:1::2/128", + "2001:1::3/128", + "2001:3::/32", + "2001:4:112::/48", + "2001:20::/28", + "2001:30::/28", + ) +) + + +def _is_global(address: IPAddress) -> bool: + # Older supported Python patch releases have stale classifications for + # these IANA protocol-assignment ranges. Keep their exceptions explicit. + # https://www.iana.org/assignments/iana-ipv4-special-registry/ + # https://www.iana.org/assignments/iana-ipv6-special-registry/ + if address in _PROTOCOL_V4: + return str(address) in ("192.0.0.9", "192.0.0.10") + if address in _PROTOCOL_V6: + return any(address in network for network in _GLOBAL_PROTOCOL_V6) + return address.is_global + + +def _validate_url(url: str) -> None: + """Reject ambiguous authorities before Requests normalizes the URL.""" + try: + parsed = urlsplit(url) + if parsed.scheme not in ("http", "https") or not parsed.hostname: + raise ValueError("an HTTP(S) URL with a hostname is required") + if parsed.username is not None or parsed.password is not None: + raise ValueError("embedded credentials are not allowed") + if ( + any(ord(char) < 32 or ord(char) == 127 for char in url) + or any(char.isspace() for char in parsed.netloc) + or "\\" in parsed.netloc + or "%" in parsed.netloc + or parsed.netloc.endswith(":") + ): + raise ValueError("invalid URL authority") + if parsed.port is not None and not 1 <= parsed.port <= 65535: + raise ValueError("invalid port") + except ValueError as exc: + raise ValueError(f"Invalid input download URL: {exc}") from exc + + +def _validated_addresses(host: str, port: int) -> tuple[IPAddress, ...]: + """Resolve once, rejecting the entire answer if any address is blocked.""" + try: + literal = ipaddress.ip_address(host) + except ValueError: + try: + answers = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM) + addresses = [ipaddress.ip_address(answer[4][0]) for answer in answers] + except (OSError, ValueError) as exc: + raise ValueError("Unable to resolve input download hostname") from exc + else: + addresses = [literal] + + if not addresses: + raise ValueError("Unable to resolve input download hostname") + normalized = [] + for address in addresses: + if isinstance(address, ipaddress.IPv6Address): + address = address.ipv4_mapped or address + if ( + not _is_global(address) + or address.is_multicast + or address.is_reserved + or any(address in network for network in _INTERNAL_NETWORKS) + # Transition mechanisms can encode an otherwise blocked IPv4 + # destination, so do not permit them as public IPv6 destinations. + or ( + isinstance(address, ipaddress.IPv6Address) + and (address.sixtofour is not None or address.teredo is not None) + ) + ): + raise ValueError("Blocked input download destination") + if address not in normalized: + normalized.append(address) + return tuple(normalized) + + +def _connect(conn: HTTPConnection, addresses: tuple[IPAddress, ...]) -> socket.socket: + """Connect to numeric addresses without invoking DNS a second time.""" + last_error = None + for address in addresses: + family = socket.AF_INET if address.version == 4 else socket.AF_INET6 + sock = None + try: + sock = socket.socket(family, socket.SOCK_STREAM) + # Requests has converted the explicit timeout=10 to a numeric + # connect timeout before urllib3 creates this connection. + sock.settimeout(cast(float, conn.timeout)) + for option in conn.socket_options or (): + sock.setsockopt(*option) + sock.connect((str(address), conn.port)) + return sock + except OSError as exc: + if sock is not None: + sock.close() + last_error = exc + raise NewConnectionError( + conn, "Unable to connect to input download destination" + ) from last_error + + +def _pinned_adapter(addresses: tuple[IPAddress, ...]) -> HTTPAdapter: + # Only the TCP socket changes. urllib3 retains the original hostname for + # Host, SNI and certificate verification, including on connection retries. + class PinnedHTTPConnection(HTTPConnection): + def _new_conn(self) -> socket.socket: + return _connect(self, addresses) + + class PinnedHTTPSConnection(HTTPSConnection): + def _new_conn(self) -> socket.socket: + return _connect(self, addresses) + + class PinnedHTTPPool(HTTPConnectionPool): + ConnectionCls = PinnedHTTPConnection + + class PinnedHTTPSPool(HTTPSConnectionPool): + ConnectionCls = PinnedHTTPSConnection + + adapter = HTTPAdapter(max_retries=0) + # Assign a new mapping: mutating the default mapping would affect unrelated + # Requests sessions and concurrent downloads. + adapter.poolmanager.pool_classes_by_scheme = { + "http": PinnedHTTPPool, + "https": PinnedHTTPSPool, + } + return adapter + + +class _NoRedirectSession(requests.Session): + def resolve_redirects( + self, resp: requests.Response, req: requests.PreparedRequest, **kwargs: Any + ) -> Iterator[requests.Response]: + # Session.send prepares response.next even with allow_redirects=False, + # which consumes the redirect body. Reject without reading that body. + return iter(()) + + +def fetch_input(url: str, headers: dict[str, str]) -> requests.Response: + """Return a streaming response; the caller owns and must close it.""" + _validate_url(url) + response = None + try: + with _NoRedirectSession() as session: + # Neither environment proxies nor netrc may override this policy. + session.trust_env = False + request = session.prepare_request( + requests.Request("GET", url, headers=headers) + ) + assert request.url is not None + _validate_url(request.url) + parsed = urlsplit(request.url) + assert parsed.hostname is not None + port = parsed.port or (443 if parsed.scheme == "https" else 80) + addresses = _validated_addresses(parsed.hostname, port) + session.mount(f"{parsed.scheme}://", _pinned_adapter(addresses)) + response = session.send( + request, stream=True, timeout=10, allow_redirects=False + ) + if 300 <= response.status_code < 400: + raise ValueError("Input download redirects are not allowed") + response.raise_for_status() + response.raw.decode_content = True + return response + except (ValueError, requests.RequestException): + if response is not None: + response.close() + raise diff --git a/python/cog/types.py b/python/cog/types.py index 8183adb0de..713fa90e21 100644 --- a/python/cog/types.py +++ b/python/cog/types.py @@ -28,7 +28,7 @@ TypeVar, ) -import requests +from ._download import fetch_input # Constants for filename handling FILENAME_ILLEGAL_CHARS = set("\u0000/") @@ -129,7 +129,7 @@ class URLFile(io.IOBase): URL that can survive pickling/unpickling. """ - __slots__ = ("__target__", "__url__", "name") + __slots__ = ("__target__", "__response__", "__url__", "name") def __init__(self, url: str, filename: Optional[str] = None) -> None: parsed = urllib.parse.urlparse(url) @@ -168,13 +168,21 @@ def __setstate__(self, state: Dict[str, Any]) -> None: # Proxy getattr/setattr/delattr through to the response object. def __setattr__(self, name: str, value: Any) -> None: - if hasattr(type(self), name): + # IOBase accesses its closed flag through Python attribute hooks. + # Keep that bookkeeping local so close() never starts a download. + if name == "__IOBase_closed" or hasattr(type(self), name): object.__setattr__(self, name, value) else: setattr(self.__wrapped__, name, value) def __getattr__(self, name: str) -> Any: - if name in ("__target__", "__wrapped__", "__url__"): + if name in ( + "__target__", + "__response__", + "__wrapped__", + "__url__", + "__IOBase_closed", + ): raise AttributeError(name) elif name == "name": return object.__getattribute__(self, "name") @@ -190,8 +198,20 @@ def __delattr__(self, name: str) -> None: def __iter__(self) -> Iterator[bytes]: return iter(self.__wrapped__) + def close(self) -> None: + """Close an existing response without triggering a lazy download.""" + try: + response = object.__getattribute__(self, "__response__") + except AttributeError: + pass + else: + response.close() + super().close() + @property def __wrapped__(self) -> Any: + if self.closed: + raise ValueError("I/O operation on closed input download") try: return object.__getattribute__(self, "__target__") except AttributeError: @@ -202,9 +222,8 @@ def __wrapped__(self) -> Any: if ua: headers["User-Agent"] = ua - resp = requests.get(url, stream=True, timeout=10, headers=headers) - resp.raise_for_status() - resp.raw.decode_content = True + resp = fetch_input(url, headers) + object.__setattr__(self, "__response__", resp) object.__setattr__(self, "__target__", resp.raw) return resp.raw @@ -284,11 +303,20 @@ def __new__(cls, *, source: str, filename: str, fileobj: io.IOBase) -> "URLPath" def convert(self) -> "Path": """Download the URL content to a temporary file and return its Path.""" if self._path is None: - # pylint: disable=consider-using-with - dest = tempfile.NamedTemporaryFile(suffix=self.filename, delete=False) - shutil.copyfileobj(self.fileobj, dest) - dest.close() - self._path = Path(dest.name) + dest_name = None + try: + with tempfile.NamedTemporaryFile( + suffix=self.filename, delete=False + ) as dest: + dest_name = dest.name + shutil.copyfileobj(self.fileobj, dest) + self._path = Path(dest_name) + except BaseException: + if dest_name is not None: + os.unlink(dest_name) + raise + finally: + self.fileobj.close() return self._path def unlink(self, missing_ok: bool = False) -> None: diff --git a/python/tests/test_download.py b/python/tests/test_download.py new file mode 100644 index 0000000000..fdaad7ad18 --- /dev/null +++ b/python/tests/test_download.py @@ -0,0 +1,255 @@ +"""Regressions for File/Path input SSRF; all HTTP listeners are local.""" + +import io +import os +import socket +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from types import SimpleNamespace + +import pytest +import requests + +from cog.types import File, Path, URLFile, URLPath + +PUBLIC_IP = "8.8.8.8" + + +@pytest.fixture(params=["file", "path", "urlfile"]) +def download(request): + if request.param == "path" and os.name == "nt": + pytest.skip("Cog Path requires POSIX") + + def read(url): + if request.param == "path": + value = Path.validate(url) + try: + return value.convert().read_bytes() + finally: + value.unlink(missing_ok=True) + value = File.validate(url) if request.param == "file" else URLFile(url) + with value: + return value.read() + + return read + + +@pytest.fixture +def listener(): + servers = [] + + def start(redirect=None): + hits = [] + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + hits.append(dict(self.headers)) + self.send_response(302 if redirect else 200) + self.send_header("Content-Length", "2") + if redirect: + self.send_header("Location", redirect) + self.end_headers() + self.wfile.write(b"ok") + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread( + target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True + ) + thread.start() + servers.append((server, thread)) + return SimpleNamespace(address=server.server_address, hits=hits) + + yield start + for server, thread in servers: + server.shutdown() + server.server_close() + thread.join() + + +def answers(*addresses): + return [ + ( + socket.AF_INET6 if ":" in ip else socket.AF_INET, + socket.SOCK_STREAM, + socket.IPPROTO_TCP, + "", + (ip, 80), + ) + for ip in addresses + ] + + +def route_public(monkeypatch, origin): + """Route only the validated numeric address to the controlled origin.""" + real_socket = socket.socket + connections = [] + + class RoutedSocket(real_socket): + def connect(self, address): + connections.append(address) + assert address[0] == PUBLIC_IP + return super().connect(origin.address) + + monkeypatch.setattr(socket, "socket", RoutedSocket) + return connections + + +@pytest.mark.parametrize( + "host", + [ + "127.0.0.1", + "localhost", + "10.0.0.1", + "169.254.169.254", + "168.63.129.16", + "224.0.0.1", + "[::1]", + "[::ffff:127.0.0.1]", + ], +) +def test_internal_input_is_rejected(download, listener, host): + blocked = listener() + with pytest.raises(ValueError, match="Blocked input download"): + download(f"http://{host}:{blocked.address[1]}/secret") + assert blocked.hits == [] + + +@pytest.mark.parametrize( + "addresses", + [ + (PUBLIC_IP, "127.0.0.1"), + ("2606:4700::1111", "::1"), + ], +) +def test_mixed_dns_is_rejected(download, listener, monkeypatch, addresses): + blocked = listener() + monkeypatch.setattr(socket, "getaddrinfo", lambda *a, **kw: answers(*addresses)) + with pytest.raises(ValueError, match="Blocked input download"): + download(f"http://download.test:{blocked.address[1]}/secret") + assert blocked.hits == [] + + +def test_resolution_failure_is_rejected(download, monkeypatch): + def resolve(*args, **kwargs): + raise socket.gaierror("test resolution failure") + + monkeypatch.setattr(socket, "getaddrinfo", resolve) + with pytest.raises(ValueError, match="Unable to resolve"): + download("http://download.test/file") + + +@pytest.mark.parametrize( + "url", + [ + "http:///file", + "http://user:password@download.test/file", + "http://download.test:65536/file", + ], +) +def test_invalid_input_url(download, url): + with pytest.raises(ValueError): + download(url) + + +def test_allowed_download_is_pinned_and_bypasses_proxies( + download, listener, monkeypatch +): + origin, blocked = listener(), listener() + connections = route_public(monkeypatch, origin) + lookups = [] + + def resolve(host, port, **kwargs): + lookups.append((host, port)) + return answers(PUBLIC_IP if len(lookups) == 1 else "127.0.0.1") + + monkeypatch.setattr(socket, "getaddrinfo", resolve) + for name in ( + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "http_proxy", + "https_proxy", + "all_proxy", + ): + monkeypatch.setenv(name, f"http://127.0.0.1:{blocked.address[1]}") + monkeypatch.setenv("NO_PROXY", "") + monkeypatch.setenv("no_proxy", "") + assert download("http://download.test:8080/file") == b"ok" + assert lookups == [("download.test", 8080)] + assert connections == [(PUBLIC_IP, 8080)] + assert origin.hits[0]["Host"] == "download.test:8080" + assert blocked.hits == [] + + +def test_redirect_is_rejected(download, listener, monkeypatch): + blocked = listener() + origin = listener(redirect=f"http://127.0.0.1:{blocked.address[1]}/secret") + route_public(monkeypatch, origin) + monkeypatch.setattr(socket, "getaddrinfo", lambda *a, **kw: answers(PUBLIC_IP)) + with pytest.raises(ValueError, match="redirects are not allowed"): + download("http://download.test/file") + assert len(origin.hits) == 1 + assert blocked.hits == [] + + +@pytest.mark.parametrize("scheme,port", [("http", 80), ("https", 443)]) +def test_fetch_preserves_request_options(monkeypatch, scheme, port): + response = requests.Response() + response.status_code = 200 + response.raw = io.BytesIO(b"ok") + lookups = [] + + def resolve(host, resolved_port, **kwargs): + lookups.append((host, resolved_port)) + return answers(PUBLIC_IP) + + def send(adapter, request, **kwargs): + assert request.url == f"{scheme}://download.test/file" + assert kwargs["verify"] is True + assert kwargs["stream"] is True + assert kwargs["timeout"] == 10 + assert kwargs["proxies"] == {} + pool = adapter.poolmanager.connection_from_url(request.url) + assert pool.host == "download.test" + assert pool.port == port + return response + + monkeypatch.setattr(socket, "getaddrinfo", resolve) + monkeypatch.setattr(requests.adapters.HTTPAdapter, "send", send) + with URLFile(f"{scheme}://download.test/file") as file: + assert file.read() == b"ok" + assert response.raw.closed + assert lookups == [("download.test", port)] + + +def test_data_and_local_inputs(tmp_path): + assert File.validate("data:text/plain;base64,b2s=").read() == b"ok" + local = tmp_path / "file.txt" + local.write_bytes(b"ok") + assert Path.validate(local) is local + if os.name != "nt": + assert Path.validate(str(local)).read_bytes() == b"ok" + value = Path.validate("data:text/plain;base64,b2s=") + try: + assert value.convert().read_bytes() == b"ok" + finally: + value.unlink() + + +def test_failed_copy_removes_partial_file(tmp_path, monkeypatch): + class BrokenFile(io.BytesIO): + def read(self, size=-1): + if self.tell(): + raise OSError("interrupted download") + return super().read(3) + + source = BrokenFile(b"partial") + value = SimpleNamespace(_path=None, filename="file.txt", fileobj=source) + monkeypatch.setattr("tempfile.tempdir", str(tmp_path)) + with pytest.raises(OSError, match="interrupted"): + URLPath.convert(value) + assert list(tmp_path.iterdir()) == [] + assert source.closed