Arbitrary File Read through Path Traversal in run-llama/llama_index

@@ -32,6 +32,7 @@
 from dataclasses_json import DataClassJsonMixin
 from deprecated import deprecated
 from typing_extensions import Self
+from PIL import Image
 
 from llama_index.core.bridge.pydantic import (
     AnyUrl,
@@ -1220,6 +1221,27 @@ def from_cloud_document(
         )
 
 
+def is_image_pil(file_path: str) -> bool:
+    try:
+        with Image.open(file_path) as img:
+            img.verify()  # Verify it's a valid image
+        return True
+    except (IOError, SyntaxError):
+        return False
+
+
+def is_image_url_pil(url: str) -> bool:
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # Raise an exception for bad status codes
+        # Open image from the response content
+        img = Image.open(BytesIO(response.content))
+        img.verify()
+        return True
+    except (requests.RequestException, IOError, SyntaxError):
+        return False
+
+
 class ImageDocument(Document):
     """Backward compatible wrapper around Document containing an image."""
 
@@ -1235,10 +1257,14 @@ def __init__(self, **kwargs: Any) -> None:
                 data=image, mimetype=image_mimetype
             )
         elif image_path:
+            if not is_image_pil(image_path):
+                raise ValueError("The specified file path is not an accessible image")
             kwargs["image_resource"] = MediaResource(
                 path=image_path, mimetype=image_mimetype
             )
         elif image_url:
+            if not is_image_url_pil(image_url):
+                raise ValueError("The specified URL is not an accessible image")
             kwargs["image_resource"] = MediaResource(
                 url=image_url, mimetype=image_mimetype
             )

@@ -2,6 +2,8 @@
 
 import pytest
 import base64
+import httpx
+from pathlib import Path
 from unittest.mock import mock_open, patch, MagicMock
 
 from llama_index.core.schema import ImageDocument
@@ -16,7 +18,9 @@
 )
 
 # Expected values
-EXP_IMAGE_URLS = ["http://example.com/image1.jpg"]
+EXP_IMAGE_URLS = [
+    "https://astrabert.github.io/hophop-science/images/whale_doing_science.png"
+]
 EXP_BASE64 = "SGVsbG8gV29ybGQ="  # "Hello World" in base64
 EXP_BINARY = b"Hello World"
 
@@ -51,22 +55,26 @@ def test_encode_image():
     assert result == EXP_BASE64
 
 
-def test_image_documents_to_base64_multiple_sources():
+def test_image_documents_to_base64_multiple_sources(tmp_path: Path):
     """Test converting multiple ImageDocuments with different source types."""
+    content = httpx.get(EXP_IMAGE_URLS[0]).content
+    expected_b64 = base64.b64encode(content).decode("utf-8")
+    fl_path = tmp_path / "test_image.png"
+    fl_path.write_bytes(content)
     documents = [
-        ImageDocument(image=EXP_BASE64),
-        ImageDocument(image_path="test.jpg"),
+        ImageDocument(image=expected_b64),
+        ImageDocument(image_path=fl_path),
         ImageDocument(metadata={"file_path": "test.jpg"}),
         ImageDocument(image_url=EXP_IMAGE_URLS[0]),
     ]
     with patch("requests.get") as mock_get:
-        mock_get.return_value.content = EXP_BINARY
+        mock_get.return_value.content = content
         with patch("os.path.isfile", return_value=True):
-            with patch("builtins.open", mock_open(read_data=EXP_BINARY)):
+            with patch("builtins.open", mock_open(read_data=content)):
                 result = image_documents_to_base64(documents)
 
     assert len(result) == 4
-    assert all(encoding == EXP_BASE64 for encoding in result)
+    assert all(encoding == expected_b64 for encoding in result)
 
 
 def test_image_documents_to_base64_failed_url():
@@ -136,11 +144,15 @@ def test_infer_image_mimetype_from_file_path():
     assert infer_image_mimetype_from_file_path("") == "image/jpeg"
 
 
-def test_set_base64_and_mimetype_for_image_docs():
+def test_set_base64_and_mimetype_for_image_docs(tmp_path: Path):
     """Test setting base64 and mimetype fields for ImageDocument objects."""
+    content = httpx.get(EXP_IMAGE_URLS[0]).content
+    expected_b64 = base64.b64encode(content).decode("utf-8")
+    fl_path = tmp_path / "test_image.png"
+    fl_path.write_bytes(content)
     image_docs = [
-        ImageDocument(image=EXP_BASE64),
-        ImageDocument(image_path="test.asdf"),
+        ImageDocument(image=expected_b64),
+        ImageDocument(image_path=fl_path.__str__()),
     ]
 
     with patch("requests.get") as mock_get:
