From c4c61e063be0134111281f4966e854fba3895ba9 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Tue, 12 May 2026 14:48:55 -0400 Subject: [PATCH 1/2] feat(files): stream uploads to disk instead of buffering Previously /api/files/create read the entire UploadFile into memory before writing to disk, peaking at ~100 MB for a 100 MB upload. Now we drain the upload in 1 MiB chunks straight to a `.part` temp file and atomically rename on success, so peak memory stays bounded regardless of file size and a failed upload never leaves a half-written file at the final path. - `OSFileSystem.stream_create_file` does the chunked write with atomic rename and cleanup on failure - `parse_multipart_request` now hands callers un-read `UploadFile` handles (instead of pre-read bytes), so streaming is possible without giving up the small-payload `.read()` path - Directories and the default-template notebook still go through the in-memory `create_file_or_directory` path; only real file content streams --- marimo/_server/api/endpoints/file_explorer.py | 32 +++- marimo/_server/api/utils.py | 37 ++-- marimo/_server/files/os_file_system.py | 153 +++++++++++++++-- marimo/_server/models/files.py | 11 +- packages/openapi/api.yaml | 6 +- packages/openapi/src/api.ts | 5 + .../api/endpoints/test_file_explorer.py | 21 +++ tests/_server/api/test_api_utils.py | 17 +- tests/_server/files/test_os_file_system.py | 160 ++++++++++++++++++ 9 files changed, 397 insertions(+), 45 deletions(-) diff --git a/marimo/_server/api/endpoints/file_explorer.py b/marimo/_server/api/endpoints/file_explorer.py index 391779c9194..4701155f319 100644 --- a/marimo/_server/api/endpoints/file_explorer.py +++ b/marimo/_server/api/endpoints/file_explorer.py @@ -4,11 +4,15 @@ from typing import TYPE_CHECKING from starlette.authentication import requires +from starlette.exceptions import HTTPException from marimo import _loggers from marimo._server.api.deps import AppState from marimo._server.api.utils import parse_multipart_request, parse_request -from marimo._server.files.os_file_system import OSFileSystem +from marimo._server.files.os_file_system import ( + OSFileSystem, + UploadTooLargeError, +) from marimo._server.models.files import ( FileCopyRequest, FileCopyResponse, @@ -120,16 +124,26 @@ async def create_file_or_directory( $ref: "#/components/schemas/FileCreateResponse" """ try: - parsed = await parse_multipart_request( + async with parse_multipart_request( request, FileCreateMultipartRequest - ) - info = file_system.create_file_or_directory( - parsed.body.path, - parsed.body.type, - parsed.body.name, - parsed.files.get("file"), - ) + ) as parsed: + upload = parsed.files.get("file") + # Stream when there's actual file content; the in-memory create + # path still handles directories and the default-template notebook. + if upload is not None and parsed.body.type in ("file", "notebook"): + info = await file_system.stream_create_file( + parsed.body.path, parsed.body.name, upload + ) + else: + info = file_system.create_file_or_directory( + parsed.body.path, parsed.body.type, parsed.body.name, None + ) return FileCreateResponse(success=True, info=info) + except UploadTooLargeError as e: + LOGGER.warning(f"Rejected oversize upload: {e}") + # 413 is the standard "payload too large" response; surface it + # explicitly rather than burying it in a 200 + success=False. + raise HTTPException(status_code=413, detail=str(e)) from e except Exception as e: LOGGER.error(f"Error creating file or directory: {e}") return FileCreateResponse(success=False, message=str(e)) diff --git a/marimo/_server/api/utils.py b/marimo/_server/api/utils.py index 57511497187..89a74845fe6 100644 --- a/marimo/_server/api/utils.py +++ b/marimo/_server/api/utils.py @@ -5,6 +5,7 @@ import subprocess import sys import webbrowser +from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path from shutil import which @@ -25,6 +26,9 @@ from marimo._utils.parse_dataclass import parse_raw if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from starlette.datastructures import UploadFile from starlette.requests import Request from marimo._session.session import Session @@ -44,20 +48,31 @@ async def parse_request( @dataclass class MultipartRequest(Generic[S]): - """Result of parsing a multipart/form-data request body.""" + """Result of parsing a multipart/form-data request body. + + `files` carries the raw `UploadFile` objects (not yet read) so callers + can stream chunks via `.read(size)`. For small payloads, callers can + simply `await upload.read()` to materialize the whole body. + """ body: S - files: dict[str, bytes] + files: dict[str, UploadFile] +@asynccontextmanager async def parse_multipart_request( request: Request, cls: type[S] -) -> MultipartRequest[S]: - """Parse a multipart/form-data body into a msgspec.Struct + file bytes. +) -> AsyncIterator[MultipartRequest[S]]: + """Parse a multipart/form-data body into a msgspec.Struct + uploads. String form fields are validated against `cls`. File upload parts are - read fully into memory and returned in `files`, keyed by form-field - name (callers look them up explicitly rather than via the struct). + yielded as `UploadFile` objects (un-read) in `files`, keyed by + form-field name, so callers can stream them to disk instead of buffering. + + Used as an async context manager so the underlying form (and its + spooled temp files / fds) are closed automatically when the caller + is done — `UploadFile` parts remain readable for the entire body of + the `async with` block. Raises msgspec.ValidationError if required string fields are missing or invalid. @@ -66,18 +81,16 @@ async def parse_multipart_request( # without starlette (e.g. pyodide). from starlette.datastructures import UploadFile - # Use as an async context manager so any spooled temp files backing - # UploadFile parts are closed after parsing. async with request.form() as form: string_payload: dict[str, Any] = {} - files: dict[str, bytes] = {} + files: dict[str, UploadFile] = {} for key, value in form.multi_items(): if isinstance(value, UploadFile): - files[key] = await value.read() + files[key] = value elif isinstance(value, str): string_payload[key] = value - body = msgspec.convert(string_payload, cls, strict=False) - return MultipartRequest(body=body, files=files) + body = msgspec.convert(string_payload, cls, strict=False) + yield MultipartRequest(body=body, files=files) @runtime_checkable diff --git a/marimo/_server/files/os_file_system.py b/marimo/_server/files/os_file_system.py index e835d6e0879..37c8e1823e7 100644 --- a/marimo/_server/files/os_file_system.py +++ b/marimo/_server/files/os_file_system.py @@ -8,9 +8,10 @@ import re import shutil import subprocess +import tempfile from collections import deque from pathlib import Path -from typing import Literal +from typing import Literal, Protocol from marimo import _loggers from marimo._server.files.file_system import FileSystem @@ -36,6 +37,33 @@ "..", ] +# 1 MiB. Large enough to amortize syscall overhead, small enough to keep +# peak memory bounded when streaming. +_STREAM_CHUNK_SIZE = 1024 * 1024 + +# Hard cap on streamed uploads. Streaming removes the implicit OOM ceiling +# that buffered uploads had, so without a cap an authenticated client could +# exhaust disk. 1 GiB covers normal notebook-data use cases with margin. +MAX_UPLOAD_BYTES = 1024 * 1024 * 1024 + + +class UploadTooLargeError(ValueError): + """Raised when a streamed upload exceeds `MAX_UPLOAD_BYTES`. + + Separate type (vs. a bare `ValueError`) so the HTTP layer can map it + to a 413 response instead of the generic error path. + """ + + +class AsyncByteSource(Protocol): + """Anything that can be drained chunk-by-chunk into a file. + + Starlette's `UploadFile` satisfies this; so does any object exposing + an async `read(size)` returning bytes. + """ + + async def read(self, size: int = -1, /) -> bytes: ... + class OSFileSystem(FileSystem): def get_root(self) -> str: @@ -133,33 +161,32 @@ def open_file(self, path: str, encoding: str | None = None) -> str | bytes: except UnicodeDecodeError: return file_path.read_bytes() - def create_file_or_directory( - self, - path: str, - file_type: Literal["file", "directory", "notebook"], - name: str, - contents: bytes | None, - ) -> FileInfo: + @staticmethod + def _validate_create_name(name: str) -> None: + """Reject names that are empty, reserved, or traverse out of the + parent. Centralized so HTTP, WASM, and streaming paths all share it. + """ if name in DISALLOWED_NAMES: raise ValueError( f"Cannot create file or directory with name {name}" ) if name.strip() == "": raise ValueError("Cannot create file or directory with empty name") - # Names that traverse out of `path` or escape via separators are - # rejected. Validation belongs here (not in the endpoint) so every - # caller of OSFileSystem — HTTP, WASM bridge, scripts — is covered. - if ( - "/" in name - or "\\" in name - or "\x00" in name - or name in (".", "..") - ): + if "/" in name or "\\" in name or "\x00" in name: raise ValueError( f"Invalid name {name!r}: must not contain path separators " "or refer to a parent directory" ) + def create_file_or_directory( + self, + path: str, + file_type: Literal["file", "directory", "notebook"], + name: str, + contents: bytes | None, + ) -> FileInfo: + self._validate_create_name(name) + full_path = Path(path) / name full_path = _generate_unique_path(full_path) @@ -192,6 +219,71 @@ def create_file_or_directory( ), ).file + async def stream_create_file( + self, + path: str, + name: str, + source: AsyncByteSource, + ) -> FileInfo: + """Stream-write an uploaded file to disk, chunk by chunk. + + Avoids loading the full payload into memory (the HTTP multipart + path can otherwise buffer 100 MB at once). Writes to a ``.part`` + temp file and atomically renames on success so a failed upload + doesn't leave a half-written file at the final path. + """ + self._validate_create_name(name) + + parent = Path(path) + os.makedirs(parent, exist_ok=True) + + # Atomically claim the destination with O_CREAT|O_EXCL so concurrent + # uploads can't both pick the same numbered suffix and clobber each + # other between `_generate_unique_path` and the final rename. + full_path = _claim_unique_path(parent / name) + + # `NamedTemporaryFile` gives us a guaranteed-unique sibling path for + # the in-progress `.part` file (same reasoning as above). + tmp = tempfile.NamedTemporaryFile( + dir=full_path.parent, + prefix=full_path.name + ".", + suffix=".part", + delete=False, + ) + tmp_path = tmp.name + try: + # Sync writes are bounded to ~1 MiB per chunk, with an `await` + # in between; event loop blockage is brief and an async file + # library would only add a dependency for marginal gain. + written = 0 + with tmp: + while chunk := await source.read(_STREAM_CHUNK_SIZE): + written += len(chunk) + if written > MAX_UPLOAD_BYTES: + raise UploadTooLargeError( + f"Upload exceeds maximum size of " + f"{MAX_UPLOAD_BYTES} bytes" + ) + tmp.write(chunk) + # Replaces our empty marker at `full_path`. Atomic on POSIX + # and Windows (Python 3.3+). + os.replace(tmp_path, full_path) + except BaseException: + # Clean up both the `.part` and the reserved empty marker. + # If `os.replace` already succeeded, the marker is gone and + # `FileNotFoundError` is the expected outcome there. + for p in (tmp_path, str(full_path)): + try: + os.unlink(p) + except FileNotFoundError: + pass + raise + + # Use the metadata-only helper: `get_details` would re-read the + # file contents (and base64-encode binary), defeating the point of + # streaming for large uploads. + return self._get_file_info(str(full_path)) + def delete_file_or_directory(self, path: str) -> bool: if os.path.isdir(path): safe_rmtree(path) @@ -485,6 +577,33 @@ def _generate_unique_path(new_path: str | Path) -> Path: i += 1 +def _claim_unique_path(target: Path) -> Path: + """Like `_generate_unique_path`, but atomically reserves the chosen name. + + Opens with O_CREAT|O_EXCL so two concurrent callers can never end up + with the same path — whichever loses the race sees `FileExistsError` + and tries the next numbered suffix. Returns the claimed path (an empty + file at that location); the caller is responsible for writing into it + (or replacing it). + """ + name_without_extension = target.stem + extension = target.suffix + candidate = target + i = 0 + while True: + try: + fd = os.open( + candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644 + ) + os.close(fd) + return candidate + except FileExistsError: + i += 1 + candidate = target.parent / ( + f"{name_without_extension}_{i}{extension}" + ) + + def _is_allowed_paths(path: str | Path, new_path: str | Path) -> bool: file_name = os.path.basename(new_path) if file_name in DISALLOWED_NAMES or not file_name.strip(): diff --git a/marimo/_server/models/files.py b/marimo/_server/models/files.py index d4cd57b26af..80f04e0b36a 100644 --- a/marimo/_server/models/files.py +++ b/marimo/_server/models/files.py @@ -58,7 +58,13 @@ class FileCreateRequest(msgspec.Struct, rename="camel"): class FileCreateMultipartRequest(msgspec.Struct, rename="camel"): - """multipart/form-data body for POST /api/files/create.""" + """multipart/form-data body for POST /api/files/create. + + Schema-only: this struct exists to describe the multipart shape in + OpenAPI. At runtime, the endpoint reads the string fields from + ``MultipartRequest.body`` and the uploaded bytes from + ``MultipartRequest.files["file"]`` — ``body.file`` is never populated. + """ # The path where to create the file or directory path: str @@ -68,6 +74,9 @@ class FileCreateMultipartRequest(msgspec.Struct, rename="camel"): name: str # The raw file bytes (optional). When omitted, an empty file is created # (or, for 'notebook' type, a default notebook template). + # NOTE: this field is OpenAPI-only — see class docstring. The + # ``format: binary`` annotation makes the generated spec emit a proper + # file-upload schema rather than a base64 string. file: Annotated[ str | None, msgspec.Meta(extra_json_schema={"format": "binary"}) ] = None diff --git a/packages/openapi/api.yaml b/packages/openapi/api.yaml index e45073164dd..79130d77ac1 100644 --- a/packages/openapi/api.yaml +++ b/packages/openapi/api.yaml @@ -1665,7 +1665,11 @@ components: title: FileCopyResponse type: object FileCreateMultipartRequest: - description: multipart/form-data body for POST /api/files/create. + description: "multipart/form-data body for POST /api/files/create.\n\n Schema-only:\ + \ this struct exists to describe the multipart shape in\n OpenAPI. At runtime,\ + \ the endpoint reads the string fields from\n ``MultipartRequest.body``\ + \ and the uploaded bytes from\n ``MultipartRequest.files[\"file\"]`` \u2014\ + \ ``body.file`` is never populated." properties: file: anyOf: diff --git a/packages/openapi/src/api.ts b/packages/openapi/src/api.ts index 7fc0cabfa0e..93a7125cbc0 100644 --- a/packages/openapi/src/api.ts +++ b/packages/openapi/src/api.ts @@ -4359,6 +4359,11 @@ export interface components { /** * FileCreateMultipartRequest * @description multipart/form-data body for POST /api/files/create. + * + * Schema-only: this struct exists to describe the multipart shape in + * OpenAPI. At runtime, the endpoint reads the string fields from + * ``MultipartRequest.body`` and the uploaded bytes from + * ``MultipartRequest.files["file"]`` — ``body.file`` is never populated. */ FileCreateMultipartRequest: { /** diff --git a/tests/_server/api/endpoints/test_file_explorer.py b/tests/_server/api/endpoints/test_file_explorer.py index e64b4b78ff8..f54869bbca3 100644 --- a/tests/_server/api/endpoints/test_file_explorer.py +++ b/tests/_server/api/endpoints/test_file_explorer.py @@ -179,6 +179,27 @@ def test_create_rejects_path_traversal_name(client: TestClient) -> None: assert not os.path.exists(parent_path) +def test_create_returns_413_when_upload_exceeds_size_cap( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + # Shrink the cap so a small payload trips it. + monkeypatch.setattr( + "marimo._server.files.os_file_system.MAX_UPLOAD_BYTES", 4 + ) + response = client.post( + "/api/files/create", + headers=HEADERS, + data={ + "path": test_dir, + "type": "file", + "name": "too_big.bin", + }, + files={"file": ("too_big.bin", b"way too much data")}, + ) + assert response.status_code == 413, response.text + assert not os.path.exists(os.path.join(test_dir, "too_big.bin")) + + def test_update_file(client: TestClient) -> None: response = client.post( "/api/files/update", diff --git a/tests/_server/api/test_api_utils.py b/tests/_server/api/test_api_utils.py index 5f4a09a4e60..1ff8a0e03ef 100644 --- a/tests/_server/api/test_api_utils.py +++ b/tests/_server/api/test_api_utils.py @@ -23,16 +23,19 @@ class _SampleForm(msgspec.Struct): def _build_app(captured: dict[str, object]) -> TestClient: async def endpoint(request: Request) -> JSONResponse: - parsed = await parse_multipart_request(request, _SampleForm) - captured["body"] = parsed.body - captured["files"] = parsed.files + async with parse_multipart_request(request, _SampleForm) as parsed: + captured["body"] = parsed.body + captured["files"] = dict(parsed.files) + upload = parsed.files.get("upload") + if upload is not None: + captured["upload_bytes"] = await upload.read() return JSONResponse({"ok": True}) app = Starlette(routes=[Route("/test", endpoint, methods=["POST"])]) return TestClient(app) -def test_parse_multipart_request_strings_and_file_bytes() -> None: +def test_parse_multipart_request_strings_and_file_upload() -> None: captured: dict[str, object] = {} client = _build_app(captured) response = client.post( @@ -45,7 +48,11 @@ def test_parse_multipart_request_strings_and_file_bytes() -> None: assert isinstance(body, _SampleForm) assert body.name == "marimo" assert body.count == 42 - assert captured["files"] == {"upload": b"\x00\x01\x02\xff"} + files = captured["files"] + assert isinstance(files, dict) + assert set(files.keys()) == {"upload"} + # File handle is returned un-read; bytes loaded inside the endpoint above. + assert captured["upload_bytes"] == b"\x00\x01\x02\xff" def test_parse_multipart_request_omitted_file_yields_empty_dict() -> None: diff --git a/tests/_server/files/test_os_file_system.py b/tests/_server/files/test_os_file_system.py index 9b7f0e508d3..49482497139 100644 --- a/tests/_server/files/test_os_file_system.py +++ b/tests/_server/files/test_os_file_system.py @@ -8,6 +8,8 @@ from marimo._server.files.os_file_system import ( OSFileSystem, + UploadTooLargeError, + _claim_unique_path, _generate_unique_path, _is_allowed_paths, ) @@ -133,6 +135,164 @@ def test_create_rejects_path_traversal( assert not (test_dir.parent / "escaped.txt").exists() +class _ChunkedSource: + """Minimal async byte source emitting predetermined chunks. + + If ``fail_after`` is set, raises ``RuntimeError`` after that many reads + have returned data — useful for simulating mid-stream upload failures. + """ + + def __init__( + self, chunks: list[bytes], *, fail_after: int | None = None + ) -> None: + self._chunks = list(chunks) + self._fail_after = fail_after + self._reads = 0 + + async def read(self, size: int = -1, /) -> bytes: + if self._fail_after is not None and self._reads >= self._fail_after: + raise RuntimeError("stream failed") + if not self._chunks: + return b"" + # Honor size by returning the next chunk if it fits, else slicing. + chunk = self._chunks[0] + if size < 0 or len(chunk) <= size: + out = self._chunks.pop(0) + else: + out, self._chunks[0] = chunk[:size], chunk[size:] + self._reads += 1 + return out + + +def _part_files(directory: Path) -> list[Path]: + return list(directory.glob("*.part")) + + +async def test_stream_create_file_writes_chunks( + test_dir: Path, fs: OSFileSystem +) -> None: + source = _ChunkedSource([b"hello ", b"streamed ", b"world"]) + info = await fs.stream_create_file(str(test_dir), "out.bin", source) + out_path = test_dir / "out.bin" + assert out_path.read_bytes() == b"hello streamed world" + assert info.path == str(out_path) + assert _part_files(test_dir) == [] + + +async def test_stream_create_file_writes_empty_file( + test_dir: Path, fs: OSFileSystem +) -> None: + info = await fs.stream_create_file( + str(test_dir), "empty.bin", _ChunkedSource([]) + ) + out_path = test_dir / "empty.bin" + assert out_path.exists() + assert out_path.read_bytes() == b"" + assert info.path == str(out_path) + assert _part_files(test_dir) == [] + + +async def test_stream_create_file_generates_unique_path( + test_dir: Path, fs: OSFileSystem +) -> None: + (test_dir / "dup.bin").write_bytes(b"existing") + info = await fs.stream_create_file( + str(test_dir), "dup.bin", _ChunkedSource([b"new"]) + ) + assert info.path == str(test_dir / "dup_1.bin") + assert (test_dir / "dup.bin").read_bytes() == b"existing" + assert (test_dir / "dup_1.bin").read_bytes() == b"new" + + +@pytest.mark.parametrize( + "name", + [ + "../escape.bin", + "../../escape.bin", + "subdir/nested.bin", + "..", + ".", + "with\\backslash.bin", + "embed\x00null.bin", + "", + " ", + ], +) +async def test_stream_create_file_rejects_traversal( + test_dir: Path, fs: OSFileSystem, name: str +) -> None: + with pytest.raises(ValueError): + await fs.stream_create_file( + str(test_dir), name, _ChunkedSource([b"x"]) + ) + # Nothing should have been written into the parent or test_dir. + assert not (test_dir.parent / "escape.bin").exists() + assert _part_files(test_dir) == [] + + +async def test_stream_create_file_cleans_up_on_immediate_failure( + test_dir: Path, fs: OSFileSystem +) -> None: + source = _ChunkedSource([], fail_after=0) + with pytest.raises(RuntimeError): + await fs.stream_create_file(str(test_dir), "no_file.bin", source) + assert not (test_dir / "no_file.bin").exists() + assert _part_files(test_dir) == [] + + +async def test_stream_create_file_cleans_up_on_mid_stream_failure( + test_dir: Path, fs: OSFileSystem +) -> None: + # The realistic case: a few chunks already written when the source + # raises (network drop, client disconnect, etc). + source = _ChunkedSource([b"first ", b"second ", b"third"], fail_after=2) + with pytest.raises(RuntimeError): + await fs.stream_create_file(str(test_dir), "partial.bin", source) + assert not (test_dir / "partial.bin").exists() + assert _part_files(test_dir) == [] + + +async def test_stream_create_file_enforces_size_cap( + test_dir: Path, fs: OSFileSystem, monkeypatch: pytest.MonkeyPatch +) -> None: + # Shrink the cap so the test is fast; the production value is 1 GiB. + monkeypatch.setattr( + "marimo._server.files.os_file_system.MAX_UPLOAD_BYTES", 8 + ) + source = _ChunkedSource([b"aaaa", b"bbbb", b"too much"]) + # Raises the specific subclass so the endpoint can map it to a 413. + with pytest.raises(UploadTooLargeError, match="exceeds maximum size"): + await fs.stream_create_file(str(test_dir), "big.bin", source) + assert not (test_dir / "big.bin").exists() + assert _part_files(test_dir) == [] + + +async def test_stream_create_file_claims_path_atomically( + test_dir: Path, fs: OSFileSystem +) -> None: + # Simulate the TOCTOU race: a concurrent caller creates the same name + # while we're streaming. With atomic O_CREAT|O_EXCL reservation we + # must NOT overwrite the concurrent file. + (test_dir / "race.bin").write_bytes(b"existing") + info = await fs.stream_create_file( + str(test_dir), "race.bin", _ChunkedSource([b"new"]) + ) + assert info.path == str(test_dir / "race_1.bin") + assert (test_dir / "race.bin").read_bytes() == b"existing" + assert (test_dir / "race_1.bin").read_bytes() == b"new" + + +def test_claim_unique_path_atomic_reservation(test_dir: Path) -> None: + # First claim takes the bare name and reserves an empty file. + claimed_a = _claim_unique_path(test_dir / "claim.bin") + assert claimed_a == test_dir / "claim.bin" + assert claimed_a.exists() + # Second claim must NOT reuse the reserved path, even though it's empty. + claimed_b = _claim_unique_path(test_dir / "claim.bin") + assert claimed_b == test_dir / "claim_1.bin" + assert claimed_b.exists() + + def test_list_files(test_dir: Path, fs: OSFileSystem) -> None: test_create_file(test_dir, fs) test_create_directory(test_dir, fs) From 6bdc9ac35f26a985e3943e7957d4eab156247fec Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Sat, 16 May 2026 14:47:57 -0700 Subject: [PATCH 2/2] cleanup comments --- marimo/_server/api/endpoints/file_explorer.py | 6 +-- marimo/_server/api/utils.py | 22 +++------- marimo/_server/files/os_file_system.py | 40 ++++++------------- marimo/_server/models/files.py | 7 +--- tests/_server/api/test_api_utils.py | 1 - tests/_server/files/test_os_file_system.py | 22 ++++------ 6 files changed, 29 insertions(+), 69 deletions(-) diff --git a/marimo/_server/api/endpoints/file_explorer.py b/marimo/_server/api/endpoints/file_explorer.py index 4701155f319..8b3c24557ce 100644 --- a/marimo/_server/api/endpoints/file_explorer.py +++ b/marimo/_server/api/endpoints/file_explorer.py @@ -128,8 +128,8 @@ async def create_file_or_directory( request, FileCreateMultipartRequest ) as parsed: upload = parsed.files.get("file") - # Stream when there's actual file content; the in-memory create - # path still handles directories and the default-template notebook. + # Directories and the default-template notebook take the + # in-memory path; only real file content streams. if upload is not None and parsed.body.type in ("file", "notebook"): info = await file_system.stream_create_file( parsed.body.path, parsed.body.name, upload @@ -141,8 +141,6 @@ async def create_file_or_directory( return FileCreateResponse(success=True, info=info) except UploadTooLargeError as e: LOGGER.warning(f"Rejected oversize upload: {e}") - # 413 is the standard "payload too large" response; surface it - # explicitly rather than burying it in a 200 + success=False. raise HTTPException(status_code=413, detail=str(e)) from e except Exception as e: LOGGER.error(f"Error creating file or directory: {e}") diff --git a/marimo/_server/api/utils.py b/marimo/_server/api/utils.py index 89a74845fe6..08a10fdda9a 100644 --- a/marimo/_server/api/utils.py +++ b/marimo/_server/api/utils.py @@ -48,12 +48,8 @@ async def parse_request( @dataclass class MultipartRequest(Generic[S]): - """Result of parsing a multipart/form-data request body. - - `files` carries the raw `UploadFile` objects (not yet read) so callers - can stream chunks via `.read(size)`. For small payloads, callers can - simply `await upload.read()` to materialize the whole body. - """ + """Parsed multipart body. `files` holds un-read `UploadFile` handles + so callers can stream large parts instead of buffering.""" body: S files: dict[str, UploadFile] @@ -65,20 +61,14 @@ async def parse_multipart_request( ) -> AsyncIterator[MultipartRequest[S]]: """Parse a multipart/form-data body into a msgspec.Struct + uploads. - String form fields are validated against `cls`. File upload parts are - yielded as `UploadFile` objects (un-read) in `files`, keyed by - form-field name, so callers can stream them to disk instead of buffering. - - Used as an async context manager so the underlying form (and its - spooled temp files / fds) are closed automatically when the caller - is done — `UploadFile` parts remain readable for the entire body of - the `async with` block. + Must be used as an async context manager: `UploadFile` parts stay + readable for the body of the `async with`, and their spooled temp + files are closed on exit. Raises msgspec.ValidationError if required string fields are missing or invalid. """ - # Imported lazily so this module stays import-safe in environments - # without starlette (e.g. pyodide). + # Lazy import so this module stays import-safe under pyodide. from starlette.datastructures import UploadFile async with request.form() as form: diff --git a/marimo/_server/files/os_file_system.py b/marimo/_server/files/os_file_system.py index 37c8e1823e7..8c1c76d7341 100644 --- a/marimo/_server/files/os_file_system.py +++ b/marimo/_server/files/os_file_system.py @@ -56,11 +56,7 @@ class UploadTooLargeError(ValueError): class AsyncByteSource(Protocol): - """Anything that can be drained chunk-by-chunk into a file. - - Starlette's `UploadFile` satisfies this; so does any object exposing - an async `read(size)` returning bytes. - """ + """Structural type for things like Starlette's `UploadFile`.""" async def read(self, size: int = -1, /) -> bytes: ... @@ -237,13 +233,10 @@ async def stream_create_file( parent = Path(path) os.makedirs(parent, exist_ok=True) - # Atomically claim the destination with O_CREAT|O_EXCL so concurrent - # uploads can't both pick the same numbered suffix and clobber each - # other between `_generate_unique_path` and the final rename. + # Atomic O_CREAT|O_EXCL reservation closes the TOCTOU window between + # picking a unique name and writing to it. full_path = _claim_unique_path(parent / name) - # `NamedTemporaryFile` gives us a guaranteed-unique sibling path for - # the in-progress `.part` file (same reasoning as above). tmp = tempfile.NamedTemporaryFile( dir=full_path.parent, prefix=full_path.name + ".", @@ -252,9 +245,8 @@ async def stream_create_file( ) tmp_path = tmp.name try: - # Sync writes are bounded to ~1 MiB per chunk, with an `await` - # in between; event loop blockage is brief and an async file - # library would only add a dependency for marginal gain. + # Sync writes are bounded to one chunk between awaits, so event + # loop blockage stays small without pulling in aiofiles. written = 0 with tmp: while chunk := await source.read(_STREAM_CHUNK_SIZE): @@ -265,13 +257,10 @@ async def stream_create_file( f"{MAX_UPLOAD_BYTES} bytes" ) tmp.write(chunk) - # Replaces our empty marker at `full_path`. Atomic on POSIX - # and Windows (Python 3.3+). os.replace(tmp_path, full_path) except BaseException: - # Clean up both the `.part` and the reserved empty marker. - # If `os.replace` already succeeded, the marker is gone and - # `FileNotFoundError` is the expected outcome there. + # Both paths may or may not exist depending on where we failed; + # FileNotFoundError on either is expected. for p in (tmp_path, str(full_path)): try: os.unlink(p) @@ -279,9 +268,8 @@ async def stream_create_file( pass raise - # Use the metadata-only helper: `get_details` would re-read the - # file contents (and base64-encode binary), defeating the point of - # streaming for large uploads. + # `get_details` would re-read the file (and base64-encode binary), + # defeating the streaming win. return self._get_file_info(str(full_path)) def delete_file_or_directory(self, path: str) -> bool: @@ -578,13 +566,9 @@ def _generate_unique_path(new_path: str | Path) -> Path: def _claim_unique_path(target: Path) -> Path: - """Like `_generate_unique_path`, but atomically reserves the chosen name. - - Opens with O_CREAT|O_EXCL so two concurrent callers can never end up - with the same path — whichever loses the race sees `FileExistsError` - and tries the next numbered suffix. Returns the claimed path (an empty - file at that location); the caller is responsible for writing into it - (or replacing it). + """Race-safe variant of `_generate_unique_path`: opens with O_EXCL and + walks numbered suffixes on collision. Returns an empty file at the + claimed path; the caller writes into or replaces it. """ name_without_extension = target.stem extension = target.suffix diff --git a/marimo/_server/models/files.py b/marimo/_server/models/files.py index 80f04e0b36a..248b713d774 100644 --- a/marimo/_server/models/files.py +++ b/marimo/_server/models/files.py @@ -72,11 +72,8 @@ class FileCreateMultipartRequest(msgspec.Struct, rename="camel"): type: FileCreateType # The name of the file or directory name: str - # The raw file bytes (optional). When omitted, an empty file is created - # (or, for 'notebook' type, a default notebook template). - # NOTE: this field is OpenAPI-only — see class docstring. The - # ``format: binary`` annotation makes the generated spec emit a proper - # file-upload schema rather than a base64 string. + # OpenAPI-only (see class docstring). `format: binary` makes the + # generated spec emit a file-upload schema rather than a base64 string. file: Annotated[ str | None, msgspec.Meta(extra_json_schema={"format": "binary"}) ] = None diff --git a/tests/_server/api/test_api_utils.py b/tests/_server/api/test_api_utils.py index 1ff8a0e03ef..56e009e543d 100644 --- a/tests/_server/api/test_api_utils.py +++ b/tests/_server/api/test_api_utils.py @@ -51,7 +51,6 @@ def test_parse_multipart_request_strings_and_file_upload() -> None: files = captured["files"] assert isinstance(files, dict) assert set(files.keys()) == {"upload"} - # File handle is returned un-read; bytes loaded inside the endpoint above. assert captured["upload_bytes"] == b"\x00\x01\x02\xff" diff --git a/tests/_server/files/test_os_file_system.py b/tests/_server/files/test_os_file_system.py index 49482497139..543fe908933 100644 --- a/tests/_server/files/test_os_file_system.py +++ b/tests/_server/files/test_os_file_system.py @@ -136,11 +136,8 @@ def test_create_rejects_path_traversal( class _ChunkedSource: - """Minimal async byte source emitting predetermined chunks. - - If ``fail_after`` is set, raises ``RuntimeError`` after that many reads - have returned data — useful for simulating mid-stream upload failures. - """ + """Async byte source emitting predetermined chunks; optionally raises + after `fail_after` reads to simulate mid-stream failures.""" def __init__( self, chunks: list[bytes], *, fail_after: int | None = None @@ -154,7 +151,6 @@ async def read(self, size: int = -1, /) -> bytes: raise RuntimeError("stream failed") if not self._chunks: return b"" - # Honor size by returning the next chunk if it fits, else slicing. chunk = self._chunks[0] if size < 0 or len(chunk) <= size: out = self._chunks.pop(0) @@ -225,7 +221,6 @@ async def test_stream_create_file_rejects_traversal( await fs.stream_create_file( str(test_dir), name, _ChunkedSource([b"x"]) ) - # Nothing should have been written into the parent or test_dir. assert not (test_dir.parent / "escape.bin").exists() assert _part_files(test_dir) == [] @@ -243,8 +238,7 @@ async def test_stream_create_file_cleans_up_on_immediate_failure( async def test_stream_create_file_cleans_up_on_mid_stream_failure( test_dir: Path, fs: OSFileSystem ) -> None: - # The realistic case: a few chunks already written when the source - # raises (network drop, client disconnect, etc). + # Source fails after writes have already landed in the .part file. source = _ChunkedSource([b"first ", b"second ", b"third"], fail_after=2) with pytest.raises(RuntimeError): await fs.stream_create_file(str(test_dir), "partial.bin", source) @@ -260,7 +254,6 @@ async def test_stream_create_file_enforces_size_cap( "marimo._server.files.os_file_system.MAX_UPLOAD_BYTES", 8 ) source = _ChunkedSource([b"aaaa", b"bbbb", b"too much"]) - # Raises the specific subclass so the endpoint can map it to a 413. with pytest.raises(UploadTooLargeError, match="exceeds maximum size"): await fs.stream_create_file(str(test_dir), "big.bin", source) assert not (test_dir / "big.bin").exists() @@ -270,9 +263,8 @@ async def test_stream_create_file_enforces_size_cap( async def test_stream_create_file_claims_path_atomically( test_dir: Path, fs: OSFileSystem ) -> None: - # Simulate the TOCTOU race: a concurrent caller creates the same name - # while we're streaming. With atomic O_CREAT|O_EXCL reservation we - # must NOT overwrite the concurrent file. + # Concurrent file at the target name must not be clobbered by the + # streaming write — the atomic claim should bump to a new suffix. (test_dir / "race.bin").write_bytes(b"existing") info = await fs.stream_create_file( str(test_dir), "race.bin", _ChunkedSource([b"new"]) @@ -283,11 +275,11 @@ async def test_stream_create_file_claims_path_atomically( def test_claim_unique_path_atomic_reservation(test_dir: Path) -> None: - # First claim takes the bare name and reserves an empty file. claimed_a = _claim_unique_path(test_dir / "claim.bin") assert claimed_a == test_dir / "claim.bin" assert claimed_a.exists() - # Second claim must NOT reuse the reserved path, even though it's empty. + # An existing (even empty) reservation forces the next caller off the + # base name — this is what closes the TOCTOU window. claimed_b = _claim_unique_path(test_dir / "claim.bin") assert claimed_b == test_dir / "claim_1.bin" assert claimed_b.exists()