Python: keep attachments close

python/samples/05-end-to-end/chatkit-integration/app.py

@@ -587,12 +587,17 @@ async def upload_file(attachment_id: str, file: UploadFile = File(...)):  # noqa
     """
     logger.info(f"Receiving file upload for attachment: {attachment_id}")
 
+    try:
+        file_path = attachment_store.get_file_path(attachment_id)
+    except ValueError:
+        logger.warning(f"Rejected invalid attachment ID: {attachment_id!r}")
+        return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."})
+
     try:
         # Read file contents
         contents = await file.read()
 
         # Save to disk
-        file_path = attachment_store.get_file_path(attachment_id)
         file_path.write_bytes(contents)
 
         logger.info(f"Saved {len(contents)} bytes to {file_path}")
@@ -625,7 +630,11 @@ async def preview_image(attachment_id: str):
 
     try:
         file_path = attachment_store.get_file_path(attachment_id)
+    except ValueError:
+        logger.warning(f"Rejected invalid attachment ID: {attachment_id!r}")
+        return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."})
 
+    try:
         if not file_path.exists():
             return JSONResponse(status_code=404, content={"error": "File not found"})
 

python/samples/05-end-to-end/chatkit-integration/attachment_store.py

@@ -48,16 +48,32 @@ def __init__(
             base_url: Base URL for generating upload and preview URLs
             data_store: Optional data store to persist attachment metadata
         """
-        self.uploads_dir = Path(uploads_dir)
+        self.uploads_dir = Path(uploads_dir).resolve()
         self.base_url = base_url.rstrip("/")
         self.data_store = data_store
 
         # Create uploads directory if it doesn't exist
         self.uploads_dir.mkdir(parents=True, exist_ok=True)
 
     def get_file_path(self, attachment_id: str) -> Path:
-        """Get the filesystem path for an attachment."""
-        return self.uploads_dir / attachment_id
+        """Get the filesystem path for an attachment.
+
+        Args:
+            attachment_id: Identifier used as the attachment filename.
+
+        Returns:
+            The resolved path within the uploads directory.
+
+        Raises:
+            ValueError: If the attachment ID does not resolve to a direct child of the uploads directory.
+        """
+        if not attachment_id or attachment_id in {".", ".."} or "/" in attachment_id or "\\" in attachment_id:
+            raise ValueError(f"Invalid attachment ID: {attachment_id!r}")
+
+        file_path = (self.uploads_dir / attachment_id).resolve()
+        if not file_path.is_relative_to(self.uploads_dir) or file_path.parent != self.uploads_dir:
+            raise ValueError(f"Invalid attachment ID: {attachment_id!r}")
+        return file_path
 
     async def delete_attachment(self, attachment_id: str, context: dict[str, Any]) -> None:
         """Delete an attachment and its file from disk."""

python/tests/samples/end_to_end/test_chatkit_attachment_store.py

@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for the ChatKit integration sample attachment store."""
+
+import importlib.util
+import json
+from io import BytesIO
+from pathlib import Path
+from types import ModuleType
+
+import agent_framework
+import pytest
+
+_ATTACHMENT_STORE_PATH = (
+    Path(__file__).parents[3] / "samples" / "05-end-to-end" / "chatkit-integration" / "attachment_store.py"
+)
+
+
+def _load_attachment_store_module() -> ModuleType:
+    spec = importlib.util.spec_from_file_location("chatkit_attachment_store", _ATTACHMENT_STORE_PATH)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+attachment_store_module = _load_attachment_store_module()
+
+
+@pytest.fixture
+def chatkit_app_module(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> ModuleType:
+    sample_dir = _ATTACHMENT_STORE_PATH.parent
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("FOUNDRY_MODEL", "test-model")
+    monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://example.com")
+    monkeypatch.setattr(agent_framework, "FunctionResultContent", object, raising=False)
+    monkeypatch.syspath_prepend(str(sample_dir))
+
+    spec = importlib.util.spec_from_file_location("chatkit_integration_app", sample_dir / "app.py")
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_get_file_path_returns_direct_child(tmp_path: Path) -> None:
+    store = attachment_store_module.FileBasedAttachmentStore(uploads_dir=str(tmp_path))
+
+    assert store.get_file_path("attachment-123") == tmp_path / "attachment-123"
+
+
+@pytest.mark.parametrize(
+    "attachment_id",
+    [
+        "../outside",
+        "nested/attachment-123",
+        "nested/../attachment-123",
+        r"nested\attachment-123",
+        r"nested\..\attachment-123",
+        "/tmp/attachment-123",
+        "",
+        ".",
+        "..",
+    ],
+)
+def test_get_file_path_rejects_non_filename_ids(tmp_path: Path, attachment_id: str) -> None:
+    store = attachment_store_module.FileBasedAttachmentStore(uploads_dir=str(tmp_path))
+
+    with pytest.raises(ValueError, match="Invalid attachment ID"):
+        store.get_file_path(attachment_id)
+
+
+async def test_attachment_routes_return_bad_request_for_invalid_id(chatkit_app_module: ModuleType) -> None:
+    upload = chatkit_app_module.UploadFile(file=BytesIO(b"contents"), filename="attachment.txt")
+
+    upload_response = await chatkit_app_module.upload_file(".", upload)
+    preview_response = await chatkit_app_module.preview_image(".")
+
+    assert upload_response.status_code == 400
+    assert json.loads(upload_response.body) == {"error": "Invalid attachment ID."}
+    assert preview_response.status_code == 400
+    assert json.loads(preview_response.body) == {"error": "Invalid attachment ID."}