@@ -151,6 +163,6 @@ def test_set_base64_and_mimetype_for_image_docs():
                 results = set_base64_and_mimetype_for_image_docs(image_docs)
 
     assert len(results) == 2
-    assert results[0].image == EXP_BASE64
+    assert results[0].image == expected_b64
     assert results[0].image_mimetype == "image/jpeg"
     assert results[1].image_mimetype == "image/jpeg"

@@ -0,0 +1,36 @@
+import httpx

+import pytest

+

+from pathlib import Path

+from llama_index.core.schema import ImageDocument

+

+

+@pytest.fixture()

+def image_url() -> str:

+    return "https://astrabert.github.io/hophop-science/images/whale_doing_science.png"

+

+

+def test_real_image_path(tmp_path: Path, image_url: str) -> None:

+    content = httpx.get(image_url).content

+    fl_path = tmp_path / "test_image.png"

+    fl_path.write_bytes(content)

+    doc = ImageDocument(image_path=fl_path.__str__())

+    assert isinstance(doc, ImageDocument)

+

+

+def test_real_image_url(image_url: str) -> None:

+    doc = ImageDocument(image_url=image_url)

+    assert isinstance(doc, ImageDocument)

+

+

+def test_non_image_path(tmp_path: Path) -> None:

+    fl_path = tmp_path / "test_file.txt"

+    fl_path.write_text("Hello world!")

+    with pytest.raises(expected_exception=ValueError):

+        doc = ImageDocument(image_path=fl_path.__str__())

+

+

+def test_non_image_url(image_url: str) -> None:

+    image_url = image_url.replace("png", "txt")

+    with pytest.raises(expected_exception=ValueError):

+        doc = ImageDocument(image_url=image_url)

@@ -1,4 +1,5 @@
 import base64
+import httpx
 import logging
 from io import BytesIO
 from pathlib import Path
@@ -250,17 +251,27 @@ def test_image_document_image():
     assert doc.image == "MTIzNDU2Nzg5"
 
 
-def test_image_document_path():
-    mock_path = Path(__file__)
-    doc = ImageDocument(id_="test", image_path=mock_path)
-    assert doc.image_path == str(mock_path)
-    doc.image_path = str(mock_path.parent)
-    assert doc.image_path == str(mock_path.parent)
+def test_image_document_path(tmp_path: Path):
+    content = httpx.get(
+        "https://astrabert.github.io/hophop-science/images/whale_doing_science.png"
+    ).content
+    fl_path = tmp_path / "test_image.png"
+    fl_path.write_bytes(content)
+    doc = ImageDocument(id_="test", image_path=fl_path)
+    assert doc.image_path == str(fl_path)
+    doc.image_path = str(fl_path.parent)
+    assert doc.image_path == str(fl_path.parent)
 
 
 def test_image_document_url():
-    doc = ImageDocument(id_="test", image_url="https://example.com/")
-    assert doc.image_url == "https://example.com/"
+    doc = ImageDocument(
+        id_="test",
+        image_url="https://astrabert.github.io/hophop-science/images/whale_doing_science.png",
+    )
+    assert (
+        doc.image_url
+        == "https://astrabert.github.io/hophop-science/images/whale_doing_science.png"
+    )
     doc.image_url = "https://foo.org"
     assert doc.image_url == "https://foo.org/"
 
@@ -281,12 +292,17 @@ def test_image_document_embeddings():
     assert doc.text_resource.embeddings == {"dense": [1.0, 2.0, 3.0]}
 
 
-def test_image_document_path_serialization():
-    doc = ImageDocument(image_path=Path("test.png"))
-    assert doc.model_dump()["image_resource"]["path"] == "test.png"
+def test_image_document_path_serialization(tmp_path: Path):
+    content = httpx.get(
+        "https://astrabert.github.io/hophop-science/images/whale_doing_science.png"
+    ).content
+    fl_path = tmp_path / "test_image.png"
+    fl_path.write_bytes(content)
+    doc = ImageDocument(image_path=fl_path)
+    assert doc.model_dump()["image_resource"]["path"] == fl_path.__str__()
 
     new_doc = ImageDocument(**doc.model_dump())
-    assert new_doc.image_resource.path == Path("test.png")
+    assert new_doc.image_resource.path == fl_path
 
 
 def test_image_block_resolve_image(png_1px: bytes, png_1px_b64: bytes):