diff --git a/fastchat/utils.py b/fastchat/utils.py index d3531928f..29213a554 100644 --- a/fastchat/utils.py +++ b/fastchat/utils.py @@ -391,6 +391,32 @@ def str_to_torch_dtype(dtype: str): raise ValueError(f"Unrecognized dtype: {dtype}") +import ipaddress +import socket +from urllib.parse import urlparse + + +def _is_safe_http_url(url: str) -> bool: + parsed = urlparse(url) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + return False + try: + port = parsed.port or (443 if parsed.scheme == "https" else 80) + for _, _, _, _, sockaddr in socket.getaddrinfo(parsed.hostname, port): + ip = ipaddress.ip_address(sockaddr[0]) + if ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_reserved + or ip.is_multicast + ): + return False + except socket.gaierror: + return False + return True + + def load_image(image_file): from PIL import Image import requests @@ -398,16 +424,18 @@ def load_image(image_file): image = None if image_file.startswith("http://") or image_file.startswith("https://"): + if not _is_safe_http_url(image_file): + raise ValueError("Refusing to fetch image from internal or private URL") timeout = int(os.getenv("REQUEST_TIMEOUT", "3")) response = requests.get(image_file, timeout=timeout) image = Image.open(BytesIO(response.content)) elif image_file.lower().endswith(("png", "jpg", "jpeg", "webp", "gif")): image = Image.open(image_file) elif image_file.startswith("data:"): - image_file = image_file.split(",")[1] - image = Image.open(BytesIO(base64.b64decode(image_file))) + image_file = image_file.split(",", 1)[1] + image = Image.open(BytesIO(base64.b64decode(image_file, validate=True))) else: - image = Image.open(BytesIO(base64.b64decode(image_file))) + image = Image.open(BytesIO(base64.b64decode(image_file, validate=True))) return image diff --git a/tests/test_load_image_security.py b/tests/test_load_image_security.py new file mode 100644 index 000000000..0d463d4ec --- /dev/null +++ b/tests/test_load_image_security.py @@ -0,0 +1,28 @@ +""" +Usage: +python3 -m unittest tests.test_load_image_security +""" + +import unittest +from unittest.mock import patch + +from fastchat.utils import _is_safe_http_url, load_image + + +class LoadImageSecurityTest(unittest.TestCase): + def test_blocks_private_http_urls(self) -> None: + self.assertFalse(_is_safe_http_url("http://127.0.0.1/image.png")) + self.assertFalse(_is_safe_http_url("http://169.254.169.254/latest/meta-data/")) + + def test_allows_public_http_urls(self) -> None: + self.assertTrue(_is_safe_http_url("https://example.com/image.png")) + + @patch("fastchat.utils.requests.get") + def test_load_image_rejects_internal_url(self, mock_get) -> None: + with self.assertRaises(ValueError): + load_image("http://127.0.0.1/secret.png") + mock_get.assert_not_called() + + +if __name__ == "__main__": + unittest.main()