diff --git a/fastchat/utils.py b/fastchat/utils.py index d3531928f..7f809ffa5 100644 --- a/fastchat/utils.py +++ b/fastchat/utils.py @@ -401,13 +401,17 @@ def load_image(image_file): 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))) + # Reject filesystem paths (CWE-22); callers pass data URLs or raw base64. + try: + image = Image.open(BytesIO(base64.b64decode(image_file, validate=True))) + except Exception as e: + raise ValueError( + "load_image expects http(s), data URI, or base64; local paths are not allowed" + ) from e return image diff --git a/tests/test_load_image_no_local_path.py b/tests/test_load_image_no_local_path.py new file mode 100644 index 000000000..50d638f2d --- /dev/null +++ b/tests/test_load_image_no_local_path.py @@ -0,0 +1,28 @@ +""" +python3 -m unittest tests.test_load_image_no_local_path +""" + +import tempfile +import unittest +from pathlib import Path + +from fastchat.utils import load_image + + +class LoadImageNoLocalPathTest(unittest.TestCase): + def test_rejects_local_image_path(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "x.png" + # minimal 1x1 PNG + path.write_bytes( + bytes.fromhex( + "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489" + "0000000a49444154789c63000100000500010d0a2db40000000049454e44ae426082" + ) + ) + with self.assertRaises(ValueError): + load_image(str(path)) + + +if __name__ == "__main__": + unittest.main()