diff --git a/projects/fal/src/fal/toolkit/file/_local_uploader.py b/projects/fal/src/fal/toolkit/file/_local_uploader.py index aed464156..b5b20542e 100644 --- a/projects/fal/src/fal/toolkit/file/_local_uploader.py +++ b/projects/fal/src/fal/toolkit/file/_local_uploader.py @@ -45,6 +45,11 @@ class LocalUploadRejected(LocalUploadRefused): """The uploader is full or draining and may accept a replay shortly.""" +class LocalUploadInvalid(LocalUploadRefused): + """The uploader refused this upload's content or metadata (HTTP 413 or 422), + not for lack of capacity.""" + + def upload( file_name: str, content: bytes | Path, @@ -176,6 +181,8 @@ def _request( status = response.status_code if status == 429: raise LocalUploadRejected("Local uploader returned HTTP 429.") + if status in (413, 422): + raise LocalUploadInvalid(f"Local uploader returned HTTP {status}.") if status != expected_status: message = f"Local uploader returned HTTP {status}." raise (LocalUploadRefused if status < 500 else LocalUploadError)(message) diff --git a/projects/fal/src/fal/toolkit/file/_upload_policy.py b/projects/fal/src/fal/toolkit/file/_upload_policy.py index 262d5c90a..cdd2d640f 100644 --- a/projects/fal/src/fal/toolkit/file/_upload_policy.py +++ b/projects/fal/src/fal/toolkit/file/_upload_policy.py @@ -3,6 +3,8 @@ Private and meant to be temporary: this is the SDK-side stopgap until infra ships a first-class upload destination, at which point this module goes away. Hence the leading underscore -- nothing here is a supported public API. +Runners opted in to the node-local uploader hand the upload to it instead and +use this path only when it refuses (``providers.local.save_with_policy``). The ``x-app-fal-upload-policy`` header carries an S3 pre-signed POST policy -- the JSON from boto3's ``generate_presigned_post()``:: @@ -556,12 +558,9 @@ def _validate_size(nbytes: int) -> None: ) -def _prepare_upload( - policy: UploadPolicy, - file_name: str, - content_type: str, -) -> tuple[str, dict[str, PolicyFieldValue]]: - """Resolve the destination URL and the POST form fields. +def _upload_file_name(policy: UploadPolicy, file_name: str, content_type: str) -> str: + """Check an output against the policy and return the name its key is built + from, the value that replaces ``${filename}``. The uuid prefix keeps concurrent uploads from colliding on a shared policy, which is signed once and typically reused for a whole request. @@ -581,26 +580,34 @@ def _prepare_upload( # generation. Re-checked here because upload_* is exported, so a caller can # hand-build an UploadPolicy that never passed parse_upload_policy. _validate_policy_url(policy.url) - fields = dict(policy.fields) - key_template = _require_key_template(fields) + _require_key_template(policy.fields) + + # Signed fields cannot be rewritten: the Content-Type is injected only if + # the caller left it out, otherwise it must agree or S3 answers an opaque 403. + for name, value in policy.fields.items(): + if _reserved_name(name) == "content-type" and value != content_type: + raise UploadPolicyInputError( + f"Invalid {UPLOAD_POLICY_KEY} header: fields.{name} is " + f"{value!r} but the file being uploaded is " + f"{content_type!r}; a signed Content-Type cannot be changed" + ) + return f"{uuid4().hex}-{file_name}" + - upload_file_name = f"{uuid4().hex}-{file_name}" - final_key = key_template.replace( - UPLOAD_POLICY_FILENAME_PLACEHOLDER, upload_file_name +def _prepare_upload( + policy: UploadPolicy, + file_name: str, + content_type: str, +) -> tuple[str, dict[str, PolicyFieldValue]]: + """Resolve the destination URL and the POST form fields.""" + name = _upload_file_name(policy, file_name, content_type) + fields = dict(policy.fields) + final_key = _require_key_template(fields).replace( + UPLOAD_POLICY_FILENAME_PLACEHOLDER, name ) fields["key"] = final_key - - # Signed fields cannot be rewritten: inject Content-Type only if the caller - # left it out, otherwise require agreement or S3 answers an opaque 403. - existing = [name for name in fields if _reserved_name(name) == "content-type"] - if not existing: + if not any(_reserved_name(field) == "content-type" for field in fields): fields["Content-Type"] = content_type - elif fields[existing[0]] != content_type: - raise UploadPolicyInputError( - f"Invalid {UPLOAD_POLICY_KEY} header: fields.{existing[0]} is " - f"{fields[existing[0]]!r} but the file being uploaded is " - f"{content_type!r}; a signed Content-Type cannot be changed" - ) encoded_key = urllib.parse.quote(final_key.lstrip("/"), safe="/~") access_url = f"{policy.url.rstrip('/')}/{encoded_key}" diff --git a/projects/fal/src/fal/toolkit/file/file.py b/projects/fal/src/fal/toolkit/file/file.py index a8336a79f..5040ee88a 100644 --- a/projects/fal/src/fal/toolkit/file/file.py +++ b/projects/fal/src/fal/toolkit/file/file.py @@ -29,8 +29,6 @@ UPLOAD_POLICY_KEY, UploadPolicyInputError, get_upload_policy, - upload_bytes_with_policy, - upload_path_with_policy, ) from fal.toolkit.file.providers.fal import ( # Re-exported for backwards compatibility; both resolve to FalFileRepositoryV3. @@ -42,7 +40,7 @@ MultipartUploadV3, ) from fal.toolkit.file.providers.gcp import GoogleStorageRepository -from fal.toolkit.file.providers.local import LocalFileRepository +from fal.toolkit.file.providers.local import LocalFileRepository, save_with_policy from fal.toolkit.file.providers.r2 import R2Repository from fal.toolkit.file.types import FileData, FileRepository, RepositoryId from fal.toolkit.utils.download_utils import download_file @@ -301,7 +299,7 @@ def from_bytes( upload_policy = get_upload_policy(request) if upload_policy is not None: _require_no_repository_conflict(repository) - url = upload_bytes_with_policy( + url = save_with_policy( upload_policy, fdata.file_name, data, fdata.content_type ) return cls( @@ -405,19 +403,15 @@ def from_path( if multipart is not None else file_size > MultipartUploadV3.MULTIPART_THRESHOLD ) - if stage_to_disk: - # Stage and stream in the background, off the request thread. - url = upload_path_with_policy( - upload_policy, file_path, file_path.name, content_type - ) - file_data = None - else: - # Reusing these bytes for file_data avoids a second full read - # and the staging copy for the common small-output case. - file_data = file_path.read_bytes() - url = upload_bytes_with_policy( - upload_policy, file_path.name, file_data, content_type - ) + # Reusing small files' bytes for file_data avoids a second full read + # and the staging copy for the common small-output case. + file_data = None if stage_to_disk else file_path.read_bytes() + url = save_with_policy( + upload_policy, + file_path.name, + file_path if file_data is None else file_data, + content_type, + ) return cls( url=url, file_data=file_data, diff --git a/projects/fal/src/fal/toolkit/file/providers/local.py b/projects/fal/src/fal/toolkit/file/providers/local.py index 923475dc6..1176f2964 100644 --- a/projects/fal/src/fal/toolkit/file/providers/local.py +++ b/projects/fal/src/fal/toolkit/file/providers/local.py @@ -1,22 +1,36 @@ """Repository adapter for opt-in, node-local uploads. -A submission the uploader refused is sent directly to the CDN instead. A -failure after the uploader may have accepted the bytes is raised, because a -second upload could publish a duplicate. Small files retain bytes for -File.as_bytes(). +A submission the uploader refused is sent directly to the CDN instead, or +in-process to the caller's bucket under an upload policy. A failure after the +uploader may have accepted the bytes is raised, because a second upload could +publish a duplicate. Small files retain bytes for File.as_bytes(). """ from __future__ import annotations +import json from functools import wraps from pathlib import Path from typing import Iterable +from urllib.parse import quote +from fal.flags import bool_envvar from fal.toolkit.file._local_uploader import ( + LocalUploadInvalid, LocalUploadRefused, upload, upload_stream, ) +from fal.toolkit.file._upload_policy import ( + UPLOAD_POLICY_KEY, + UploadPolicy, + UploadPolicyError, + _current_request_id, + _upload_file_name, + _validate_size, + upload_bytes_with_policy, + upload_path_with_policy, +) from fal.toolkit.file.providers.fal import ( FalFileRepositoryV3, MultipartUploadV3, @@ -119,3 +133,45 @@ def save_stream( return upload_stream( file_name, chunks, headers, wait_for_completion=wait_for_completion ) + + +def save_with_policy( + policy: UploadPolicy, file_name: str, content: bytes | Path, content_type: str +) -> str: + """Upload to the caller's bucket and return the URL the object will have. + + Opted-in runners hand the upload to the node-local uploader. Any other + runner, or a refusal, uploads in-process. Neither waits for the bucket. + """ + if bool_envvar("FAL_USE_LOCAL_UPLOADER"): + try: + return _upload_with_policy(policy, file_name, content, content_type) + except LocalUploadInvalid as exc: + raise UploadPolicyError( + f"Upload via {UPLOAD_POLICY_KEY} was refused: {exc}" + ) from None + except LocalUploadRefused as exc: + print(f"{exc} Uploading in-process.") + if isinstance(content, bytes): + return upload_bytes_with_policy(policy, file_name, content, content_type) + return upload_path_with_policy(policy, content, file_name, content_type) + + +def _upload_with_policy( + policy: UploadPolicy, file_name: str, content: bytes | Path, content_type: str +) -> str: + _validate_size( + len(content) if isinstance(content, bytes) else content.stat().st_size + ) + # Percent-encoded so a non-ASCII name can travel in the header; the + # uploader decodes it, so the key matches the in-process path. + name = quote(_upload_file_name(policy, file_name, content_type), safe="") + # No credentials: an uploader without policy support then refuses the + # upload rather than send it to the fal CDN. + headers = { + "Content-Type": content_type, + UPLOAD_POLICY_KEY: json.dumps({"url": policy.url, "fields": policy.fields}), + } + if request_id := _current_request_id(): + headers["X-Fal-Request-ID"] = request_id + return upload(name, content, headers) diff --git a/projects/fal/tests/integration/toolkit/test_local_uploader.py b/projects/fal/tests/integration/toolkit/test_local_uploader.py index bf4483f37..af6c2c36a 100644 --- a/projects/fal/tests/integration/toolkit/test_local_uploader.py +++ b/projects/fal/tests/integration/toolkit/test_local_uploader.py @@ -6,6 +6,7 @@ from __future__ import annotations +import base64 import hashlib import json import os @@ -19,13 +20,17 @@ from types import SimpleNamespace from unittest.mock import Mock from urllib.parse import urlparse +from uuid import UUID import httpx import pytest +from starlette.requests import Request from fal.toolkit.file import File, _local_uploader +from fal.toolkit.file import _upload_policy as upload_policy from fal.toolkit.file._local_uploader import LocalUploadError from fal.toolkit.file.providers import fal as remote +from fal.toolkit.file.providers import local from fal.toolkit.file.providers.local import LocalFileRepository pytestmark = pytest.mark.skipif( @@ -56,7 +61,9 @@ def uploader(monkeypatch): release = threading.Event() receiving = threading.Event() done = threading.Event() + connected = threading.Event() state = { + "connects": [], "parts": {}, "multipart": False, "headers": None, @@ -129,6 +136,14 @@ def do_GET(self): } ) + def do_CONNECT(self): + # The uploader's HTTPS proxy: records where a policy upload goes, + # then refuses the tunnel, since no fake can pass as an S3 host. + state["connects"].append(self.path) + connected.set() + self.send_response(502) + self.end_headers() + def log_message(self, *args): pass @@ -153,6 +168,7 @@ def log_message(self, *args): "CDN_UPLOADER_MAX_QUEUE_AGE_SECONDS": "60", "CDN_UPLOADER_FAILED_RETENTION_SECONDS": "60", "CDN_UPLOADER_SWEEP_INTERVAL_SECONDS": "1", + "HTTPS_PROXY": origin, } env.pop("CDN_UPLOADER_CDN_BASE_URL", None) with (root / "uploader.log").open("w+") as log: @@ -165,7 +181,11 @@ def log_message(self, *args): monkeypatch.setenv("FAL_USE_LOCAL_UPLOADER", "1") monkeypatch.setenv("FAL_KEY", "local:test") yield SimpleNamespace( - release=release, receiving=receiving, done=done, state=state + release=release, + receiving=receiving, + done=done, + connected=connected, + state=state, ) finally: release.set() @@ -269,3 +289,47 @@ def test_generated_stream_waits_for_completion(uploader): assert url.endswith("/file/1") assert uploader.done.is_set() assert uploader.state["parts"] == {1: b"hello world"} + + +def _policy_request(policy): + header = (upload_policy.UPLOAD_POLICY_KEY.encode(), json.dumps(policy).encode()) + return Request({"type": "http", "headers": [header]}) + + +def test_upload_policy_posts_to_the_caller_bucket(uploader, monkeypatch): + monkeypatch.delenv("FAL_KEY") # a policy upload needs no fal credentials + monkeypatch.setattr(upload_policy, "uuid4", lambda: UUID(int=7)) + in_process = Mock() + monkeypatch.setattr(local, "upload_bytes_with_policy", in_process) + policy = { + "url": "https://bucket.s3.us-west-1.amazonaws.com/", + "fields": {"key": "outputs/${filename}"}, + } + result = File.from_bytes( + b"hello", "text/plain", "vidéo (1).txt", request=_policy_request(policy) + ) + expected, _ = upload_policy._prepare_upload( + upload_policy.UploadPolicy(**policy), "vidéo (1).txt", "text/plain" + ) + assert result.url == expected + assert uploader.connected.wait(10) + assert uploader.state["connects"][0] == "bucket.s3.us-west-1.amazonaws.com:443" + assert uploader.state["reservations"] == 0 + in_process.assert_not_called() + + +def test_expired_upload_policy_is_refused(uploader, monkeypatch): + in_process = Mock() + monkeypatch.setattr(local, "upload_bytes_with_policy", in_process) + document = {"expiration": "2020-01-01T00:00:00Z", "conditions": []} + policy = { + "url": "https://bucket.s3.amazonaws.com/", + "fields": { + "key": "outputs/${filename}", + "policy": base64.b64encode(json.dumps(document).encode()).decode(), + }, + } + with pytest.raises(upload_policy.UploadPolicyError, match="HTTP 422"): + File.from_bytes(b"hello", request=_policy_request(policy)) + in_process.assert_not_called() + assert not uploader.state["connects"] diff --git a/projects/fal/tests/unit/toolkit/file/test_local_uploader.py b/projects/fal/tests/unit/toolkit/file/test_local_uploader.py index e62949b2c..219f27a6d 100644 --- a/projects/fal/tests/unit/toolkit/file/test_local_uploader.py +++ b/projects/fal/tests/unit/toolkit/file/test_local_uploader.py @@ -4,6 +4,8 @@ import sys from types import SimpleNamespace from unittest.mock import Mock +from urllib.parse import quote, unquote +from uuid import UUID import httpx import pytest @@ -11,6 +13,7 @@ from fal.auth import AuthCredentials from fal.toolkit.file import _local_uploader +from fal.toolkit.file import _upload_policy as upload_policy from fal.toolkit.file import file as files from fal.toolkit.file._local_uploader import LocalUploadError, LocalUploadRefused from fal.toolkit.file._upload_policy import UPLOAD_POLICY_KEY @@ -137,26 +140,165 @@ def test_explicit_repository_objects_and_destinations_are_preserved( assert not local_upload -def test_upload_policy_overrides_explicit_local_repository(monkeypatch, local_upload): - policy = json.dumps( - { - "url": "https://bucket.s3.amazonaws.com/", - "fields": {"key": "outputs/${filename}"}, - } +POLICY = { + "url": "https://bucket.s3.amazonaws.com/", + "fields": {"key": "outputs/${filename}", "policy": "signed", "x-amz-flag": True}, +} + + +def _policy_request(): + header = (UPLOAD_POLICY_KEY.encode(), json.dumps(POLICY).encode()) + return Request({"type": "http", "headers": [header]}) + + +@pytest.fixture +def policy_upload(monkeypatch, transport): + """Opt in; record in-process uploads and fail on any fal CDN upload.""" + monkeypatch.setenv("FAL_USE_LOCAL_UPLOADER", "1") + monkeypatch.setattr(SLEEP, Mock()) + in_process = Mock(return_value="https://bucket.s3.amazonaws.com/outputs/direct") + monkeypatch.setattr(local, "upload_bytes_with_policy", in_process) + monkeypatch.setattr(local, "upload_path_with_policy", in_process) + for repository in (remote.FalFileRepositoryV3, remote.FalFileRepository): + for method in ("save", "save_file"): + monkeypatch.setattr(repository, method, Mock(side_effect=AssertionError)) + + def install(response): + respond = isinstance(response, (httpx.Response, type)) + handler = _respond(response) if respond else response + return transport(handler), in_process + + return install + + +def _accept_like_the_uploader(request): + """Substitute ${filename} and build file_url as isolate-cloud's s3_post.rs does.""" + policy = json.loads(request.headers[UPLOAD_POLICY_KEY]) + key = policy["fields"]["key"].replace( + "${filename}", unquote(request.headers["x-fal-file-name"]) + ) + file_url = policy["url"].rstrip("/") + "/" + quote(key, safe="/~") + return httpx.Response(202, json={**ACCEPTED, "file_url": file_url}) + + +@pytest.mark.parametrize( + "file_name,header", + [ + ("cat.txt", "cat.txt"), + ("vidéo (1).mp4", "vid%C3%A9o%20%281%29.mp4"), + ], +) +@pytest.mark.parametrize("explicit", [False, True]) +def test_upload_policy_goes_through_the_local_uploader( + monkeypatch, policy_upload, caller, explicit, file_name, header +): + requests, in_process = policy_upload(_accept_like_the_uploader) + caller({"x-fal-cdn-token": "caller-token"}, "request-id") + monkeypatch.setattr(upload_policy, "get_current_app", remote.get_current_app) + monkeypatch.setattr(upload_policy, "uuid4", lambda: UUID(int=7)) + kwargs = {"repository": local.LocalFileRepository()} if explicit else {} + result = files.File.from_bytes( + b"hello", "text/plain", file_name, request=_policy_request(), **kwargs ) - request = Request( - {"type": "http", "headers": [(UPLOAD_POLICY_KEY.encode(), policy.encode())]} + in_process_url, fields = upload_policy._prepare_upload( + upload_policy.UploadPolicy(**POLICY), file_name, "text/plain" + ) + assert result.url == in_process_url + assert result.as_bytes() == b"hello" + (request,) = requests + assert request.url.path == "/uploads" + assert request.content == b"hello" + assert json.loads(request.headers[UPLOAD_POLICY_KEY]) == POLICY + assert request.headers["x-fal-file-name"] == f"{UUID(int=7).hex}-{header}" + assert fields["key"] == f"outputs/{UUID(int=7).hex}-{file_name}" + assert request.headers["content-type"] == "text/plain" + assert request.headers["x-fal-request-id"] == "request-id" + # Without credentials an uploader that predates policies cannot use the CDN. + assert "authorization" not in request.headers + assert "x-fal-cdn-token" not in request.headers + in_process.assert_not_called() + + +@pytest.mark.parametrize("stage_to_disk", [False, True]) +def test_upload_policy_files_stream_from_disk_when_large( + policy_upload, tmp_path, stage_to_disk +): + requests, _ = policy_upload(httpx.Response(202, json=ACCEPTED)) + path = tmp_path / "out.bin" + path.write_bytes(b"data") + result = files.File.from_path( + path, multipart=stage_to_disk, request=_policy_request() ) - upload = Mock(return_value="https://bucket.s3.amazonaws.com/outputs/file.txt") - monkeypatch.setattr(files, "upload_bytes_with_policy", upload) + assert result.url == ACCEPTED["file_url"] + assert result.file_data == (None if stage_to_disk else b"data") + assert requests[0].content == b"data" + + +@pytest.mark.parametrize( + "response,attempts", + [ + (httpx.Response(429), 3), + (httpx.Response(401), 1), # an uploader that predates policies + (httpx.ConnectError, 1), + ], +) +def test_upload_policy_refusals_upload_in_process(policy_upload, response, attempts): + requests, in_process = policy_upload(response) result = files.File.from_bytes( - b"hello", - file_name="file.txt", - request=request, - repository=local.LocalFileRepository(), + b"hi", "text/plain", "cat.txt", request=_policy_request() ) - assert result.url == upload.return_value - assert not local_upload + assert result.url == in_process.return_value + assert len(requests) == attempts + in_process.assert_called_once_with( + upload_policy.UploadPolicy(**POLICY), "cat.txt", b"hi", "text/plain" + ) + + +@pytest.mark.parametrize("status", [413, 422]) +def test_upload_policy_rejections_are_raised(policy_upload, status): + _, in_process = policy_upload(httpx.Response(status, text="secret")) + with pytest.raises(upload_policy.UploadPolicyError) as caught: + files.File.from_bytes(b"hi", request=_policy_request()) + assert caught.value.status_code == 424 + assert "secret" not in str(caught.value) + in_process.assert_not_called() + + +@pytest.mark.parametrize("response", [httpx.Response(503), httpx.ReadError]) +def test_upload_policy_possibly_accepted_failures_are_raised(policy_upload, response): + _, in_process = policy_upload(response) + with pytest.raises(LocalUploadError) as caught: + files.File.from_bytes(b"hi", request=_policy_request()) + assert not isinstance(caught.value, LocalUploadRefused) + in_process.assert_not_called() + + +@pytest.mark.parametrize( + "kwargs", + [ + {"file_name": "a/b.txt"}, + {"content_type": "image/png", "file_name": "b.txt"}, + ], +) +def test_upload_policy_input_is_checked_before_handoff( + monkeypatch, policy_upload, kwargs +): + requests, in_process = policy_upload(httpx.Response(202, json=ACCEPTED)) + signed = {**POLICY, "fields": {**POLICY["fields"], "Content-Type": "text/plain"}} + header = (UPLOAD_POLICY_KEY.encode(), json.dumps(signed).encode()) + request = Request({"type": "http", "headers": [header]}) + with pytest.raises(upload_policy.UploadPolicyInputError): + files.File.from_bytes(b"hi", request=request, **kwargs) + assert not requests + in_process.assert_not_called() + + +def test_upload_policy_without_opt_in_uploads_in_process(monkeypatch, policy_upload): + monkeypatch.delenv("FAL_USE_LOCAL_UPLOADER") + requests, in_process = policy_upload(httpx.Response(202, json=ACCEPTED)) + result = files.File.from_bytes(b"hi", request=_policy_request()) + assert result.url == in_process.return_value + assert not requests @pytest.mark.parametrize("multipart", [False, True, None])