diff --git a/fastvideo/entrypoints/openai/media_policy.py b/fastvideo/entrypoints/openai/media_policy.py index 34f62030c..90389296f 100644 --- a/fastvideo/entrypoints/openai/media_policy.py +++ b/fastvideo/entrypoints/openai/media_policy.py @@ -26,6 +26,7 @@ import httpx MAX_REDIRECTS = 5 READ_TIMEOUT_S = 60.0 DATA_URL_HEADER_MAX_CHARS = 256 +DATA_URL_CHUNK_CHARS = 1 << 16 # at least 3, so a percent escape fits in one chunk class MediaPolicyError(ValueError): @@ -94,14 +95,37 @@ def data_url_media_type(source: str) -> str: return parameters[0].strip().lower() +def _base64_payload(source: str, start: int, max_bytes: int) -> str: + """``source[start:]`` without percent escapes or whitespace, normalized in bounded chunks. + + Stops as soon as the payload must decode to more than ``max_bytes``, so an oversized request costs memory + in proportion to the limit, not to its own size. + """ + chunks: list[str] = [] + length = 0 + while start < len(source): + end = min(start + DATA_URL_CHUNK_CHARS, len(source)) + escape = source.find("%", max(start + 1, end - 2), end) if end < len(source) else -1 + if escape >= 0: + end = escape # an escape that may continue past the chunk starts the next one + # Same leniency as SGLang Diffusion: percent escapes and line-wrapped payloads decode. + chunk = source[start:end] + if "%" in chunk: + chunk = unquote(chunk) + chunk = "".join(chunk.split()) + length += len(chunk) + # Base64 decodes to 3 bytes per 4 characters, less at most 2 bytes of padding. + if length // 4 * 3 - 2 > max_bytes: + raise _too_large(max_bytes) + chunks.append(chunk) + start = end + return "".join(chunks) + + def decode_data_url(source: str, target: str, *, max_bytes: int) -> None: """Decode a base64 ``data:`` URL to ``target``; nothing is written if it is invalid or over ``max_bytes``.""" data_url_media_type(source) - payload = source[source.find(",") + 1:] - # Same leniency as SGLang Diffusion: percent escapes and line-wrapped payloads decode. - if "%" in payload: - payload = unquote(payload) - payload = "".join(payload.split()) + payload = _base64_payload(source, source.find(",") + 1, max_bytes) padding = len(payload) - len(payload.rstrip("=")) if len(payload) // 4 * 3 - padding > max_bytes: raise _too_large(max_bytes) diff --git a/fastvideo/entrypoints/openai/request_adapter.py b/fastvideo/entrypoints/openai/request_adapter.py index 4f07146ab..7d112dde9 100644 --- a/fastvideo/entrypoints/openai/request_adapter.py +++ b/fastvideo/entrypoints/openai/request_adapter.py @@ -152,7 +152,8 @@ async def _materialize_media(source: str, source = source.strip() if not source: raise RequestAdaptationError(f"{kind.capitalize()} references must not be empty.") - if source.lower().startswith(("http://", "https://")): + # Only the scheme is lower-cased: copying a whole data: URL would double its memory. + if source[:8].lower().startswith(("http://", "https://")): suffix = os.path.splitext(urlparse(source).path)[1].lower() if suffix not in _MEDIA_SUFFIXES[kind]: suffix = _DEFAULT_SUFFIXES[kind] @@ -164,7 +165,7 @@ async def _materialize_media(source: str, except Exception as error: raise RequestAdaptationError(f"Unable to fetch {kind} reference: {error}") from error return target, True - if source.lower().startswith("data:"): + if source[:5].lower() == "data:": return await _decode_data_media(source, target_base, kind, max_bytes=max_bytes), True return _local_media(source, policy), False @@ -213,7 +214,7 @@ async def prepare_reference_media( local_path = track(await save_base64_image_to_path(source, target)) except Exception as error: raise RequestAdaptationError(f"Unable to decode image reference: {error}") from error - elif source.lower().startswith("data:"): + elif source[:5].lower() == "data:": raise RequestAdaptationError("Image references need a data:image/... URL.") elif is_remote(source): local_path = track(await _materialize_image_url(source, target, max_bytes=max_bytes, policy=policy)) diff --git a/fastvideo/tests/entrypoints/test_openai_serving_engine.py b/fastvideo/tests/entrypoints/test_openai_serving_engine.py index 5cf069684..fa3ed59f1 100644 --- a/fastvideo/tests/entrypoints/test_openai_serving_engine.py +++ b/fastvideo/tests/entrypoints/test_openai_serving_engine.py @@ -17,6 +17,7 @@ from fastapi.testclient import TestClient from fastvideo.api.parser import parse_config from fastvideo.api.schema import GenerationRequest +from fastvideo.entrypoints.openai import media_policy, request_adapter from fastvideo.entrypoints.openai.media_policy import MediaPolicy from fastvideo.entrypoints.openai.protocol import UrlAudioReference, VideoGenerationRequest from fastvideo.entrypoints.openai.request_adapter import ( @@ -293,11 +294,17 @@ def test_inline_video_and_audio_references_are_decoded_in_order(tmp_path: Path) assert not _uploads(tmp_path) -def test_inline_reference_payload_may_be_wrapped_and_percent_escaped(tmp_path: Path) -> None: - track = b"wrapped-track" * 8 - encoded = base64.b64encode(track).decode().replace("=", "%3D") - wrapped = "\n".join(encoded[start:start + 16] for start in range(0, len(encoded), 16)) - request = VideoGenerationRequest(prompt="waves", audio_reference={"audio_url": f"data:audio/flac;base64,{wrapped}"}) +@pytest.mark.parametrize("chunk_chars", [media_policy.DATA_URL_CHUNK_CHARS, 3, 4, 5, 7]) +def test_inline_reference_payload_may_be_wrapped_and_percent_escaped(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + chunk_chars: int) -> None: + # Small chunks put chunk boundaries inside percent escapes and line breaks. + monkeypatch.setattr(media_policy, "DATA_URL_CHUNK_CHARS", chunk_chars) + track = bytes(range(256)) * 2 + encoded = base64.b64encode(track).decode() + lines = [encoded[start:start + 16] for start in range(0, len(encoded), 16)] + wrapped = "".join(line + ("%0A" if number % 2 else "\r\n") for number, line in enumerate(lines)) + escaped = wrapped.replace("+", "%2B").replace("/", "%2f").replace("=", "%3D") + request = VideoGenerationRequest(prompt="waves", audio_reference={"audio_url": f"data:audio/flac;base64,{escaped}"}) created = asyncio.run(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy())) assert Path(created[0]).read_bytes() == track assert created[0].endswith(".flac") @@ -333,6 +340,24 @@ def test_inline_reference_over_the_media_limit_is_refused(tmp_path: Path) -> Non assert Path(created[0]).read_bytes() == b"1234" +def test_oversized_inline_payload_is_refused_before_it_is_all_read(tmp_path: Path, + monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(media_policy, "DATA_URL_CHUNK_CHARS", 8) + unescaped: list[int] = [] + original_unquote = media_policy.unquote + + def counting_unquote(chunk: str) -> str: + unescaped.append(len(chunk)) + return original_unquote(chunk) + + monkeypatch.setattr(media_policy, "unquote", counting_unquote) + source = "data:video/mp4;base64," + "QU%4AB\n" * 10_000 + with pytest.raises(media_policy.MediaPolicyError, match="4-byte limit"): + media_policy.decode_data_url(source, str(tmp_path / "clip.mp4"), max_bytes=4) + assert max(unescaped) <= 8 and sum(unescaped) < 64 + assert not (tmp_path / "clip.mp4").exists() + + def test_decoded_references_are_removed_when_a_later_reference_fails(tmp_path: Path) -> None: request = VideoGenerationRequest( prompt="waves", @@ -344,6 +369,41 @@ def test_decoded_references_are_removed_when_a_later_reference_fails(tmp_path: P assert not _uploads(tmp_path) +def test_cancelled_decoding_removes_its_file_once_the_thread_finishes(tmp_path: Path, + monkeypatch: pytest.MonkeyPatch) -> None: + asyncio.run(_assert_cancelled_decoding_removes_its_file(tmp_path, monkeypatch)) + + +async def _assert_cancelled_decoding_removes_its_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + started, release, finished = threading.Event(), threading.Event(), threading.Event() + written: list[bool] = [] + + def decode_after_release(source: str, target: str, *, max_bytes: int) -> None: + started.set() + release.wait(timeout=5) + try: + media_policy.decode_data_url(source, target, max_bytes=max_bytes) + written.append(os.path.exists(target)) + finally: + finished.set() + + monkeypatch.setattr(request_adapter, "decode_data_url", decode_after_release) + request = VideoGenerationRequest(prompt="waves", video_reference={"video_url": _data_url("video/mp4", b"clip")}) + preparing = asyncio.create_task(prepare_reference_media("request", request, str(tmp_path), policy=MediaPolicy())) + assert await asyncio.to_thread(started.wait, 5) + preparing.cancel() + with pytest.raises(asyncio.CancelledError): + await preparing + # The thread writes its file only after the request is gone, so cleanup has to wait for the thread. + release.set() + assert await asyncio.to_thread(finished.wait, 5) + assert written == [True] + deadline = time.monotonic() + 5 + while _uploads(tmp_path) and time.monotonic() < deadline: + await asyncio.sleep(0.01) + assert not _uploads(tmp_path) + + def test_local_audio_paths_need_a_configured_root(tmp_path: Path) -> None: track = tmp_path / "track.wav" track.write_bytes(b"x")