Serve _upload files with content-disposition: attachment

packages/reflex-components-core/src/reflex_components_core/core/_upload.py

@@ -6,7 +6,13 @@
 import contextlib
 import dataclasses
 from collections import deque
-from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
+from collections.abc import (
+    AsyncGenerator,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    MutableMapping,
+)
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, BinaryIO, cast
 
@@ -22,7 +28,7 @@
 from typing_extensions import Self
 
 if TYPE_CHECKING:
-    from reflex_base.utils.types import Receive, Scope, Send
+    from reflex_base.utils.types import ASGIApp, Receive, Scope, Send
 
     from reflex.app import App
 
@@ -597,6 +603,45 @@ async def _upload_chunk_file(
     return Response(status_code=202)
 
 
+class UploadedFilesHeadersMiddleware:
+    """ASGI middleware that adds security headers to uploaded file responses."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        """Wrap an ASGI application with upload security headers.
+
+        Args:
+            app: The ASGI application to wrap.
+        """
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Add Content-Disposition and X-Content-Type-Options headers.
+
+        Args:
+            scope: The ASGI scope.
+            receive: The ASGI receive callable.
+            send: The ASGI send callable.
+        """
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        is_pdf = scope.get("path", "").lower().endswith(".pdf")
+
+        async def send_with_headers(message: MutableMapping[str, Any]) -> None:
+            if message["type"] == "http.response.start":
+                headers = list(message.get("headers", []))
+                headers.append((b"x-content-type-options", b"nosniff"))
+                if is_pdf:
+                    headers.append((b"content-type", b"application/pdf"))
+                else:
+                    headers.append((b"content-disposition", b"attachment"))
+                message = {**message, "headers": headers}
+            await send(message)
+
+        await self.app(scope, receive, send_with_headers)
+
+
 def upload(app: App):
     """Upload files, dispatching to buffered or streaming handling.
 

reflex/app.py

@@ -68,8 +68,8 @@
 from starlette.staticfiles import StaticFiles
 from typing_extensions import Unpack
 
+from reflex._upload import UploadedFilesHeadersMiddleware, upload
 from reflex._upload import UploadFile as UploadFile
-from reflex._upload import upload
 from reflex.admin import AdminDash
 from reflex.app_mixins import AppMixin, LifespanMixin, MiddlewareMixin
 from reflex.compiler import compiler
@@ -721,7 +721,7 @@ def _add_optional_endpoints(self):
             # To access uploaded files.
             self._api.mount(
                 str(constants.Endpoint.UPLOAD),
-                StaticFiles(directory=get_upload_dir()),
+                UploadedFilesHeadersMiddleware(StaticFiles(directory=get_upload_dir())),
                 name="uploaded_files",
             )
 

tests/integration/test_media.py

@@ -55,6 +55,15 @@ def img_from_url(self) -> Image.Image:
             img_bytes = img_resp.content
             return Image.open(io.BytesIO(img_bytes))
 
+        @rx.var
+        def generated_image(self) -> str:
+            # Generate a 150x150 red PNG and write it to the upload directory.
+            img = Image.new("RGB", (150, 150), "red")
+            upload_dir = rx.get_upload_dir()
+            upload_dir.mkdir(parents=True, exist_ok=True)
+            img.save(upload_dir / "generated.png")
+            return "generated.png"
+
     app = rx.App()
 
     @app.add_page
@@ -72,19 +81,53 @@ def index():
             rx.image(src=State.img_gif, alt="GIF image", id="gif"),
             rx.image(src=State.img_webp, alt="WEBP image", id="webp"),
             rx.image(src=State.img_from_url, alt="Image from URL", id="from_url"),
+            rx.image(
+                src=rx.get_upload_url(State.generated_image),
+                alt="Uploaded image",
+                id="uploaded",
+            ),
         )
 
 
+def check_image_loaded(
+    driver, img, expected_width: int = 200, expected_height: int = 200
+) -> bool:
+    """Check whether an image element has fully loaded with expected dimensions.
+
+    Args:
+        driver: WebDriver instance.
+        img: The image WebElement.
+        expected_width: Expected natural width.
+        expected_height: Expected natural height.
+
+    Returns:
+        True if the image is complete and matches the expected dimensions.
+    """
+    return driver.execute_script(
+        "return arguments[0].complete "
+        '&& typeof arguments[0].naturalWidth != "undefined" '
+        "&& arguments[0].naturalWidth === arguments[1] "
+        '&& typeof arguments[0].naturalHeight != "undefined" '
+        "&& arguments[0].naturalHeight === arguments[2]",
+        img,
+        expected_width,
+        expected_height,
+    )
+
+
 @pytest.fixture
-def media_app(tmp_path) -> Generator[AppHarness, None, None]:
+def media_app(tmp_path, monkeypatch) -> Generator[AppHarness, None, None]:
     """Start MediaApp app at tmp_path via AppHarness.
 
     Args:
         tmp_path: pytest tmp_path fixture
+        monkeypatch: pytest monkeypatch fixture
 
     Yields:
         running AppHarness instance
     """
+    monkeypatch.setenv("REFLEX_UPLOADED_FILES_DIR", str(tmp_path / "uploads"))
+
     with AppHarness.create(
         root=tmp_path,
         app_source=MediaApp,
@@ -116,52 +159,45 @@ def test_media_app(media_app: AppHarness):
     gif_img = driver.find_element(By.ID, "gif")
     webp_img = driver.find_element(By.ID, "webp")
     from_url_img = driver.find_element(By.ID, "from_url")
-
-    def check_image_loaded(img, check_width=" == 200", check_height=" == 200"):
-        return driver.execute_script(
-            "console.log(arguments); return arguments[1].complete "
-            '&& typeof arguments[1].naturalWidth != "undefined" '
-            f"&& arguments[1].naturalWidth {check_width} ",
-            '&& typeof arguments[1].naturalHeight != "undefined" '
-            f"&& arguments[1].naturalHeight {check_height} ",
-            img,
-        )
+    uploaded_img = driver.find_element(By.ID, "uploaded")
 
     default_img_src = default_img.get_attribute("src")
     assert default_img_src is not None
     assert default_img_src.startswith("data:image/png;base64")
-    assert check_image_loaded(default_img)
+    assert check_image_loaded(driver, default_img)
 
     bmp_img_src = bmp_img.get_attribute("src")
     assert bmp_img_src is not None
     assert bmp_img_src.startswith("data:image/bmp;base64")
-    assert check_image_loaded(bmp_img)
+    assert check_image_loaded(driver, bmp_img)
 
     jpg_img_src = jpg_img.get_attribute("src")
     assert jpg_img_src is not None
     assert jpg_img_src.startswith("data:image/jpeg;base64")
-    assert check_image_loaded(jpg_img)
+    assert check_image_loaded(driver, jpg_img)
 
     png_img_src = png_img.get_attribute("src")
     assert png_img_src is not None
     assert png_img_src.startswith("data:image/png;base64")
-    assert check_image_loaded(png_img)
+    assert check_image_loaded(driver, png_img)
 
     gif_img_src = gif_img.get_attribute("src")
     assert gif_img_src is not None
     assert gif_img_src.startswith("data:image/gif;base64")
-    assert check_image_loaded(gif_img)
+    assert check_image_loaded(driver, gif_img)
 
     webp_img_src = webp_img.get_attribute("src")
     assert webp_img_src is not None
     assert webp_img_src.startswith("data:image/webp;base64")
-    assert check_image_loaded(webp_img)
+    assert check_image_loaded(driver, webp_img)
 
     from_url_img_src = from_url_img.get_attribute("src")
     assert from_url_img_src is not None
     assert from_url_img_src.startswith("data:image/jpeg;base64")
-    assert check_image_loaded(
-        from_url_img,
-        check_width=" == 200",
-        check_height=" == 300",
-    )
+    assert check_image_loaded(driver, from_url_img, expected_height=300)
+
+    breakpoint()
+    uploaded_img_src = uploaded_img.get_attribute("src")
+    assert uploaded_img_src is not None
+    assert "generated.png" in uploaded_img_src
+    assert check_image_loaded(driver, uploaded_img, 150, 150)

tests/integration/test_upload.py

@@ -743,6 +743,63 @@ def test_upload_download_file(
     assert driver.find_element(by=By.TAG_NAME, value="body").text == exp_contents
 
 
+def test_uploaded_file_security_headers(
+    tmp_path,
+    upload_file: AppHarness,
+    driver: WebDriver,
+):
+    """Upload an HTML file and verify security headers prevent inline rendering.
+
+    Args:
+        tmp_path: pytest tmp_path fixture
+        upload_file: harness for UploadFile app.
+        driver: WebDriver instance.
+    """
+    import httpx
+    from reflex_base.config import get_config
+
+    assert upload_file.app_instance is not None
+    poll_for_token(driver, upload_file)
+    clear_btn = driver.find_element(By.ID, "clear_uploads")
+    clear_btn.click()
+
+    upload_box = driver.find_elements(By.XPATH, "//input[@type='file']")[2]
+    upload_button = driver.find_element(By.ID, "upload_button_tertiary")
+
+    exp_name = "malicious.html"
+    exp_contents = "<html><body><script>alert('xss')</script></body></html>"
+    target_file = tmp_path / exp_name
+    target_file.write_text(exp_contents)
+
+    upload_box.send_keys(str(target_file))
+    upload_button.click()
+
+    upload_done = driver.find_element(By.ID, "upload_done")
+    assert upload_file.poll_for_value(upload_done, exp_not_equal="false") == "true"
+
+    # Fetch the uploaded file directly via httpx and check security headers.
+    upload_url = f"{get_config().api_url}/{Endpoint.UPLOAD.value}/{exp_name}"
+    resp = httpx.get(upload_url)
+    assert resp.status_code == 200
+    assert resp.text == exp_contents
+    assert resp.headers["content-disposition"] == "attachment"
+    assert resp.headers["x-content-type-options"] == "nosniff"
+
+    # Navigate to the uploaded HTML file in the browser and verify the script
+    # does not execute (Content-Disposition: attachment prevents rendering).
+    driver.get(upload_url)
+    # If the browser rendered the HTML, an alert('xss') dialog would appear.
+    # Verify no alert is present — the file should be downloaded, not rendered.
+    from selenium.common.exceptions import NoAlertPresentException
+
+    try:
+        alert = driver.switch_to.alert
+        alert.dismiss()
+        pytest.fail("Browser rendered the HTML and triggered an alert dialog")
+    except NoAlertPresentException:
+        pass  # Expected: no alert because the file was downloaded, not rendered
+
+
 def test_on_drop(
     tmp_path,
     upload_file: AppHarness,