diff --git a/invokeai/app/api/routers/videos.py b/invokeai/app/api/routers/videos.py index cc8526e6f54..9e9a9e890d2 100644 --- a/invokeai/app/api/routers/videos.py +++ b/invokeai/app/api/routers/videos.py @@ -6,13 +6,16 @@ from pathlib import Path from typing import Annotated, BinaryIO, Optional -from fastapi import Body, HTTPException, Query, Request, Response, UploadFile +from fastapi import Body, HTTPException, Query, Request, Response from fastapi import Path as PathParam from fastapi.responses import StreamingResponse from fastapi.routing import APIRouter from PIL import Image as PILImage from pydantic import BaseModel, Field, StringConstraints, ValidationError +from python_multipart.exceptions import MultipartParseError +from python_multipart.multipart import MultipartParser, parse_options_header from starlette.concurrency import run_in_threadpool +from starlette.requests import ClientDisconnect from invokeai.app.api.auth_dependencies import CurrentMediaUserOrDefault, CurrentUserOrDefault from invokeai.app.api.dependencies import ApiDependencies @@ -52,17 +55,19 @@ # Per-chunk size for HTTP Range responses (1 MB) RANGE_CHUNK_SIZE = 1024 * 1024 -# Upload streaming chunk size (1 MB) and a coarse per-upload size cap. The cap is generous +# Coarse per-upload size cap, enforced against the file part as it streams in. Generous # because Wan-generated MP4s for long sequences can run into the hundreds of megabytes; -# the goal is to prevent a single client from exhausting RAM, not to be a content policy. -UPLOAD_CHUNK_SIZE = 1024 * 1024 +# the goal is to prevent a single client from exhausting RAM/disk, not to be a content policy. MAX_UPLOAD_SIZE = 1024 * 1024 * 1024 # 1 GB # Pre-parse ingress cap enforced by VideoUploadLimitASGIMiddleware, applied to the whole -# request body *before* the multipart parser spools it to temp storage. Slightly larger -# than MAX_UPLOAD_SIZE to allow for multipart framing and the metadata form field. +# request body *before* the upload route parses it. Slightly larger than MAX_UPLOAD_SIZE +# to allow for multipart framing and the metadata form field. MAX_UPLOAD_REQUEST_SIZE = MAX_UPLOAD_SIZE + 10 * 1024 * 1024 -# Global bound on concurrent video uploads — each in-flight upload can hold up to two -# copies of the file in temp storage (the multipart spool + the route's own tmp file). +# The `metadata` form field is a stringified JSON dict; it is buffered in memory while the +# body streams, so it gets its own (generous) cap. +MAX_UPLOAD_METADATA_SIZE = 1024 * 1024 +# Global bound on concurrent video uploads — each in-flight upload holds one full-size copy +# of the file in temp storage until probe/thumbnail/create finish with it. MAX_CONCURRENT_VIDEO_UPLOADS = 2 # Per-user bound (multiuser mode only): keeps one tenant's slow uploads from holding # every global slot and starving the other users into 429s. @@ -178,14 +183,174 @@ def _assert_video_read_access(video_name: str, current_user: CurrentUserOrDefaul raise HTTPException(status_code=403, detail="Not authorized to access this video") -def _is_accepted_video_upload(file: UploadFile) -> bool: - if file.content_type and file.content_type.startswith(ACCEPTED_VIDEO_MIME_PREFIXES): +def _is_accepted_video_upload(filename: Optional[str], content_type: Optional[str]) -> bool: + if content_type and content_type.startswith(ACCEPTED_VIDEO_MIME_PREFIXES): return True - if file.filename: - return file.filename.lower().endswith(ACCEPTED_VIDEO_EXTENSIONS) + if filename: + return filename.lower().endswith(ACCEPTED_VIDEO_EXTENSIONS) return False +class _VideoUploadStreamParser: + """Parses the multipart upload body, writing the file part straight to `destination`. + + Declaring `file: UploadFile` on the route makes Starlette parse the body into its own + spooled temp file first, so the route's copy to a named temp file was a SECOND + full-size copy: every in-flight upload occupied up to 2 x MAX_UPLOAD_SIZE of temp + storage (x MAX_CONCURRENT_VIDEO_UPLOADS) for the whole probe/thumbnail/create phase. + The spool's path cannot be reused instead — once rolled over it is an unlinked + anonymous file, and both ffmpeg and videos.create need a real path. + + Parsing the stream ourselves keeps exactly one copy on disk. It also lets the + file-type and size checks fire while the body is still arriving, rather than after the + whole thing has been written somewhere. + + Callbacks run inside `MultipartParser.write`, which the route calls in a worker thread + — the disk writes must not happen on the event loop. + """ + + def __init__(self, destination: BinaryIO) -> None: + self._destination = destination + self._header_field = bytearray() + self._header_value = bytearray() + self._headers: dict[bytes, bytes] = {} + self._part_name: Optional[bytes] = None + self._metadata_chunks: list[bytes] = [] + self._metadata_size = 0 + self.filename: Optional[str] = None + self.content_type: Optional[str] = None + self.metadata: Optional[str] = None + self.file_size = 0 + self.saw_file_part = False + # `MultipartParser.finalize()` is a no-op that does NOT check the parser reached its + # end state, so a body that stops after the file bytes — truncated upload, aborted + # client, or the parser's own silent `max_size` truncation — would otherwise look + # exactly like a complete one and get probed and persisted. `on_end` fires only when + # the closing `--boundary--` is parsed, so it is the proof of a complete body. + self.saw_end = False + + @property + def callbacks(self) -> dict[str, object]: + return { + "on_part_begin": self._on_part_begin, + "on_part_data": self._on_part_data, + "on_part_end": self._on_part_end, + "on_header_field": self._on_header_field, + "on_header_value": self._on_header_value, + "on_header_end": self._on_header_end, + "on_headers_finished": self._on_headers_finished, + "on_end": self._on_end, + } + + def _on_part_begin(self) -> None: + self._headers = {} + self._header_field = bytearray() + self._header_value = bytearray() + self._part_name = None + self._metadata_chunks = [] + self._metadata_size = 0 + + def _on_header_field(self, data: bytes, start: int, end: int) -> None: + self._header_field.extend(data[start:end]) + + def _on_header_value(self, data: bytes, start: int, end: int) -> None: + self._header_value.extend(data[start:end]) + + def _on_header_end(self) -> None: + self._headers[bytes(self._header_field).lower()] = bytes(self._header_value) + self._header_field = bytearray() + self._header_value = bytearray() + + def _on_headers_finished(self) -> None: + _, options = parse_options_header(self._headers.get(b"content-disposition", b"")) + self._part_name = options.get(b"name") + if self._part_name != b"file": + return + if self.saw_file_part: + raise HTTPException(status_code=422, detail="Expected exactly one video file") + self.saw_file_part = True + filename = options.get(b"filename") + self.filename = filename.decode("utf-8", errors="replace") if filename is not None else None + content_type = self._headers.get(b"content-type") + self.content_type = content_type.decode("latin-1") if content_type is not None else None + # Reject the wrong kind of file before any of its bytes reach the disk. + if not _is_accepted_video_upload(self.filename, self.content_type): + raise HTTPException(status_code=415, detail="Not a supported video file") + + def _on_part_data(self, data: bytes, start: int, end: int) -> None: + chunk = data[start:end] + if self._part_name == b"file": + self.file_size += len(chunk) + if self.file_size > MAX_UPLOAD_SIZE: + raise HTTPException( + status_code=413, + detail=f"Video upload exceeds maximum size ({MAX_UPLOAD_SIZE} bytes)", + ) + self._destination.write(chunk) + elif self._part_name == b"metadata": + self._metadata_size += len(chunk) + if self._metadata_size > MAX_UPLOAD_METADATA_SIZE: + raise HTTPException( + status_code=413, + detail=f"Video metadata exceeds maximum size ({MAX_UPLOAD_METADATA_SIZE} bytes)", + ) + self._metadata_chunks.append(chunk) + # Any other field is dropped rather than buffered: an unknown part must not be a + # way to make the server hold arbitrary bytes in memory. + + def _on_part_end(self) -> None: + if self._part_name == b"metadata": + try: + self.metadata = b"".join(self._metadata_chunks).decode("utf-8") + except UnicodeDecodeError as error: + raise HTTPException(status_code=422, detail="Metadata must be UTF-8 encoded") from error + self._part_name = None + self._metadata_chunks = [] + self._metadata_size = 0 + + def _on_end(self) -> None: + self.saw_end = True + + +async def _stream_video_upload(request: Request, destination: BinaryIO) -> _VideoUploadStreamParser: + """Streams the request body through the multipart parser into `destination`.""" + media_type, options = parse_options_header(request.headers.get("content-type", "")) + boundary = options.get(b"boundary") + # Content-Type is case-insensitive (RFC 7231) and parse_options_header preserves case. + if media_type.lower() != b"multipart/form-data" or boundary is None: + raise HTTPException(status_code=422, detail="Expected a multipart/form-data video upload") + + parser_state = _VideoUploadStreamParser(destination) + # max_size is the pre-parse ingress cap the middleware already enforces; repeating it + # here bounds the parser itself for any path that reaches it directly. The parser + # silently *truncates* past it rather than erroring, so on such a path the body would + # look short rather than rejected — the saw_end check below is what catches that. + parser = MultipartParser(boundary, parser_state.callbacks, max_size=MAX_UPLOAD_REQUEST_SIZE) + try: + async for chunk in request.stream(): + # Parsing writes to disk, so it belongs in the thread pool alongside the rest of + # the blocking upload work. + await run_in_threadpool(parser.write, chunk) + await run_in_threadpool(parser.finalize) + except MultipartParseError as error: + # A malformed body is the client's fault, not a server error. + raise HTTPException(status_code=422, detail="Malformed multipart body") from error + except ClientDisconnect as error: + # The client went away, or VideoUploadLimitASGIMiddleware cut it off for exceeding the + # ingress cap, the idle timeout or the duration cap. Left to propagate this surfaces as + # a 500 raised above the middleware, which both misreports a client-side abort as a + # server fault and bypasses the middleware's connection-close handling. + raise HTTPException(status_code=400, detail="Upload ended before the body was complete") from error + + if not parser_state.saw_end: + # No closing boundary: the body was truncated. Proceeding would probe and persist a + # partial file whenever the truncated bytes happen to survive the MP4 checks. + raise HTTPException(status_code=422, detail="Incomplete multipart body") + if not parser_state.saw_file_part: + raise HTTPException(status_code=422, detail="Expected a video file in the upload") + return parser_state + + def _is_mp4_file(path: Path) -> bool: try: with open(path, "rb") as video_file: @@ -245,33 +410,57 @@ def _probe_decodable_video(path: Path) -> tuple[tuple[int, int, float, Optional[ }, status_code=201, response_model=VideoDTO, + # The body is parsed by hand (see _stream_video_upload) so the file lands in exactly + # one temp file, which means FastAPI cannot infer the request schema from the + # signature. This spells out the same multipart body the `file` + `metadata` + # parameters used to generate, so the documented contract is unchanged. + openapi_extra={ + "requestBody": { + "required": True, + "content": { + "multipart/form-data": { + "schema": { + "properties": { + # Key order and shape mirror what FastAPI generates for the sibling + # upload routes (see Body_upload_image), so the documented contract + # stays byte-identical to what `file` + `metadata` produced. + "file": { + "type": "string", + "contentMediaType": "application/octet-stream", + "title": "File", + }, + "metadata": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "title": "Metadata", + "description": "The metadata to associate with the video, must be a stringified JSON dict", + }, + }, + "type": "object", + "required": ["file"], + "title": "Body_upload_video", + } + } + }, + } + }, ) async def upload_video( current_user: CurrentUserOrDefault, - file: UploadFile, request: Request, response: Response, video_category: ImageCategory = Query(description="The category of the video"), is_intermediate: bool = Query(description="Whether this is an intermediate video"), board_id: Optional[str] = Query(default=None, description="The board to add this video to, if any"), session_id: Optional[str] = Query(default=None, description="The session ID associated with this upload, if any"), - metadata: Optional[str] = Body( - default=None, - description="The metadata to associate with the video, must be a stringified JSON dict", - embed=True, - ), ) -> VideoDTO: """Uploads a video for the current user.""" - if metadata is not None: - try: - MetadataFieldValidator.validate_json(metadata) - except ValidationError as e: - raise HTTPException(status_code=422, detail="Metadata must be a JSON object") from e - # Check board access for uploads to a specific board. if board_id is not None: from invokeai.app.services.board_records.board_records_common import BoardVisibility + # This rejects before any of the body is read. VideoUploadLimitASGIMiddleware is what + # keeps that from becoming a quota hole: it closes the connection when the app + # answers early, so the client's in-flight upload dies with the response. try: board = await run_in_threadpool(ApiDependencies.invoker.services.boards.get_dto, board_id=board_id) except Exception: @@ -283,35 +472,24 @@ async def upload_video( ): raise HTTPException(status_code=403, detail="Not authorized to upload to this board") - if not _is_accepted_video_upload(file): - raise HTTPException(status_code=415, detail="Not a supported video file") - - # Stream the upload to a tmp file so we can probe and then hand its path to the service. - # Reading the full body into memory first risked exhausting RAM on multi-GB uploads; - # chunk-stream instead and enforce a hard size cap. Filesystem writes, container - # validation, ffmpeg probing, and thumbnail extraction are all blocking — run them in - # the thread pool so a slow (or hostile) file can't stall the event loop and every - # other API request with it. + # Stream the upload straight into a tmp file so we can probe it and then hand its path + # to the service. Reading the full body into memory first risked exhausting RAM on + # multi-GB uploads; the parser streams it instead and enforces a hard size cap as the + # bytes arrive. Filesystem writes, container validation, ffmpeg probing, and thumbnail + # extraction are all blocking — run them in the thread pool so a slow (or hostile) file + # can't stall the event loop and every other API request with it. tmp = tempfile.NamedTemporaryFile(prefix="invokeai_upload_", suffix=".mp4", delete=False) tmp_path = Path(tmp.name) try: - total = 0 - while chunk := await file.read(UPLOAD_CHUNK_SIZE): - total += len(chunk) - if total > MAX_UPLOAD_SIZE: - tmp.close() - raise HTTPException( - status_code=413, - detail=f"Video upload exceeds maximum size ({MAX_UPLOAD_SIZE} bytes)", - ) - await run_in_threadpool(tmp.write, chunk) + upload = await _stream_video_upload(request, tmp) tmp.close() - # Release the multipart spool now that the body is copied: each in-flight - # upload otherwise holds TWO on-disk copies (Starlette's spool + our tmp file) - # through the probe/thumbnail/create phase — up to 2 x MAX_UPLOAD_SIZE x - # MAX_CONCURRENT_VIDEO_UPLOADS of temp disk. Closing shrinks the double-copy - # window to the copy loop itself. - await file.close() + + metadata = upload.metadata + if metadata is not None: + try: + MetadataFieldValidator.validate_json(metadata) + except ValidationError as e: + raise HTTPException(status_code=422, detail="Metadata must be a JSON object") from e if not await run_in_threadpool(_is_mp4_file, tmp_path): raise HTTPException(status_code=415, detail="Not an MP4 video file") diff --git a/invokeai/app/api_app.py b/invokeai/app/api_app.py index f35a590e2d0..dee6988f7a8 100644 --- a/invokeai/app/api_app.py +++ b/invokeai/app/api_app.py @@ -349,6 +349,28 @@ class VideoUploadLimitASGIMiddleware: chunked bodies that exceed the cap mid-stream, and bounds concurrent uploads both globally and per user (so one tenant's slow uploads cannot starve the others into 429s). + + It also asks the server to close the connection on any response sent before the request + body has been read to completion. The leases below are released as soon as the app + returns, and the route answers plenty of requests without reading the body (a forbidden + board, a filename that isn't .mp4) — FastAPI's own query-param validation answers 422 + before the route body runs at all. A client that kept streaming after such a response + would hold ingress with no slot charged against it, outside the 429 bound, the idle + timeout and the duration cap; closing ends that upload along with the response. + + Whether the body was read is the only thing that can be known here, so any early answer + closes — including when the client had in fact already finished sending. That costs a + fresh connection per rejected upload, which is the conservative side to err on and is + what servers generally do when a response is sent without consuming the body. + + Two limits worth knowing. Draining the body instead would also close the hole, but it + would pin one of the very few upload slots for the whole duration cap per rejection, + which is a cheaper denial of service than the hole it closes. And behind a reverse proxy + that buffers request bodies (nginx's default, per the multi-user admin guide) `Connection` + is hop-by-hop, so this only closes the proxy-to-app hop — there the proxy has already + absorbed the whole upload before the app is invoked, so the hole does not arise. Responses + generated above this middleware (Starlette's ServerErrorMiddleware 500) do not pass + through it; uvicorn closes the transport on those itself. """ def __init__( @@ -379,6 +401,11 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if not (scope.get("method") == "POST" and route_path == "/api/v1/videos/upload"): return await self.app(scope, receive, send) + # `connection` is a hop-by-hop header and illegal in HTTP/2+, so every use below is + # gated on HTTP/1. + is_http1 = str(scope.get("http_version", "1.1")).startswith("1.") + close_header = {"connection": "close"} if is_http1 else {} + per_user_key: str | None = None if self.identify_user is not None: identity = self.identify_user(scope) @@ -389,7 +416,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: response = JSONResponse( {"detail": "Authentication required"}, status_code=401, - headers={"WWW-Authenticate": "Bearer"}, + headers={"WWW-Authenticate": "Bearer", **close_header}, ) return await response(scope, receive, send) @@ -398,6 +425,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: response = JSONResponse( {"detail": f"Video upload exceeds maximum request size ({self.max_body_bytes} bytes)"}, status_code=413, + headers=close_header, ) return await response(scope, receive, send) @@ -405,7 +433,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: response = JSONResponse( {"detail": "Too many concurrent video uploads; try again shortly"}, status_code=429, - headers={"Retry-After": "5"}, + headers={"Retry-After": "5", **close_header}, ) return await response(scope, receive, send) @@ -417,7 +445,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: response = JSONResponse( {"detail": "Too many concurrent video uploads for this user; try again shortly"}, status_code=429, - headers={"Retry-After": "5"}, + headers={"Retry-After": "5", **close_header}, ) return await response(scope, receive, send) @@ -425,6 +453,13 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if per_user_key is not None: self._active_by_user[per_user_key] = self._active_by_user.get(per_user_key, 0) + 1 received = 0 + # Only reading the body to its end proves the client has finished sending. This must + # NOT be seeded from Content-Length: h11 accepts `Content-Length: 0` alongside + # `Transfer-Encoding: chunked` (the chunked framing wins), so trusting the header let + # a client suppress the close and then stream indefinitely with no lease held and + # none of the caps below applying — they all live in limited_receive, which an app + # that answers early never calls again. + body_finished = False upload_started_at = asyncio.get_running_loop().time() async def limited_receive() -> Message: @@ -432,10 +467,13 @@ async def limited_receive() -> Message: # streamed body and abort the request once it exceeds the cap, so the multipart # parser stops spooling. A clean 413 isn't possible mid-parse; the aborted # request surfaces to the client as a dropped connection. - nonlocal received + nonlocal received, body_finished remaining_duration = self.max_upload_duration_seconds - ( asyncio.get_running_loop().time() - upload_started_at ) + # The three aborts below synthesize a disconnect precisely because the client is + # still uploading, so they deliberately leave body_finished alone: whatever the + # app answers must still close the connection. if remaining_duration <= 0: return {"type": "http.disconnect"} try: @@ -446,10 +484,24 @@ async def limited_receive() -> Message: received += len(message.get("body", b"")) if received > self.max_body_bytes: return {"type": "http.disconnect"} + if not message.get("more_body", False): + body_finished = True + elif message["type"] == "http.disconnect": + # The client is already gone; there is nothing left to close. + body_finished = True return message + async def close_if_answered_early(message: Message) -> None: + # See the class docstring: answering before the body has been read to its end must + # not leave the client uploading into an already-sent response. + if message["type"] == "http.response.start" and not body_finished and is_http1: + headers = [(name, value) for name, value in message.get("headers", []) if name.lower() != b"connection"] + headers.append((b"connection", b"close")) + message = {**message, "headers": headers} + await send(message) + try: - await self.app(scope, limited_receive, send) + await self.app(scope, limited_receive, close_if_answered_early) finally: self._active -= 1 if per_user_key is not None: diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index ccd0f318694..cc470ea5c3d 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -6087,16 +6087,6 @@ "description": "The session ID associated with this upload, if any" } ], - "requestBody": { - "required": true, - "content": { - "multipart/form-data": { - "schema": { - "$ref": "#/components/schemas/Body_upload_video" - } - } - } - }, "responses": { "201": { "description": "The video was uploaded successfully", @@ -6121,6 +6111,37 @@ } } } + }, + "requestBody": { + "required": true, + "content": { + "multipart/form-data": { + "schema": { + "properties": { + "file": { + "type": "string", + "contentMediaType": "application/octet-stream", + "title": "File" + }, + "metadata": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Metadata", + "description": "The metadata to associate with the video, must be a stringified JSON dict" + } + }, + "type": "object", + "required": ["file"], + "title": "Body_upload_video" + } + } + } } } }, @@ -16619,30 +16640,6 @@ "required": ["file"], "title": "Body_upload_image" }, - "Body_upload_video": { - "properties": { - "file": { - "type": "string", - "contentMediaType": "application/octet-stream", - "title": "File" - }, - "metadata": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Metadata", - "description": "The metadata to associate with the video, must be a stringified JSON dict" - } - }, - "type": "object", - "required": ["file"], - "title": "Body_upload_video" - }, "BooleanCollectionInvocation": { "category": "primitives", "class": "invocation", diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index c19832d09a1..2365dd00afa 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -5147,16 +5147,6 @@ export type components = { */ metadata?: string | null; }; - /** Body_upload_video */ - Body_upload_video: { - /** File */ - file: Blob; - /** - * Metadata - * @description The metadata to associate with the video, must be a stringified JSON dict - */ - metadata?: string | null; - }; /** * Boolean Collection Primitive * @description A collection of boolean primitive values @@ -45273,7 +45263,15 @@ export interface operations { }; requestBody: { content: { - "multipart/form-data": components["schemas"]["Body_upload_video"]; + "multipart/form-data": { + /** File */ + file: Blob; + /** + * Metadata + * @description The metadata to associate with the video, must be a stringified JSON dict + */ + metadata?: string | null; + }; }; }; responses: { diff --git a/tests/app/api/test_video_upload_limits.py b/tests/app/api/test_video_upload_limits.py index 1d1df14fd3f..726ecca835a 100644 --- a/tests/app/api/test_video_upload_limits.py +++ b/tests/app/api/test_video_upload_limits.py @@ -1,8 +1,10 @@ -"""Tests for VideoUploadLimitASGIMiddleware (PR #9163 review fix). +"""Tests for VideoUploadLimitASGIMiddleware and the upload route's body handling. -The upload route's MAX_UPLOAD_SIZE check runs only after FastAPI has parsed (and spooled) -the entire multipart body, so oversized/chunked/concurrent requests could exhaust temp -storage before rejection. The middleware bounds ingress before the parser runs. +The middleware (PR #9163 review fix) bounds ingress before any parsing happens, so +oversized, chunked or too-many-concurrent requests are rejected without touching temp +storage at all. The route itself then parses the multipart body and streams the file part +straight into one temp file, enforcing MAX_UPLOAD_SIZE as the bytes arrive — declaring +`file: UploadFile` used to make Starlette spool a second full-size copy first. """ import asyncio @@ -11,11 +13,12 @@ from pathlib import Path from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest -from fastapi import FastAPI, Response, UploadFile +from fastapi import FastAPI, HTTPException, Response from fastapi.testclient import TestClient +from python_multipart.multipart import MultipartParser, MultipartState from starlette.datastructures import Headers from invokeai.app import api_app @@ -28,6 +31,7 @@ _identify_video_upload_user_async, ) from invokeai.app.services.auth.token_service import TokenData, create_access_token, set_jwt_secret +from invokeai.app.services.board_records.board_records_common import BoardVisibility from invokeai.app.services.image_records.image_records_common import ImageCategory from invokeai.app.util.video_thumbnails import VideoDecodeTimeoutError @@ -35,10 +39,15 @@ MAX_CONCURRENT = 2 -def test_configured_upload_slots_bound_peak_double_spool_usage() -> None: +def test_configured_upload_slots_bound_peak_temp_storage_usage() -> None: + """Peak temp storage is now one copy per in-flight upload, not two. + + The route parses the multipart body itself and writes the file part straight into its + own temp file, so Starlette's spool no longer holds a second full-size copy. + """ assert videos.MAX_CONCURRENT_VIDEO_UPLOADS <= 2 assert videos.MAX_CONCURRENT_VIDEO_UPLOADS_PER_USER <= 1 - assert 2 * videos.MAX_UPLOAD_SIZE * videos.MAX_CONCURRENT_VIDEO_UPLOADS <= 4 * 1024 * 1024 * 1024 + assert videos.MAX_UPLOAD_SIZE * videos.MAX_CONCURRENT_VIDEO_UPLOADS <= 2 * 1024 * 1024 * 1024 def test_upload_probe_requires_a_decodable_frame(monkeypatch: pytest.MonkeyPatch): @@ -443,6 +452,98 @@ async def send(_message: dict[str, Any]) -> None: assert offloaded == [_identify_video_upload_user] +BOUNDARY = "testboundary" + + +def _multipart_body( + file_bytes: bytes, + filename: str = "video.mp4", + metadata: str | None = None, + content_type: str = "video/mp4", +) -> bytes: + parts = [] + if metadata is not None: + parts.append(f'--{BOUNDARY}\r\nContent-Disposition: form-data; name="metadata"\r\n\r\n{metadata}\r\n'.encode()) + parts.append( + f'--{BOUNDARY}\r\nContent-Disposition: form-data; name="file"; filename="{filename}"\r\n' + f"Content-Type: {content_type}\r\n\r\n".encode() + + file_bytes + + b"\r\n" + ) + parts.append(f"--{BOUNDARY}--\r\n".encode()) + return b"".join(parts) + + +def test_python_multipart_contract_the_upload_route_depends_on(): + """Pins the `python_multipart` behaviours the hand-rolled parsing relies on. + + The route drives `MultipartParser` directly instead of going through Starlette, so a + version bump that changes any of this silently changes upload behaviour. If one of + these ever fails, revisit `_stream_video_upload` rather than just updating the numbers. + """ + body = _multipart_body(b"abc") + closing = f"--{BOUNDARY}--\r\n".encode() + + # 1. `on_end` fires exactly once, at the closing boundary. `saw_end` is built on it. + seen: list[str] = [] + parser = MultipartParser( + BOUNDARY.encode(), + {"on_part_begin": lambda: seen.append("part_begin"), "on_end": lambda: seen.append("end")}, + ) + parser.write(body) + parser.finalize() + assert seen.count("end") == 1 + assert seen[-1] == "end" + assert parser.state == MultipartState.END + + # 2. `finalize()` does NOT verify the parser reached its end state — it is documented as + # a no-op with a TODO. A body cut off before the closing boundary finalizes cleanly, + # which is exactly why the route needs its own end-of-body check. + truncated_end: list[str] = [] + truncated = MultipartParser(BOUNDARY.encode(), {"on_end": lambda: truncated_end.append("end")}) + truncated.write(body[: -len(closing)]) + truncated.finalize() # does not raise + assert truncated_end == [] + assert truncated.state != MultipartState.END + + # 3. `max_size` silently *truncates* rather than raising, so an over-cap body arrives at + # the same end-of-body check instead of erroring out of `write()`. + capped_end: list[str] = [] + capped = MultipartParser( + BOUNDARY.encode(), {"on_end": lambda: capped_end.append("end")}, max_size=len(body) - len(closing) + ) + capped.write(body) + capped.finalize() # does not raise + assert capped_end == [] + + +def _fake_upload_request(body: bytes, chunk_size: int = 8) -> MagicMock: + """A Request stand-in that dribbles the body out in small chunks.""" + request = MagicMock() + request.headers = {"content-type": f"multipart/form-data; boundary={BOUNDARY}"} + + async def stream(): + for start in range(0, len(body), chunk_size): + yield body[start : start + chunk_size] + + request.stream = stream + return request + + +def _run_upload(request: MagicMock) -> Any: + return asyncio.run( + upload_video( + current_user=TokenData(user_id="user", email="user@example.com", is_admin=False), + request=request, + response=Response(), + video_category=ImageCategory.GENERAL, + is_intermediate=False, + board_id=None, + session_id=None, + ) + ) + + def test_upload_video_closes_tmp_handle_when_stream_copy_fails(): captured_handles: list[Any] = [] real_named_tmp = tempfile.NamedTemporaryFile @@ -456,30 +557,13 @@ def failing_named_tmp(*args: Any, **kwargs: Any): captured_handles.append(handle) return handle - upload = MagicMock(spec=UploadFile) - upload.filename = "video.mp4" - upload.content_type = "video/mp4" - upload.read = AsyncMock(side_effect=[b"not-empty", b""]) - try: with ( patch("invokeai.app.api.routers.videos.tempfile.NamedTemporaryFile", side_effect=failing_named_tmp), patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), pytest.raises(OSError, match="disk full"), ): - asyncio.run( - upload_video( - current_user=TokenData(user_id="user", email="user@example.com", is_admin=False), - file=upload, - request=MagicMock(), - response=Response(), - video_category=ImageCategory.GENERAL, - is_intermediate=False, - board_id=None, - session_id=None, - metadata=None, - ) - ) + _run_upload(_fake_upload_request(_multipart_body(b"not-empty"))) assert len(captured_handles) == 1 assert captured_handles[0].closed @@ -490,6 +574,482 @@ def failing_named_tmp(*args: Any, **kwargs: Any): Path(handle.name).unlink(missing_ok=True) +def test_upload_video_writes_exactly_one_copy_of_the_body(): + """The file part is written straight to the route's temp file, in order, once.""" + payload = bytes(range(256)) * 8 + written: list[bytes] = [] + captured_handles: list[Any] = [] + real_named_tmp = tempfile.NamedTemporaryFile + + async def run_immediately(func: Any, *args: Any): + return func(*args) + + def recording_named_tmp(*args: Any, **kwargs: Any): + handle = real_named_tmp(*args, **kwargs) + real_write = handle.write + handle.write = lambda chunk: (written.append(bytes(chunk)), real_write(chunk))[1] + captured_handles.append(handle) + return handle + + try: + with ( + patch("invokeai.app.api.routers.videos.tempfile.NamedTemporaryFile", side_effect=recording_named_tmp), + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + patch("invokeai.app.api.routers.videos._is_mp4_file", return_value=False), + pytest.raises(HTTPException) as error, + ): + _run_upload(_fake_upload_request(_multipart_body(payload), chunk_size=13)) + + # Stops at the container check — the point is what reached the disk before that. + assert error.value.status_code == 415 + assert b"".join(written) == payload + finally: + for handle in captured_handles: + handle.close() + Path(handle.name).unlink(missing_ok=True) + + +def test_upload_video_rejects_bad_file_part_before_any_of_it_reaches_disk(): + """A rejected upload must not put any of its payload on disk. + + The old route could only reject after Starlette had spooled the whole thing; the + filename/MIME gate now fires from the part headers, before the payload streams in. + """ + payload = b"x" * 4096 + written: list[bytes] = [] + captured_handles: list[Any] = [] + real_named_tmp = tempfile.NamedTemporaryFile + body = _multipart_body(payload, filename="notes.txt", content_type="text/plain") + + async def run_immediately(func: Any, *args: Any): + return func(*args) + + def recording_named_tmp(*args: Any, **kwargs: Any): + handle = real_named_tmp(*args, **kwargs) + real_write = handle.write + handle.write = lambda chunk: (written.append(bytes(chunk)), real_write(chunk))[1] + captured_handles.append(handle) + return handle + + try: + with ( + patch("invokeai.app.api.routers.videos.tempfile.NamedTemporaryFile", side_effect=recording_named_tmp), + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + pytest.raises(HTTPException) as error, + ): + _run_upload(_fake_upload_request(body, chunk_size=64)) + + assert error.value.status_code == 415 + assert written == [] + finally: + for handle in captured_handles: + handle.close() + Path(handle.name).unlink(missing_ok=True) + + +def test_upload_video_rejects_bad_file_part_without_finishing_the_body(): + """A rejected upload must not require reading the rest of the body first. + + The old route could only reject after Starlette had spooled the whole thing; the + filename/MIME gate now fires from the part headers, before the payload streams in. + """ + consumed: list[int] = [] + body = _multipart_body(b"x" * 4096, filename="notes.txt", content_type="text/plain") + + async def run_immediately(func: Any, *args: Any): + return func(*args) + + request = MagicMock() + request.headers = {"content-type": f"multipart/form-data; boundary={BOUNDARY}"} + + async def stream(): + for start in range(0, len(body), 64): + consumed.append(start) + yield body[start : start + 64] + + request.stream = stream + + with ( + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + pytest.raises(HTTPException) as error, + ): + _run_upload(request) + + assert error.value.status_code == 415 + # Rejected from the part headers: only the first chunks were ever read. + assert len(consumed) * 64 < len(body) + + +@pytest.mark.parametrize("header", ["multipart/form-data", "MULTIPART/Form-Data"]) +def test_upload_video_accepts_the_content_type_in_any_case(header: str): + """Content-Type is case-insensitive (RFC 7231) and parse_options_header preserves case, + so an upper-case media type used to be a 422 instead of a normal upload.""" + + async def run_immediately(func: Any, *args: Any): + return func(*args) + + request = MagicMock() + request.headers = {"content-type": f"{header}; boundary={BOUNDARY}"} + body = _multipart_body(b"abc") + + async def stream(): + yield body + + request.stream = stream + + with ( + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + patch("invokeai.app.api.routers.videos._is_mp4_file", return_value=False), + pytest.raises(HTTPException) as error, + ): + _run_upload(request) + + # Got as far as the container check either way — i.e. the body was parsed, not refused. + assert error.value.status_code == 415 + + +def test_upload_video_rejects_a_body_with_no_closing_boundary(): + """`MultipartParser.finalize()` does not check the parser reached its end state. + + Without the explicit end-of-body check a body that stops right after the file bytes + looks complete, and a truncated file gets probed and persisted whenever the partial + bytes happen to survive the MP4 checks. + """ + complete = _multipart_body(b"y" * 512) + truncated = complete[: -len(f"--{BOUNDARY}--\r\n".encode())] + is_mp4 = MagicMock(return_value=True) + + async def run_immediately(func: Any, *args: Any): + return func(*args) + + with ( + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + patch("invokeai.app.api.routers.videos._is_mp4_file", is_mp4), + pytest.raises(HTTPException) as error, + ): + _run_upload(_fake_upload_request(truncated, chunk_size=64)) + + assert error.value.status_code == 422 + # Stopped before the container check, so nothing downstream (probe, create) ran either. + is_mp4.assert_not_called() + + # Control: the same body with its closing boundary gets past the completeness check. + with ( + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + patch("invokeai.app.api.routers.videos._is_mp4_file", return_value=False), + pytest.raises(HTTPException) as error, + ): + _run_upload(_fake_upload_request(complete, chunk_size=64)) + + assert error.value.status_code == 415 + + +def _forbidden_board_deps() -> MagicMock: + deps = MagicMock() + deps.invoker.services.boards.get_dto.return_value = SimpleNamespace( + user_id="someone-else", board_visibility=BoardVisibility.Private + ) + return deps + + +def _upload_scope(query_string: bytes = b"video_category=general&is_intermediate=false") -> dict[str, Any]: + return { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/api/v1/videos/upload", + "raw_path": b"/api/v1/videos/upload", + "query_string": query_string, + "headers": [(b"content-type", f"multipart/form-data; boundary={BOUNDARY}".encode())], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "root_path": "", + } + + +def _drive_upload( + middleware: VideoUploadLimitASGIMiddleware, + body: bytes, + scope: dict[str, Any] | None = None, + chunk_delay: float = 0.0, +) -> tuple[int, list[tuple[bytes, bytes]], int, int]: + """POST `body` in small chunks; return (status, response headers, chunks sent, chunks total).""" + chunks = [body[start : start + 64] for start in range(0, len(body), 64)] + sent = 0 + result: dict[str, Any] = {} + + async def receive() -> dict[str, Any]: + nonlocal sent + if chunk_delay: + await asyncio.sleep(chunk_delay) + if sent >= len(chunks): + return {"type": "http.request", "body": b"", "more_body": False} + chunk = chunks[sent] + sent += 1 + return {"type": "http.request", "body": chunk, "more_body": True} + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.start": + result["status"] = message["status"] + result["headers"] = [(name.lower(), value.lower()) for name, value in message["headers"]] + result["sent_at_response"] = sent + + asyncio.run(asyncio.wait_for(middleware(_upload_scope() if scope is None else scope, receive, send), timeout=5)) # type: ignore[arg-type] + return result["status"], result["headers"], result["sent_at_response"], len(chunks) + + +def _real_upload_app() -> FastAPI: + """A FastAPI app carrying the real videos router, with only authentication stubbed.""" + from invokeai.app.api.auth_dependencies import get_current_user_or_default + + app = FastAPI() + app.include_router(videos.videos_router, prefix="/api") + app.dependency_overrides[get_current_user_or_default] = lambda: TokenData( + user_id="user", email="user@example.com", is_admin=False + ) + return app + + +@pytest.mark.parametrize( + "query_string,expected_status", + [ + (b"video_category=general&is_intermediate=false&board_id=not-mine", 403), # rejected inside the route + (b"video_category=bogus&is_intermediate=false", 422), # rejected by FastAPI, before it + (b"", 422), # missing required query params — also before the route body + ], + ids=["forbidden-board", "invalid-query", "missing-query"], +) +def test_upload_answered_before_the_body_ends_closes_the_connection(query_string: bytes, expected_status: int): + """An upload answered while the client is still sending must not leave it streaming. + + VideoUploadLimitASGIMiddleware releases its global and per-user leases as soon as the app + returns, so a client that keeps sending afterwards would hold ingress with no slot + charged against it — outside the 429 bound, the idle timeout and the duration cap. + `Connection: close` ends that upload with the response instead. + + Draining the body would also close the hole, but it cannot: FastAPI answers an invalid + query string before the route body runs at all (the `invalid-query`/`missing-query` + cases), so no amount of route-level draining reaches those paths. It would also pin one + of only MAX_CONCURRENT_VIDEO_UPLOADS slots for the full duration cap per rejection, + making each early rejection a cheaper denial of service than the hole it closed. + """ + body = _multipart_body(b"z" * 2048) + scope = _upload_scope(query_string) + middleware = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_body_bytes=10 * len(body), max_concurrent=1) + + with patch.object(videos, "ApiDependencies", _forbidden_board_deps()): + status, headers, sent_at_response, total = _drive_upload(middleware, body, scope) + + assert status == expected_status + assert (b"connection", b"close") in headers + # The point of closing rather than draining: answered immediately, slot given straight back. + assert sent_at_response < total + assert middleware._active == 0 + + +def test_close_does_not_trust_a_client_supplied_content_length(): + """`Content-Length: 0` must not be taken as proof the client has finished sending. + + h11 accepts `Content-Length: 0` alongside `Transfer-Encoding: chunked` — the chunked + framing wins — so a client that sends both could otherwise suppress the close and then + stream indefinitely with no lease held and none of the caps applying (they all live in + `limited_receive`, which an app that answered early never calls again). + """ + body = _multipart_body(b"z" * 2048) + scope = _upload_scope(b"video_category=general&is_intermediate=false&board_id=not-mine") + scope["headers"] = [*scope["headers"], (b"content-length", b"0"), (b"transfer-encoding", b"chunked")] + middleware = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_body_bytes=10 * len(body), max_concurrent=1) + + with patch.object(videos, "ApiDependencies", _forbidden_board_deps()): + status, headers, _sent, _total = _drive_upload(middleware, body, scope) + + assert status == 403 + assert (b"connection", b"close") in headers + + +@pytest.mark.parametrize( + "limits,chunk_delay", + [ + ({"max_body_bytes": 128}, 0.0), + ({"max_body_bytes": 10**6, "max_upload_duration_seconds": 0.0}, 0.0), + ({"max_body_bytes": 10**6, "idle_timeout_seconds": 0.01}, 0.2), + ], + ids=["byte-cap", "duration-cap", "idle-timeout"], +) +def test_aborting_a_still_uploading_client_closes_the_connection(limits: dict[str, Any], chunk_delay: float): + """The byte cap, duration cap and idle timeout all synthesize a disconnect *because* the + client is still uploading, so the response they produce must close the connection. + + Marking the body finished on those paths would suppress exactly the close they need. + """ + body = _multipart_body(b"z" * 2048) + middleware = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_concurrent=1, **limits) + + with patch.object(videos, "ApiDependencies", MagicMock()): + status, headers, sent_at_response, total = _drive_upload(middleware, body, chunk_delay=chunk_delay) + + assert sent_at_response < total # aborted mid-body + # A client-side abort, reported as such rather than as a server fault. + assert status == 400 + assert (b"connection", b"close") in headers + assert middleware._active == 0 + + +def test_close_is_not_sent_on_http2(): + """`connection` is hop-by-hop and illegal in HTTP/2+.""" + body = _multipart_body(b"z" * 2048) + scope = _upload_scope(b"video_category=general&is_intermediate=false&board_id=not-mine") + scope["http_version"] = "2" + middleware = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_body_bytes=10 * len(body), max_concurrent=1) + + with patch.object(videos, "ApiDependencies", _forbidden_board_deps()): + status, headers, _sent, _total = _drive_upload(middleware, body, scope) + assert status == 403 + assert not any(name == b"connection" for name, _ in headers) + + # The middleware's own rejections are gated the same way. + saturated = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_body_bytes=10 * len(body), max_concurrent=0) + status, headers, _sent, _total = _drive_upload(saturated, body, scope) + assert status == 429 + assert not any(name == b"connection" for name, _ in headers) + + +def test_close_replaces_rather_than_duplicates_an_existing_connection_header(): + """Two `connection` headers on one response is not a well-formed message.""" + body = _multipart_body(b"z" * 2048) + + async def app(scope: Any, receive: Any, send: Any) -> None: + await send( + { + "type": "http.response.start", + "status": 403, + "headers": [(b"content-length", b"0"), (b"connection", b"keep-alive")], + } + ) + await send({"type": "http.response.body", "body": b""}) + + middleware = VideoUploadLimitASGIMiddleware(app, max_body_bytes=10 * len(body), max_concurrent=1) + _status, headers, _sent, _total = _drive_upload(middleware, body) + + assert [value for name, value in headers if name == b"connection"] == [b"close"] + + +def test_unauthenticated_rejection_closes_the_connection(): + """The 401 answers without reading the body too.""" + body = _multipart_body(b"z" * 2048) + middleware = VideoUploadLimitASGIMiddleware( + _real_upload_app(), + max_body_bytes=10 * len(body), + max_concurrent=1, + identify_user=lambda scope: (False, None), + ) + + status, headers, _sent, _total = _drive_upload(middleware, body) + assert status == 401 + assert (b"connection", b"close") in headers + + +def test_completed_upload_does_not_close_the_connection(): + """Control: a client that finished sending gets a normal keep-alive response. + + Without this, the close would be unconditional and every upload would cost a new + connection. + """ + body = _multipart_body(b"z" * 2048) + deps = MagicMock() + deps.invoker.services.videos.create.return_value = SimpleNamespace( + video_url="/api/v1/videos/i/v.mp4/full", video_name="v.mp4" + ) + middleware = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_body_bytes=10 * len(body), max_concurrent=1) + + with ( + patch.object(videos, "ApiDependencies", deps), + patch.object(videos, "_is_mp4_file", return_value=False), + ): + status, headers, sent_at_response, total = _drive_upload(middleware, body) + + # 415 from the container check — reached only after the whole body was read. + assert status == 415 + assert sent_at_response == total + assert not any(name == b"connection" for name, _ in headers) + assert middleware._active == 0 + + +def test_rejections_the_middleware_makes_itself_also_close_the_connection(): + """The 429/413 paths answer without reading the body too.""" + body = _multipart_body(b"z" * 2048) + middleware = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_body_bytes=10 * len(body), max_concurrent=0) + + status, headers, _sent, _total = _drive_upload(middleware, body) + assert status == 429 + assert (b"connection", b"close") in headers + + oversized = VideoUploadLimitASGIMiddleware(_real_upload_app(), max_body_bytes=8, max_concurrent=1) + scope = _upload_scope() + scope["headers"] = [*scope["headers"], (b"content-length", str(len(body)).encode())] + status, headers, _sent, _total = _drive_upload(oversized, body, scope) + assert status == 413 + assert (b"connection", b"close") in headers + + +def test_upload_video_requires_a_file_part(): + async def run_immediately(func: Any, *args: Any): + return func(*args) + + body = f'--{BOUNDARY}\r\nContent-Disposition: form-data; name="metadata"\r\n\r\n{{}}\r\n--{BOUNDARY}--\r\n'.encode() + + with ( + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + pytest.raises(HTTPException) as error, + ): + _run_upload(_fake_upload_request(body)) + + assert error.value.status_code == 422 + + +def test_upload_video_rejects_oversized_file_part_mid_stream(monkeypatch: pytest.MonkeyPatch): + """The size cap fires while the body streams, not after it has all landed on disk.""" + monkeypatch.setattr(videos, "MAX_UPLOAD_SIZE", 512) + written = 0 + + async def run_immediately(func: Any, *args: Any): + return func(*args) + + real_named_tmp = tempfile.NamedTemporaryFile + captured_handles: list[Any] = [] + + def recording_named_tmp(*args: Any, **kwargs: Any): + handle = real_named_tmp(*args, **kwargs) + real_write = handle.write + + def counting_write(chunk: bytes) -> int: + nonlocal written + written += len(chunk) + return real_write(chunk) + + handle.write = counting_write + captured_handles.append(handle) + return handle + + try: + with ( + patch("invokeai.app.api.routers.videos.tempfile.NamedTemporaryFile", side_effect=recording_named_tmp), + patch("invokeai.app.api.routers.videos.run_in_threadpool", side_effect=run_immediately), + pytest.raises(HTTPException) as error, + ): + _run_upload(_fake_upload_request(_multipart_body(b"y" * 4096), chunk_size=64)) + + assert error.value.status_code == 413 + # Only the bytes up to the cap (plus at most one chunk) were ever written. + assert written <= 512 + 64 + finally: + for handle in captured_handles: + handle.close() + Path(handle.name).unlink(missing_ok=True) + + @pytest.mark.parametrize("preserve", [True, False], ids=["preserve", "strip"]) def test_route_matching_is_root_path_aware(preserve: bool): """Behind a sub-path proxy the public path carries the prefix; the size cap must still