feat(#320): convert upload_artifact to true streaming multipart parsing with mid-stream 413 abort

backend/app/api/external_results.py

@@ -18,8 +18,10 @@
 import uuid as _uuid_module
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, File, Form, HTTPException, Response, UploadFile, status
+from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
 from sqlalchemy.ext.asyncio import AsyncSession
+from streaming_form_data import StreamingFormDataParser
+from streaming_form_data.targets import FileTarget, ValueTarget
 
 from app.auth.dependencies import (
     get_runner_or_user_auth,
@@ -51,11 +53,11 @@
 _DEFAULT_RUNNER = "@bgstm/playwright-core@unknown"
 
 # ---------------------------------------------------------------------------
-# Artifact upload constants
+# Artifact upload constants and helpers
 # ---------------------------------------------------------------------------
 
-# Read the upload stream in 64 KiB chunks.  Tests may monkeypatch this value
-# to a smaller number to exercise the streaming / partial-write path.
+# Kept as a module attribute so existing monkeypatch calls in tests don't fail
+# (the streaming implementation uses network-driven chunk sizes, not this value).
 _ARTIFACT_CHUNK_SIZE: int = 65_536  # 64 KiB
 
 # Content-type allowlist (global).  ``artifact_kind.other`` bypasses this check.
@@ -427,16 +429,45 @@ def _safe_unlink(path: str) -> None:
         pass
 
 
+class _SizeLimitExceeded(Exception):
+    """Raised inside ``_SizeLimitedFileTarget.on_data_received`` when the
+    running byte total exceeds ``max_bytes``.  The stream loop catches this
+    sentinel and stops reading immediately — bytes past the limit are never
+    consumed from the request stream.
+    """
+
+
+class _SizeLimitedFileTarget(FileTarget):
+    """``FileTarget`` subclass that aborts mid-stream on size-limit violation."""
+
+    def __init__(self, filename: str, *, max_bytes: int) -> None:
+        super().__init__(filename)
+        self._max_bytes = max_bytes
+        self.size_bytes: int = 0
+        # Redeclare with an explicit type so mypy can resolve it (FileTarget sets it
+        # to None in __init__ but the stubs don't expose its type).
+        self._fd = None  # type: ignore[assignment]
+
+    def on_data_received(self, chunk: bytes) -> None:
+        self.size_bytes += len(chunk)
+        if self.size_bytes > self._max_bytes:
+            # Close the file descriptor before aborting so the caller can safely
+            # unlink the temp file on all platforms.
+            fd = self._fd  # type: ignore[has-type]
+            if fd is not None:
+                fd.close()
+                self._fd = None  # type: ignore[has-type]
+            raise _SizeLimitExceeded()
+        super().on_data_received(chunk)
+
+
 @router.post(
     "/external-results/artifact",
     response_model=ArtifactResponse,
     status_code=status.HTTP_201_CREATED,
 )
 async def upload_artifact(
-    case_result_id: str = Form(...),
-    kind: str = Form(...),
-    filename: str = Form(...),
-    file: UploadFile = File(...),
+    request: Request,
     db: AsyncSession = Depends(get_db),
     token: RunnerToken = Depends(require_runner_scope(_WRITE_SCOPE)),
 ) -> ArtifactResponse:
@@ -448,11 +479,75 @@ async def upload_artifact(
     - ``filename``       — original filename including extension.
     - ``file``           — binary body; its ``Content-Type`` part header is used as the
                            artifact content-type.
+
+    The handler uses ``streaming-form-data`` to parse the multipart body chunk by
+    chunk.  As soon as the cumulative byte count of the ``file`` part exceeds
+    ``BGSTM_ARTIFACT_MAX_BYTES``, a ``_SizeLimitExceeded`` sentinel is raised
+    inside the part-data callback, the stream loop exits immediately (no further
+    reads), the temp file is deleted, and 413 is returned.
     """
+    max_bytes: int = settings.BGSTM_ARTIFACT_MAX_BYTES
+
+    # Create the temp file upfront; ``FileTarget`` will reopen it via ``on_start``.
+    fd, tmp_path = tempfile.mkstemp(prefix="bgstm_artifact_")
+    os.close(fd)
+
+    case_result_id_target = ValueTarget()
+    kind_target = ValueTarget()
+    filename_target = ValueTarget()
+    file_target = _SizeLimitedFileTarget(tmp_path, max_bytes=max_bytes)
+
+    parser = StreamingFormDataParser(headers=request.headers)
+    parser.register("case_result_id", case_result_id_target)
+    parser.register("kind", kind_target)
+    parser.register("filename", filename_target)
+    parser.register("file", file_target)
+
+    try:
+        async for chunk in request.stream():
+            parser.data_received(chunk)
+    except _SizeLimitExceeded:
+        _safe_unlink(tmp_path)
+        raise HTTPException(
+            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
+            detail={
+                "code": "artifact.too_large",
+                "message": f"Artifact exceeds the maximum allowed size of {max_bytes} bytes.",
+                "details": None,
+            },
+        )
+    except Exception as exc:
+        _safe_unlink(tmp_path)
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail={
+                "code": "validation_error",
+                "message": f"Failed to parse multipart body: {exc}",
+                "details": None,
+            },
+        ) from exc
+
+    # --- Decode text fields ---
+    try:
+        case_result_id = case_result_id_target.value.decode("utf-8")
+        kind = kind_target.value.decode("utf-8")
+        filename = filename_target.value.decode("utf-8")
+    except UnicodeDecodeError as exc:
+        _safe_unlink(tmp_path)
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail={
+                "code": "validation_error",
+                "message": f"Multipart text field contains invalid UTF-8: {exc}",
+                "details": None,
+            },
+        ) from exc
+
     # --- Validate kind ---
     try:
         artifact_kind = ArtifactKind(kind)
     except ValueError:
+        _safe_unlink(tmp_path)
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
             detail={
@@ -463,11 +558,9 @@ async def upload_artifact(
         )
 
     # --- Sanitize and validate filename (path-traversal defense) ---
-    # Reject if the filename differs from its own basename OR fails the allowlist regex.
-    # This catches directory components (`../`, `subdir/`, `/etc/`) as well as
-    # dangerous characters (null bytes, backslashes, spaces, etc.).
     safe_filename = os.path.basename(filename)
     if safe_filename != filename or not _SAFE_FILENAME_RE.fullmatch(safe_filename):
+        _safe_unlink(tmp_path)
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
             detail={
@@ -484,6 +577,7 @@ async def upload_artifact(
     try:
         case_result_uuid = _uuid_module.UUID(case_result_id)
     except ValueError:
+        _safe_unlink(tmp_path)
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
             detail={
@@ -496,6 +590,7 @@ async def upload_artifact(
     # --- Verify the case result exists ---
     case_result = await get_case_result(db, case_result_uuid)
     if case_result is None:
+        _safe_unlink(tmp_path)
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail={
@@ -506,10 +601,12 @@ async def upload_artifact(
         )
 
     # --- Derive content-type from the upload part header ---
-    content_type: str = (file.content_type or "application/octet-stream").split(";")[0].strip().lower()
+    raw_ct = getattr(file_target, "multipart_content_type", None) or "application/octet-stream"
+    content_type: str = raw_ct.split(";")[0].strip().lower()
 
     # --- Validate content-type (bypass for kind=other) ---
     if artifact_kind != ArtifactKind.other and content_type not in _ALLOWED_CONTENT_TYPES:
+        _safe_unlink(tmp_path)
         raise HTTPException(
             status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
             detail={
@@ -522,35 +619,6 @@ async def upload_artifact(
             },
         )
 
-    # --- Stream to a temp file, enforcing max size ---
-    max_bytes: int = settings.BGSTM_ARTIFACT_MAX_BYTES
-    fd, tmp_path = tempfile.mkstemp(prefix="bgstm_artifact_")
-    total_bytes = 0
-    try:
-        with os.fdopen(fd, "wb") as fp:
-            while True:
-                chunk = await file.read(_ARTIFACT_CHUNK_SIZE)
-                if not chunk:
-                    break
-                fp.write(chunk)
-                total_bytes += len(chunk)
-                if total_bytes > max_bytes:
-                    # Partial data written — clean up and reject
-                    raise HTTPException(
-                        status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
-                        detail={
-                            "code": "artifact.too_large",
-                            "message": (f"Artifact exceeds the maximum allowed size of {max_bytes} bytes."),
-                            "details": None,
-                        },
-                    )
-    except HTTPException:
-        _safe_unlink(tmp_path)
-        raise
-    except Exception:
-        _safe_unlink(tmp_path)
-        raise
-
     # --- Persist via storage backend ---
     storage = get_storage()
     storage_key = f"{case_result_id}/{_uuid_module.uuid4().hex}/{safe_filename}"

backend/requirements.txt

@@ -17,6 +17,9 @@ PyJWT==2.11.0
 bcrypt>=4.0.0
 email-validator==2.3.0
 reportlab==4.4.10
+# Streaming multipart parser — used by upload_artifact for true mid-wire size-limit abort.
+# Pinned to <2.0 to avoid breaking API changes; Dependabot will bump the minor.
+streaming-form-data>=1.16,<2.0
 
 # Optional dependencies for LLM embeddings
 # Uncomment the lines below if using LLM-based similarity

backend/tests/test_external_results_artifact.py

@@ -214,17 +214,14 @@ async def test_upload_screenshot_returns_201(self, db_session, write_token, tmp_
 class TestSizeEnforcement:
     @pytest.mark.asyncio
     async def test_oversized_file_returns_413_and_cleans_up(self, db_session, write_token, tmp_path, monkeypatch):
-        """413 on oversized upload; Starlette-managed temp file is cleaned up; no DB record created.
+        """413 on oversized upload; streaming abort is now the actual behavior.
 
-        Note: FastAPI/Starlette fully parses the multipart body via python-multipart
-        **before** the handler is invoked — by the time the chunk loop runs, ``file``
-        is already a SpooledTemporaryFile containing the entire upload.  The 413
-        enforcement happens post-buffer (reading from the spooled file in chunks),
-        not mid-wire.  True in-stream early-abort is a follow-up improvement; deploy
-        behind a reverse-proxy ``client_max_body_size`` for first-line DoS protection.
+        The handler uses streaming-form-data to parse the multipart body and raises
+        _SizeLimitExceeded mid-stream as soon as cumulative bytes exceed
+        BGSTM_ARTIFACT_MAX_BYTES — bytes past the limit are never read from the
+        connection.
 
-        This test verifies:
-          - 413 is returned when the spooled content exceeds BGSTM_ARTIFACT_MAX_BYTES.
+        This test exercises the cleanup / DB-row / artifact-dir guarantees post-abort:
           - The handler's own temp file (bgstm_artifact_*) is cleaned up on 413.
           - No artifact row is written to the DB.
           - No file is left in the artifacts directory.
@@ -281,12 +278,130 @@ def _fake_mkstemp(*args, **kwargs):
         artifact_files = list(tmp_path.rglob("*"))
         assert artifact_files == [], f"Unexpected files in artifacts dir: {artifact_files}"
 
+    @pytest.mark.asyncio
+    async def test_oversized_upload_aborts_stream_without_reading_full_body(
+        self, db_session, write_token, tmp_path, monkeypatch
+    ):
+        """Verify that bytes past BGSTM_ARTIFACT_MAX_BYTES are NEVER read from the
+        request stream. This is the load-bearing test for #320 — it fails against
+        the old buffer-then-reject implementation and passes only when streaming
+        abort is correctly wired.
+        """
+        _token_model, plaintext = write_token
+        monkeypatch.setattr(settings, "BGSTM_ARTIFACT_MAX_BYTES", 1024)
+        monkeypatch.setattr(settings, "BGSTM_ARTIFACTS_DIR", str(tmp_path))
+        monkeypatch.setattr(settings, "BGSTM_ARTIFACT_URL_PREFIX", "http://testserver/artifacts")
+        monkeypatch.setattr(settings, "BGSTM_STORAGE_BACKEND", "local")
+
+        with TestClient(app) as sync_client:
+            session_id = _create_session(sync_client, plaintext)
+            case_result_id = _create_case_result(sync_client, plaintext, session_id)
+
+        # Build a multipart body: small text fields + 100 KiB file (100x the 1024 limit).
+        LIMIT = 1024
+        FILE_SIZE = 100 * LIMIT  # 100 KiB — well past the limit
+        boundary = b"bgstmtestboundary"
+        file_data = b"X" * FILE_SIZE
+
+        def _field_part(name: str, value: str) -> bytes:
+            return (
+                b"--"
+                + boundary
+                + b"\r\n"
+                + b'Content-Disposition: form-data; name="'
+                + name.encode()
+                + b'"\r\n'
+                + b"\r\n"
+                + value.encode()
+                + b"\r\n"
+            )
+
+        full_body = (
+            _field_part("case_result_id", case_result_id)
+            + _field_part("kind", "screenshot")
+            + _field_part("filename", "big.png")
+            + b"--"
+            + boundary
+            + b"\r\n"
+            + b'Content-Disposition: form-data; name="file"; filename="big.png"\r\n'
+            + b"Content-Type: image/png\r\n"
+            + b"\r\n"
+            + file_data
+            + b"\r\n"
+            + b"--"
+            + boundary
+            + b"--\r\n"
+        )
+        body_size = len(full_body)
+
+        # Track how many bytes our generator has yielded — each yield corresponds
+        # to one receive() call from the ASGI server (via ASGITransport).
+        bytes_yielded = 0
+        CHUNK_SIZE = 4096
+
+        async def streaming_body():
+            nonlocal bytes_yielded
+            for i in range(0, len(full_body), CHUNK_SIZE):
+                chunk = full_body[i : i + CHUNK_SIZE]
+                bytes_yielded += len(chunk)
+                yield chunk
+
+        from httpx import ASGITransport, AsyncClient
+
+        async with AsyncClient(transport=ASGITransport(app=app), base_url="http://testserver") as client:
+            resp = await client.post(
+                "/api/v1/external-results/artifact",
+                content=streaming_body(),
+                headers={
+                    "Authorization": f"Bearer {plaintext}",
+                    "Content-Type": f"multipart/form-data; boundary={boundary.decode()}",
+                },
+            )
+
+        assert resp.status_code == 413, resp.text
+        assert resp.json()["detail"]["code"] == "artifact.too_large"
+
+        # The server must have stopped reading well before the full body was sent.
+        # Allow generous slack (limit + 256 KiB for framing + chunks), but assert
+        # well below total body size (100 KiB file ≫ slack).
+        assert bytes_yielded < body_size // 2, (
+            f"Server read {bytes_yielded} of {body_size} body bytes — "
+            "streaming abort is not actually aborting; bytes past the limit are still being read."
+        )
+
 
 # ---------------------------------------------------------------------------
-# Content-type enforcement (415)
+# Malformed multipart body (422 + code=validation_error)
 # ---------------------------------------------------------------------------
 
 
+class TestMalformedMultipart:
+    @pytest.mark.asyncio
+    async def test_malformed_multipart_returns_422(self, db_session, write_token, tmp_path, monkeypatch):
+        """Posting a body that is not valid multipart must return 422 with
+        code=validation_error, never 500.
+        """
+        _token_model, plaintext = write_token
+        monkeypatch.setattr(settings, "BGSTM_ARTIFACTS_DIR", str(tmp_path))
+        monkeypatch.setattr(settings, "BGSTM_STORAGE_BACKEND", "local")
+
+        from httpx import ASGITransport, AsyncClient
+
+        async with AsyncClient(transport=ASGITransport(app=app), base_url="http://testserver") as client:
+            resp = await client.post(
+                "/api/v1/external-results/artifact",
+                content=b"this is not multipart data at all \x00\x01\x02",
+                headers={
+                    "Authorization": f"Bearer {plaintext}",
+                    # Valid multipart content-type but body is garbage
+                    "Content-Type": "multipart/form-data; boundary=correctboundary",
+                },
+            )
+
+        assert resp.status_code == 422, resp.text
+        assert resp.json()["detail"]["code"] == "validation_error"
+
+
 class TestContentTypeEnforcement:
     @pytest.mark.asyncio
     async def test_disallowed_content_type_returns_415(self, db_session, write_token, tmp_path, monkeypatch):