diff --git a/changes/11767.fix.md b/changes/11767.fix.md new file mode 100644 index 00000000000..77ea0fb310d --- /dev/null +++ b/changes/11767.fix.md @@ -0,0 +1 @@ +Fix multi-proxy storage TUS uploads from corrupting shared NFS-mounted files by storing each PATCH as a metadata-tracked per-offset chunk and assembling atomically on completion. diff --git a/src/ai/backend/storage/api/client.py b/src/ai/backend/storage/api/client.py index 634c8d2d330..4caec1a80b9 100644 --- a/src/ai/backend/storage/api/client.py +++ b/src/ai/backend/storage/api/client.py @@ -8,7 +8,7 @@ import logging import os import urllib.parse -from collections.abc import AsyncGenerator, Iterator, Mapping, MutableMapping +from collections.abc import AsyncGenerator, Iterator, Mapping from contextlib import AbstractAsyncContextManager from datetime import UTC, datetime from http import HTTPStatus @@ -38,20 +38,26 @@ ArchiveDownloadQueryParams, ArchiveDownloadTokenData, ) -from ai.backend.common.files import AsyncFileWriter from ai.backend.common.json import dump_json_str from ai.backend.common.metrics.http import build_api_metric_middleware from ai.backend.common.middlewares.exception import general_exception_middleware from ai.backend.common.typed_validators import PydanticJWTValidator -from ai.backend.common.types import BinarySize, VFolderID +from ai.backend.common.types import BinarySize, TusSessionId, VFolderID from ai.backend.logging import BraceStyleAdapter from ai.backend.storage import __version__ from ai.backend.storage.dto.context import StorageRootCtx -from ai.backend.storage.errors import InvalidAPIParameters, UploadOffsetMismatchError +from ai.backend.storage.errors import ( + InvalidAPIParameters, + TusSessionNotFoundError, + UploadChunkExceedsTotalSizeError, + UploadOffsetMismatchError, +) from ai.backend.storage.services.file_stream.zip import ( ZipArchiveStreamReader, ) -from ai.backend.storage.types import SENTINEL +from ai.backend.storage.services.upload.tus_session import TusUploadSession +from ai.backend.storage.services.upload.types import TusUploadSessionArgs +from ai.backend.storage.types import SENTINEL, TusChunkUploadStreamReader from ai.backend.storage.utils import ( CheckParamSource, build_attachment_headers, @@ -60,7 +66,6 @@ if TYPE_CHECKING: from ai.backend.storage.context import RootContext - from ai.backend.storage.volumes.abc import AbstractVolume log = BraceStyleAdapter(logging.getLogger(__spec__.name)) @@ -96,7 +101,7 @@ class UploadTokenData(TypedDict): volume: str vfid: VFolderID relpath: str - session: str + session: TusSessionId size: int @@ -329,14 +334,30 @@ class Params(TypedDict): ) as params: token_data = params["token"] async with ctx.get_volume(token_data["volume"]) as volume: - headers = await prepare_tus_session_headers(request, token_data, volume) + session_dir = ( + volume.mangle_vfpath(token_data["vfid"]) / ".upload" / token_data["session"] + ) + session = TusUploadSession( + TusUploadSessionArgs( + session_dir=session_dir, + session_id=token_data["session"], + total_size=int(token_data["size"]), + valkey_client=ctx.valkey_tus_client, + lock_factory=ctx.tus_lock_factory, + ) + ) + if not await session.exists(): + raise TusSessionNotFoundError + state = await session.read_state() + headers = _prepare_tus_session_headers( + upload_offset=state.committed_offset, + upload_length=int(token_data["size"]), + ) return web.Response(headers=headers) async def tus_upload_part(request: web.Request) -> web.Response: - """ - Perform the chunk upload. - """ + """Perform a TUS PATCH chunk upload.""" ctx: RootContext = request.app["ctx"] secret = ctx.local_config.storage_proxy.secret @@ -361,61 +382,110 @@ class Params(TypedDict): ), ) as params: token_data = params["token"] - async with ctx.get_volume(token_data["volume"]) as volume: - headers = await prepare_tus_session_headers(request, token_data, volume) - vfpath = volume.mangle_vfpath(token_data["vfid"]) - upload_temp_path: Path = vfpath / ".upload" / token_data["session"] + total_size = int(token_data["size"]) - # TUS protocol requires Upload-Offset validation before appending data - upload_offset_header = request.headers.get("Upload-Offset") - if upload_offset_header is None: - raise InvalidAPIParameters( - "Missing required Upload-Offset header for TUS PATCH request" + upload_offset_header = request.headers.get("Upload-Offset") + if upload_offset_header is None: + raise InvalidAPIParameters( + "Missing required Upload-Offset header for TUS PATCH request" + ) + try: + client_offset = int(upload_offset_header) + except ValueError as e: + raise InvalidAPIParameters( + f"Invalid Upload-Offset header value: {upload_offset_header}" + ) from e + if client_offset < 0 or client_offset > total_size: + raise UploadOffsetMismatchError( + f"Upload-Offset {client_offset} is out of range [0, {total_size}]" + ) + + async with ctx.get_volume(token_data["volume"]) as volume: + session_dir = ( + volume.mangle_vfpath(token_data["vfid"]) / ".upload" / token_data["session"] + ) + session = TusUploadSession( + TusUploadSessionArgs( + session_dir=session_dir, + session_id=token_data["session"], + total_size=total_size, + valkey_client=ctx.valkey_tus_client, + lock_factory=ctx.tus_lock_factory, ) + ) + if not await session.exists(): + raise TusSessionNotFoundError + upload_stream = TusChunkUploadStreamReader( + request.content, request.content_type, DEFAULT_CHUNK_SIZE + ) + written = await session.write_temp_chunk(client_offset, upload_stream) try: - client_offset = int(upload_offset_header) - except ValueError as e: - raise InvalidAPIParameters( - f"Invalid Upload-Offset header value: {upload_offset_header}" - ) from e - - actual_offset = int(headers["Upload-Offset"]) - if client_offset != actual_offset: - raise UploadOffsetMismatchError( - f"Upload offset mismatch: expected {actual_offset}, got {client_offset}" + if client_offset + written.length > total_size: + raise UploadChunkExceedsTotalSizeError( + f"Chunk at offset {client_offset} with length {written.length} " + f"exceeds declared size {total_size}" + ) + commit_result = await session.commit_chunk( + offset=client_offset, + chunk_path=written.path, + length=written.length, + sha256=written.sha256, ) - - async with AsyncFileWriter( - target_filename=upload_temp_path, - access_mode="ab", - max_chunks=DEFAULT_INFLIGHT_CHUNKS, - ) as writer: - while not request.content.at_eof(): - chunk = await request.content.read(DEFAULT_CHUNK_SIZE) - await writer.write(chunk) - - current_size = Path(upload_temp_path).stat().st_size - if current_size >= int(token_data["size"]): - parent_dir = vfpath + except BaseException: + if written.path.exists(): + await asyncio.to_thread(written.path.unlink) + raise + + state = commit_result.state + if commit_result.is_final_commit: + parent_dir = volume.mangle_vfpath(token_data["vfid"]) if (dst_dir := params["dst_dir"]) is not None: - parent_dir = vfpath / dst_dir - target_path: Path = parent_dir / token_data["relpath"] - if not target_path.parent.exists(): - target_path.parent.mkdir(parents=True, exist_ok=True) - upload_temp_path.rename(target_path) - try: - loop = asyncio.get_running_loop() - await loop.run_in_executor( - None, - lambda: upload_temp_path.parent.rmdir(), - ) - except OSError: - pass - headers["Upload-Offset"] = str(current_size) + parent_dir = parent_dir / dst_dir + target_path = parent_dir / token_data["relpath"] + await session.assemble(target_path) + await session.cleanup() + + headers = _prepare_tus_session_headers( + upload_offset=state.committed_offset, + upload_length=total_size, + ) return web.Response(status=HTTPStatus.NO_CONTENT, headers=headers) +# TUS-related request/response headers that the storage-proxy accepts and emits, +# advertised via Access-Control-{Allow,Expose}-Headers so browser clients can use +# them across origins. +# +# - Tus-Resumable: TUS protocol version (1.0.0). Required on every TUS request +# and response — the spec's version-negotiation handshake. +# - Upload-Length: total declared size of the file in bytes. Set on session +# creation; echoed in HEAD responses. +# - Upload-Metadata: comma-separated key/value pairs (value is base64-encoded +# per the spec) carrying user-provided metadata such as filename or content +# type. Optional; this implementation does not currently consume it but +# surfaces it to clients. +# - Upload-Offset: byte position from which the next PATCH must continue; +# required on PATCH requests and present on HEAD/PATCH responses. The core +# resumability primitive of TUS. +# - Content-Type: PATCH bodies must be `application/offset+octet-stream` per +# the TUS spec. +_TUS_HEADER_LIST = "Tus-Resumable, Upload-Length, Upload-Metadata, Upload-Offset, Content-Type" + + +def _prepare_tus_session_headers(*, upload_offset: int, upload_length: int) -> dict[str, str]: + return { + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Headers": _TUS_HEADER_LIST, + "Access-Control-Expose-Headers": _TUS_HEADER_LIST, + "Access-Control-Allow-Methods": "*", + "Cache-Control": "no-store", + "Tus-Resumable": "1.0.0", + "Upload-Offset": str(upload_offset), + "Upload-Length": str(upload_length), + } + + async def tus_options(request: web.Request) -> web.Response: """ Let clients discover the supported features of our tus.io server-side implementation. @@ -439,39 +509,6 @@ async def tus_options(request: web.Request) -> web.Response: return web.Response(headers=headers) -async def prepare_tus_session_headers( - _request: web.Request, - token_data: Mapping[str, Any], - volume: AbstractVolume, -) -> MutableMapping[str, str]: - vfpath = volume.mangle_vfpath(token_data["vfid"]) - upload_temp_path = vfpath / ".upload" / token_data["session"] - if not Path(upload_temp_path).exists(): - raise web.HTTPNotFound( - body=dump_json_str( - { - "title": "No such upload session", - "type": "https://api.backend.ai/probs/storage/no-such-upload-session", - }, - ), - content_type="application/problem+json", - ) - headers = {} - headers["Access-Control-Allow-Origin"] = "*" - headers["Access-Control-Allow-Headers"] = ( - "Tus-Resumable, Upload-Length, Upload-Metadata, Upload-Offset, Content-Type" - ) - headers["Access-Control-Expose-Headers"] = ( - "Tus-Resumable, Upload-Length, Upload-Metadata, Upload-Offset, Content-Type" - ) - headers["Access-Control-Allow-Methods"] = "*" - headers["Cache-Control"] = "no-store" - headers["Tus-Resumable"] = "1.0.0" - headers["Upload-Offset"] = str(Path(upload_temp_path).stat().st_size) - headers["Upload-Length"] = str(token_data["size"]) - return headers - - class DownloadHandler: """Handler class for download operations following manager's api_handler pattern. diff --git a/src/ai/backend/storage/api/manager.py b/src/ai/backend/storage/api/manager.py index c3bd27fef30..8a3e41f3282 100644 --- a/src/ai/backend/storage/api/manager.py +++ b/src/ai/backend/storage/api/manager.py @@ -69,6 +69,7 @@ ItemResult, QuotaScopeID, ResultSet, + TusSessionId, VolumeMountableNodeType, ) from ai.backend.logging import BraceStyleAdapter @@ -89,6 +90,7 @@ StorageProxyError, VFolderNotFoundError, ) +from ai.backend.storage.services.upload.types import TusSessionState from ai.backend.storage.types import QuotaConfig, VFolderID from ai.backend.storage.utils import check_params, log_manager_api_entry from ai.backend.storage.watcher import ChownTask, MountTask, UmountTask @@ -1189,6 +1191,14 @@ class Params(TypedDict): ctx: RootContext = request.app["ctx"] async with ctx.get_volume(params["volume"]) as volume: session_id = await volume.prepare_upload(params["vfid"]) + # Register the session in Valkey so HEAD/PATCH can authoritatively tell + # "session exists" from a Valkey state lookup rather than poking the + # filesystem. + tus_session_id = TusSessionId(session_id) + await ctx.valkey_tus_client.set_session_state( + tus_session_id, + TusSessionState.empty(tus_session_id, int(params["size"])).model_dump_json(), + ) token_data = { "op": "upload", "volume": params["volume"], diff --git a/src/ai/backend/storage/errors/__init__.py b/src/ai/backend/storage/errors/__init__.py index dbcaf0fde34..47a76c8e901 100644 --- a/src/ai/backend/storage/errors/__init__.py +++ b/src/ai/backend/storage/errors/__init__.py @@ -64,6 +64,8 @@ ) from .upload import ( ChunkConflictError, + TusSessionNotFoundError, + UploadChunkExceedsTotalSizeError, UploadSessionCorruptedError, ) from .vfolder import ( @@ -96,7 +98,9 @@ "UploadOffsetMismatchError", # upload "ChunkConflictError", + "UploadChunkExceedsTotalSizeError", "UploadSessionCorruptedError", + "TusSessionNotFoundError", # vfolder "VFolderNotFoundError", "InvalidSubpathError", diff --git a/src/ai/backend/storage/errors/upload.py b/src/ai/backend/storage/errors/upload.py index 1e7155b67b4..7c1aab40d2b 100644 --- a/src/ai/backend/storage/errors/upload.py +++ b/src/ai/backend/storage/errors/upload.py @@ -32,6 +32,41 @@ def error_code(self) -> ErrorCode: ) +class TusSessionNotFoundError(BackendAIError, web.HTTPNotFound): + """ + Raised when a TUS handler is invoked for a session that is not registered + in Valkey (never created, or its state expired by TTL). + """ + + error_type = "https://api.backend.ai/probs/storage/no-such-upload-session" + error_title = "No such upload session" + + def error_code(self) -> ErrorCode: + return ErrorCode( + domain=ErrorDomain.STORAGE_PROXY, + operation=ErrorOperation.READ, + error_detail=ErrorDetail.NOT_FOUND, + ) + + +class UploadChunkExceedsTotalSizeError(BackendAIError, web.HTTPConflict): + """ + Raised when a PATCH chunk's offset+length would write past the declared + ``Upload-Length`` (409 Conflict). The Upload-Offset header itself is in + range; the chunk's body simply overruns the remaining slot. + """ + + error_type = "https://api.backend.ai/probs/storage/upload-chunk-exceeds-total-size" + error_title = "Upload Chunk Exceeds Total Size" + + def error_code(self) -> ErrorCode: + return ErrorCode( + domain=ErrorDomain.STORAGE_PROXY, + operation=ErrorOperation.UPDATE, + error_detail=ErrorDetail.CONFLICT, + ) + + class UploadSessionCorruptedError(BackendAIError, web.HTTPInternalServerError): """ Raised when the upload session metadata stored in Valkey cannot be parsed, diff --git a/src/ai/backend/storage/types.py b/src/ai/backend/storage/types.py index b04c114138e..69467a9fec9 100644 --- a/src/ai/backend/storage/types.py +++ b/src/ai/backend/storage/types.py @@ -10,6 +10,7 @@ import attrs import trafaret as t from aiohttp import BodyPartReader, MultipartReader, web +from aiohttp import StreamReader as AiohttpStreamReader from ai.backend.common import validators as tx from ai.backend.common.types import QuotaConfig, StreamReader, VFolderID @@ -148,3 +149,33 @@ async def read(self) -> AsyncIterator[bytes]: @override def content_type(self) -> str | None: return self._content_type + + +class TusChunkUploadStreamReader(StreamReader): + """ + Adapts a raw TUS PATCH request body (``application/offset+octet-stream``) + into the common :class:`StreamReader` by reading it in fixed-size chunks, + so the upload writer holds only one chunk in memory at a time. + """ + + def __init__( + self, + content: AiohttpStreamReader, + content_type: str | None, + chunk_size: int, + ) -> None: + self._content = content + self._content_type = content_type + self._chunk_size = chunk_size + + @override + async def read(self) -> AsyncIterator[bytes]: + while True: + chunk = await self._content.read(self._chunk_size) + if not chunk: + break + yield chunk + + @override + def content_type(self) -> str | None: + return self._content_type diff --git a/src/ai/backend/storage/volumes/vfs/__init__.py b/src/ai/backend/storage/volumes/vfs/__init__.py index 9113d817698..7b30e04a597 100644 --- a/src/ai/backend/storage/volumes/vfs/__init__.py +++ b/src/ai/backend/storage/volumes/vfs/__init__.py @@ -624,8 +624,8 @@ async def prepare_upload(self, vfid: VFolderID) -> str: def _create_target() -> None: upload_base_path = vfpath / ".upload" upload_base_path.mkdir(exist_ok=True) - upload_target_path = upload_base_path / session_id - upload_target_path.touch() + session_dir = upload_base_path / session_id + session_dir.mkdir(parents=True, exist_ok=True) loop = asyncio.get_running_loop() await loop.run_in_executor(None, _create_target) diff --git a/tests/unit/storage/api/test_tus_upload.py b/tests/unit/storage/api/test_tus_upload.py index bc95cd886f2..9565353e209 100644 --- a/tests/unit/storage/api/test_tus_upload.py +++ b/tests/unit/storage/api/test_tus_upload.py @@ -1,259 +1,443 @@ """ -Tests for TUS upload offset validation in tus_upload_part(). +Tests for the rewired TUS handlers (``tus_check_session``, ``tus_upload_part``). + +The handlers run against a real Valkey (the ``valkey_tus_client`` fixture, backed +by a redis container) and ``tmp_path`` chunk storage; only the volume/context +plumbing is mocked. This covers header-level guard checks (Upload-Offset parsing, +range bounds) plus the end-to-end happy path that writes through to +``TusUploadSession``. """ from __future__ import annotations -from collections.abc import Generator +import dataclasses +import secrets +from collections.abc import AsyncIterator from pathlib import Path from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from aiohttp import web - +from redis.asyncio import Redis + +from ai.backend.common import config +from ai.backend.common.clients.valkey_client.valkey_tus import ValkeyTusClient +from ai.backend.common.defs import REDIS_STREAM_LOCK, REDIS_TUS_DB +from ai.backend.common.lock import DistributedLockFactory +from ai.backend.common.typed_validators import HostPortPair as HostPortPairModel +from ai.backend.common.types import RedisConnectionInfo, TusSessionId, ValkeyTarget from ai.backend.storage.api.client import tus_upload_part -from ai.backend.storage.errors import InvalidAPIParameters, UploadOffsetMismatchError +from ai.backend.storage.errors import ( + InvalidAPIParameters, + UploadChunkExceedsTotalSizeError, + UploadOffsetMismatchError, +) +from ai.backend.storage.services.upload.lock import create_tus_lock_factory +from ai.backend.storage.services.upload.types import TusSessionState +from ai.backend.testutils.bootstrap import redis_container # noqa: F401 + + +@pytest.fixture +async def valkey_tus_client( + redis_container: tuple[str, HostPortPairModel], # noqa: F811 +) -> AsyncIterator[ValkeyTusClient]: + hostport_pair = redis_container[1] + client = await ValkeyTusClient.create( + ValkeyTarget(addr=hostport_pair.address), + db_id=REDIS_TUS_DB, + human_readable_name="test.tus.api", + ) + try: + yield client + finally: + await client.close() + + +@pytest.fixture +async def tus_lock_factory( + redis_container: tuple[str, HostPortPairModel], # noqa: F811 +) -> AsyncIterator[DistributedLockFactory]: + hostport_pair = redis_container[1] + lock_redis = RedisConnectionInfo( + Redis.from_url(f"redis://{hostport_pair.address}/{REDIS_STREAM_LOCK}"), + sentinel=None, + name="test.tus.api.lock", + service_name=None, + redis_helper_config=config.redis_helper_default_config, + ) + try: + yield create_tus_lock_factory(lock_redis) + finally: + await lock_redis.close() + + +@dataclasses.dataclass(slots=True) +class _PatchEnv: + vfpath: Path + session_dir: Path + session_id: str + + +def _build_request( + *, + vfpath: Path, + session_id: str, + total_size: int, + body: bytes | None, + offset_header: str | None, + valkey_client: ValkeyTusClient, + lock_factory: DistributedLockFactory, +) -> MagicMock: + volume = MagicMock() + volume.mangle_vfpath.return_value = vfpath + + ctx = MagicMock() + ctx.local_config.storage_proxy.secret = "test-secret" + ctx.get_volume.return_value.__aenter__ = AsyncMock(return_value=volume) + ctx.get_volume.return_value.__aexit__ = AsyncMock(return_value=None) + ctx.valkey_tus_client = valkey_client + ctx.tus_lock_factory = lock_factory + + request = MagicMock(spec=web.Request) + request.app = {"ctx": ctx} + request.headers = {} + if offset_header is not None: + request.headers["Upload-Offset"] = offset_header + request.query = {"token": "test-token"} + + if body is None: + content = AsyncMock() + content.read = AsyncMock(return_value=b"") + request.content = content + else: + body_state = {"pos": 0} + async def _read(_n: int) -> bytes: + if body_state["pos"] >= len(body): + return b"" + chunk = body[body_state["pos"] :] + body_state["pos"] = len(body) + return chunk -class TestTusUploadPartOffsetValidation: - """Tests for Upload-Offset header validation in tus_upload_part().""" + content = AsyncMock() + content.read = AsyncMock(side_effect=_read) + request.content = content - # ========================================================================= - # Low-level fixtures (internal building blocks) - # ========================================================================= + return request - @pytest.fixture - def token_data(self) -> dict[str, Any]: - """Create mock token data.""" - return { - "volume": "test-volume", - "vfid": MagicMock(), - "session": "test-session", - "size": "10240", - "relpath": "test-file.txt", - } - def _create_mock_request( - self, - tmp_path: Path, - client_offset: str | None, - ) -> MagicMock: - """Create a fully configured mock request.""" - # Mock volume - volume = MagicMock() - volume.mangle_vfpath.return_value = tmp_path - - # Mock context - ctx = MagicMock() - ctx.local_config.storage_proxy.secret = "test-secret" - ctx.get_volume.return_value.__aenter__ = AsyncMock(return_value=volume) - ctx.get_volume.return_value.__aexit__ = AsyncMock(return_value=None) - - # Mock request - request = MagicMock(spec=web.Request) - request.app = {"ctx": ctx} - request.headers = {} - request.query = {"token": "test-token"} - - if client_offset is not None: - request.headers["Upload-Offset"] = client_offset - - # Mock content reader (returns EOF immediately) - content = AsyncMock() - content.at_eof.return_value = True - request.content = content +@pytest.fixture +def tus_session_id() -> str: + # Unique per test so each test is isolated within the shared Valkey. + return f"test-session-{secrets.token_hex(8)}" - return request - - # ========================================================================= - # Scenario fixtures (composed, one per test scenario) - # ========================================================================= - - @pytest.fixture - def request_without_offset_header( - self, tmp_path: Path, token_data: dict[str, Any] - ) -> Generator[MagicMock, None, None]: - """Scenario: Missing Upload-Offset header.""" - request = self._create_mock_request(tmp_path, client_offset=None) - - with ( - patch("ai.backend.storage.api.client.check_params") as mock_check_params, - patch("ai.backend.storage.api.client.prepare_tus_session_headers") as mock_headers, - ): - mock_check_params.return_value.__aenter__ = AsyncMock( - return_value={"token": token_data, "dst_dir": None} - ) - mock_check_params.return_value.__aexit__ = AsyncMock(return_value=None) - mock_headers.return_value = {"Upload-Offset": "1024"} - - yield request - - @pytest.fixture - def request_with_invalid_offset( - self, tmp_path: Path, token_data: dict[str, Any] - ) -> Generator[MagicMock, None, None]: - """Scenario: Invalid (non-integer) Upload-Offset header.""" - request = self._create_mock_request(tmp_path, client_offset="not-a-number") - - with ( - patch("ai.backend.storage.api.client.check_params") as mock_check_params, - patch("ai.backend.storage.api.client.prepare_tus_session_headers") as mock_headers, - ): - mock_check_params.return_value.__aenter__ = AsyncMock( - return_value={"token": token_data, "dst_dir": None} - ) - mock_check_params.return_value.__aexit__ = AsyncMock(return_value=None) - mock_headers.return_value = {"Upload-Offset": "1024"} - - yield request - - @pytest.fixture - def request_offset_mismatch_client_behind( - self, tmp_path: Path, token_data: dict[str, Any] - ) -> Generator[MagicMock, None, None]: - """Scenario: Client offset (512) behind server offset (1024).""" - request = self._create_mock_request(tmp_path, client_offset="512") - - with ( - patch("ai.backend.storage.api.client.check_params") as mock_check_params, - patch("ai.backend.storage.api.client.prepare_tus_session_headers") as mock_headers, - ): - mock_check_params.return_value.__aenter__ = AsyncMock( - return_value={"token": token_data, "dst_dir": None} - ) - mock_check_params.return_value.__aexit__ = AsyncMock(return_value=None) - mock_headers.return_value = {"Upload-Offset": "1024"} - - yield request - - @pytest.fixture - def request_offset_mismatch_client_ahead( - self, tmp_path: Path, token_data: dict[str, Any] - ) -> Generator[MagicMock, None, None]: - """Scenario: Client offset (2048) ahead of server offset (1024).""" - request = self._create_mock_request(tmp_path, client_offset="2048") - - with ( - patch("ai.backend.storage.api.client.check_params") as mock_check_params, - patch("ai.backend.storage.api.client.prepare_tus_session_headers") as mock_headers, - ): - mock_check_params.return_value.__aenter__ = AsyncMock( - return_value={"token": token_data, "dst_dir": None} - ) - mock_check_params.return_value.__aexit__ = AsyncMock(return_value=None) - mock_headers.return_value = {"Upload-Offset": "1024"} - - yield request - - @pytest.fixture - def request_with_matching_offset( - self, tmp_path: Path, token_data: dict[str, Any] - ) -> Generator[MagicMock, None, None]: - """Scenario: Client offset matches server offset (both 1024).""" - request = self._create_mock_request(tmp_path, client_offset="1024") - - # Create upload temp file - upload_parent = tmp_path / ".upload" - upload_parent.mkdir(parents=True, exist_ok=True) - temp_file = upload_parent / token_data["session"] - temp_file.write_bytes(b"x" * 1024) - - with ( - patch("ai.backend.storage.api.client.check_params") as mock_check_params, - patch("ai.backend.storage.api.client.prepare_tus_session_headers") as mock_headers, - patch("ai.backend.storage.api.client.AsyncFileWriter") as mock_writer, - ): - mock_check_params.return_value.__aenter__ = AsyncMock( - return_value={"token": token_data, "dst_dir": None} - ) - mock_check_params.return_value.__aexit__ = AsyncMock(return_value=None) - mock_headers.return_value = {"Upload-Offset": "1024"} - - writer = AsyncMock() - mock_writer.return_value.__aenter__ = AsyncMock(return_value=writer) - mock_writer.return_value.__aexit__ = AsyncMock(return_value=None) - - yield request - - @pytest.fixture - def request_with_zero_offset( - self, tmp_path: Path, token_data: dict[str, Any] - ) -> Generator[MagicMock, None, None]: - """Scenario: New file upload with zero offset.""" - request = self._create_mock_request(tmp_path, client_offset="0") - - # Create empty upload temp file - upload_parent = tmp_path / ".upload" - upload_parent.mkdir(parents=True, exist_ok=True) - temp_file = upload_parent / token_data["session"] - temp_file.write_bytes(b"") - - with ( - patch("ai.backend.storage.api.client.check_params") as mock_check_params, - patch("ai.backend.storage.api.client.prepare_tus_session_headers") as mock_headers, - patch("ai.backend.storage.api.client.AsyncFileWriter") as mock_writer, - ): - mock_check_params.return_value.__aenter__ = AsyncMock( - return_value={"token": token_data, "dst_dir": None} - ) - mock_check_params.return_value.__aexit__ = AsyncMock(return_value=None) - mock_headers.return_value = {"Upload-Offset": "0"} - writer = AsyncMock() - mock_writer.return_value.__aenter__ = AsyncMock(return_value=writer) - mock_writer.return_value.__aexit__ = AsyncMock(return_value=None) +@pytest.fixture +def patch_env(tmp_path: Path, tus_session_id: str) -> _PatchEnv: + vfpath = tmp_path / "vfpath" + session_dir = vfpath / ".upload" / tus_session_id + session_dir.mkdir(parents=True) + return _PatchEnv(vfpath=vfpath, session_dir=session_dir, session_id=tus_session_id) + + +def _token_data(*, session_id: str, total_size: int, relpath: str) -> dict[str, Any]: + return { + "volume": "test-volume", + "vfid": MagicMock(), + "session": session_id, + "size": total_size, + "relpath": relpath, + } + - yield request +def _patch_handler_params(token_data: dict[str, Any]) -> Any: + cp = patch("ai.backend.storage.api.client.check_params") + mock_check_params = cp.start() + mock_check_params.return_value.__aenter__ = AsyncMock( + return_value={"token": token_data, "dst_dir": None} + ) + mock_check_params.return_value.__aexit__ = AsyncMock(return_value=None) + return cp - # ========================================================================= - # Test methods (minimal fixture injection, Act & Assert only) - # ========================================================================= - async def test_missing_upload_offset_header_raises_bad_request( +async def _register_session_state( + valkey: ValkeyTusClient, session_id: str, total_size: int +) -> None: + """Mimic ``create_upload_session`` having pre-registered the session in Valkey.""" + sid = TusSessionId(session_id) + await valkey.set_session_state(sid, TusSessionState.empty(sid, total_size).model_dump_json()) + + +class TestUploadOffsetHeaderValidation: + async def test_missing_offset_header_raises( self, - request_without_offset_header: MagicMock, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, ) -> None: - """When Upload-Offset header is missing, should raise InvalidAPIParameters (400).""" - with pytest.raises(InvalidAPIParameters): - await tus_upload_part(request_without_offset_header) - - async def test_invalid_upload_offset_header_raises_bad_request( + request = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=1024, + body=None, + offset_header=None, + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin") + cp = _patch_handler_params(token_data) + try: + with pytest.raises(InvalidAPIParameters): + await tus_upload_part(request) + finally: + cp.stop() + + async def test_non_integer_offset_header_raises( self, - request_with_invalid_offset: MagicMock, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, ) -> None: - """When Upload-Offset header is not a valid integer, should raise InvalidAPIParameters (400).""" - with pytest.raises(InvalidAPIParameters): - await tus_upload_part(request_with_invalid_offset) - - async def test_offset_mismatch_client_behind_raises_conflict( + request = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=1024, + body=None, + offset_header="not-a-number", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin") + cp = _patch_handler_params(token_data) + try: + with pytest.raises(InvalidAPIParameters): + await tus_upload_part(request) + finally: + cp.stop() + + async def test_negative_offset_raises_conflict( self, - request_offset_mismatch_client_behind: MagicMock, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, ) -> None: - """When client offset is behind server file size, should raise UploadOffsetMismatchError (409).""" - with pytest.raises(UploadOffsetMismatchError): - await tus_upload_part(request_offset_mismatch_client_behind) - - async def test_offset_mismatch_client_ahead_raises_conflict( + request = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=1024, + body=None, + offset_header="-1", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin") + cp = _patch_handler_params(token_data) + try: + with pytest.raises(UploadOffsetMismatchError): + await tus_upload_part(request) + finally: + cp.stop() + + async def test_offset_above_total_size_raises_conflict( + self, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, + ) -> None: + request = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=1024, + body=None, + offset_header="2048", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + token_data = _token_data(session_id=patch_env.session_id, total_size=1024, relpath="f.bin") + cp = _patch_handler_params(token_data) + try: + with pytest.raises(UploadOffsetMismatchError): + await tus_upload_part(request) + finally: + cp.stop() + + +class TestSessionNotFound: + async def test_missing_valkey_state_raises_not_found( + self, + tmp_path: Path, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, + ) -> None: + # The session has never been registered in Valkey (no + # `create_upload_session` call). Existence is determined by the Valkey + # state, so the handler must 404 regardless of any filesystem layout. + vfpath = tmp_path / "vfpath" + vfpath.mkdir() + + request = _build_request( + vfpath=vfpath, + session_id=f"missing-{secrets.token_hex(8)}", + total_size=1024, + body=b"", + offset_header="0", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + token_data = _token_data(session_id="missing-session", total_size=1024, relpath="f.bin") + cp = _patch_handler_params(token_data) + try: + with pytest.raises(web.HTTPNotFound): + await tus_upload_part(request) + finally: + cp.stop() + + +class TestHappyPath: + async def test_single_chunk_upload_completes_and_assembles( + self, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, + ) -> None: + payload = b"hello world" * 100 + await _register_session_state(valkey_tus_client, patch_env.session_id, len(payload)) + request = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=len(payload), + body=payload, + offset_header="0", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + token_data = _token_data( + session_id=patch_env.session_id, + total_size=len(payload), + relpath="result.bin", + ) + cp = _patch_handler_params(token_data) + try: + response = await tus_upload_part(request) + finally: + cp.stop() + + assert response.headers["Upload-Offset"] == str(len(payload)) + final_path = patch_env.vfpath / "result.bin" + assert final_path.read_bytes() == payload + # After assembly the chunk payloads are reclaimed; the completed marker + # is kept in Valkey (so a late duplicate PATCH observes completion). + assert list((patch_env.session_dir / "chunks").glob("*.dat")) == [] + + async def test_two_chunks_assemble_in_order( self, - request_offset_mismatch_client_ahead: MagicMock, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, ) -> None: - """When client offset is ahead of server file size, should raise UploadOffsetMismatchError (409).""" - with pytest.raises(UploadOffsetMismatchError): - await tus_upload_part(request_offset_mismatch_client_ahead) + await _register_session_state(valkey_tus_client, patch_env.session_id, 2048) + token_data = _token_data( + session_id=patch_env.session_id, total_size=2048, relpath="result.bin" + ) + + cp = _patch_handler_params(token_data) + try: + first = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=2048, + body=b"A" * 1024, + offset_header="0", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + await tus_upload_part(first) + + second = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=2048, + body=b"B" * 1024, + offset_header="1024", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + response = await tus_upload_part(second) + finally: + cp.stop() - async def test_matching_offset_proceeds_with_upload( + assert response.headers["Upload-Offset"] == "2048" + final_path = patch_env.vfpath / "result.bin" + assert final_path.read_bytes() == b"A" * 1024 + b"B" * 1024 + + async def test_duplicate_chunk_replay_is_idempotent( self, - request_with_matching_offset: MagicMock, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, ) -> None: - """When client offset matches server file size, upload should proceed successfully.""" - # No exception should be raised - await tus_upload_part(request_with_matching_offset) + await _register_session_state(valkey_tus_client, patch_env.session_id, 2048) + token_data = _token_data( + session_id=patch_env.session_id, total_size=2048, relpath="result.bin" + ) + cp = _patch_handler_params(token_data) + try: + payload = b"A" * 1024 + first = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=2048, + body=payload, + offset_header="0", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + await tus_upload_part(first) + + # Replay the same chunk; must not change committed offset. + replay = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=2048, + body=payload, + offset_header="0", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + response = await tus_upload_part(replay) + finally: + cp.stop() + + assert response.headers["Upload-Offset"] == "1024" - async def test_new_file_upload_with_zero_offset( + async def test_chunk_exceeding_declared_size_raises( self, - request_with_zero_offset: MagicMock, + patch_env: _PatchEnv, + valkey_tus_client: ValkeyTusClient, + tus_lock_factory: DistributedLockFactory, ) -> None: - """When both client offset and file size are 0, upload should proceed successfully.""" - # No exception should be raised - await tus_upload_part(request_with_zero_offset) + await _register_session_state(valkey_tus_client, patch_env.session_id, 10) + token_data = _token_data( + session_id=patch_env.session_id, total_size=10, relpath="result.bin" + ) + request = _build_request( + vfpath=patch_env.vfpath, + session_id=patch_env.session_id, + total_size=10, + body=b"too-much-data", # 13 bytes > 10 + offset_header="0", + valkey_client=valkey_tus_client, + lock_factory=tus_lock_factory, + ) + cp = _patch_handler_params(token_data) + try: + with pytest.raises(UploadChunkExceedsTotalSizeError): + await tus_upload_part(request) + finally: + cp.stop() + + # The aborted temp chunk file must be cleaned up. + chunks_dir = patch_env.session_dir / "chunks" + if chunks_dir.exists(): + assert list(chunks_dir.glob("*.tmp")) == []