diff --git a/s3proxy/Dockerfile b/s3proxy/Dockerfile index 891ad48..aa9038c 100644 --- a/s3proxy/Dockerfile +++ b/s3proxy/Dockerfile @@ -36,4 +36,5 @@ ENV S3PROXY_IP=0.0.0.0 \ EXPOSE 4433 -CMD ["uvicorn", "s3proxy.main:app", "--host", "0.0.0.0", "--port", "4433", "--loop", "uvloop"] +# Use the CLI entrypoint so TLS settings from Settings apply. +CMD ["python", "-m", "s3proxy.main"] diff --git a/s3proxy/client/types.py b/s3proxy/client/types.py index 0ab70fa..96fafea 100644 --- a/s3proxy/client/types.py +++ b/s3proxy/client/types.py @@ -13,6 +13,33 @@ class S3Credentials: service: str = "s3" +@dataclass(slots=True) +class ParsedHeaderAuth: + """Parsed SigV4 Authorization header, before signature verification.""" + + credentials: S3Credentials | None + signed_headers: list[str] + signature: str + date_stamp: str + region: str + service: str + amz_date: str + error: str | None = None + + @classmethod + def fail(cls, error: str, *, credentials: S3Credentials | None = None) -> ParsedHeaderAuth: + return cls( + credentials=credentials, + signed_headers=[], + signature="", + date_stamp="", + region="", + service="", + amz_date="", + error=error, + ) + + @dataclass(slots=True) class ParsedRequest: """Parsed S3 request information.""" diff --git a/s3proxy/client/verifier.py b/s3proxy/client/verifier.py index 397b926..acc4e34 100644 --- a/s3proxy/client/verifier.py +++ b/s3proxy/client/verifier.py @@ -13,7 +13,7 @@ import structlog from structlog.stdlib import BoundLogger -from .types import ParsedRequest, S3Credentials +from .types import ParsedHeaderAuth, ParsedRequest, S3Credentials logger: BoundLogger = structlog.get_logger(__name__) @@ -95,10 +95,8 @@ def verify(self, request: ParsedRequest, path: str) -> tuple[bool, S3Credentials return False, None, "No AWS signature found" - def _verify_header_signature( - self, request: ParsedRequest, path: str, auth_header: str - ) -> tuple[bool, S3Credentials | None, str]: - """Verify Authorization header signature.""" + def _parse_header_auth(self, request: ParsedRequest, auth_header: str) -> ParsedHeaderAuth: + """Parse Authorization header fields without verifying the signature.""" try: parts = auth_header.replace("AWS4-HMAC-SHA256 ", "").split(",") auth_parts = {} @@ -112,41 +110,126 @@ def _verify_header_signature( credentials, date_stamp, region, service, error = self._parse_v4_credential(credential) if error: - return False, None, error + return ParsedHeaderAuth.fail(error) amz_date = request.headers.get("x-amz-date", "") if not amz_date: - return False, credentials, "Missing x-amz-date header" - - try: - request_time = datetime.strptime(amz_date, "%Y%m%dT%H%M%SZ").replace(tzinfo=UTC) - if abs(datetime.now(UTC) - request_time) > CLOCK_SKEW_TOLERANCE: - return False, credentials, "Request time too skewed" - except ValueError: - return False, credentials, "Invalid x-amz-date format" + return ParsedHeaderAuth.fail( + "Missing x-amz-date header", + credentials=credentials, + ) - canonical_request = self._build_canonical_request( - request, path, signed_headers.split(";") - ) - calculated_sig = self._compute_v4_signature( - canonical_request, amz_date, date_stamp, region, service, credentials.secret_key + return ParsedHeaderAuth( + credentials=credentials, + signed_headers=signed_headers.split(";"), + signature=signature, + date_stamp=date_stamp, + region=region, + service=service, + amz_date=amz_date, ) + except (KeyError, ValueError, IndexError) as e: + return ParsedHeaderAuth.fail(f"Invalid Authorization header: {e}") - if hmac.compare_digest(calculated_sig, signature): - return True, credentials, "" + def _check_request_time(self, amz_date: str) -> str | None: + try: + request_time = datetime.strptime(amz_date, "%Y%m%dT%H%M%SZ").replace(tzinfo=UTC) + if abs(datetime.now(UTC) - request_time) > CLOCK_SKEW_TOLERANCE: + return "Request time too skewed" + except ValueError: + return "Invalid x-amz-date format" + return None + + def prepare_header_auth( + self, request: ParsedRequest, auth_header: str + ) -> tuple[S3Credentials | None, str]: + """Parse credentials and clock skew. Signature verified separately.""" + parsed = self._parse_header_auth(request, auth_header) + if parsed.error: + return None, parsed.error + assert parsed.credentials is not None + time_error = self._check_request_time(parsed.amz_date) + if time_error: + return parsed.credentials, time_error + return parsed.credentials, "" + + def verify_with_payload_hash( + self, + request: ParsedRequest, + path: str, + auth_header: str, + payload_hash: str, + ) -> tuple[bool, S3Credentials | None, str]: + """Verify SigV4 using an explicit payload hash (stream-computed).""" + auth = self._parse_header_auth(request, auth_header) + if auth.error: + return False, auth.credentials, auth.error + assert auth.credentials is not None + + time_error = self._check_request_time(auth.amz_date) + if time_error: + return False, auth.credentials, time_error + + canonical_request = self._build_canonical_request( + request, path, auth.signed_headers, payload_hash=payload_hash + ) + calculated_sig = self._compute_v4_signature( + canonical_request, + auth.amz_date, + auth.date_stamp, + auth.region, + auth.service, + auth.credentials.secret_key, + ) - logger.debug( - "Signature verification failed", - method=request.method, - path=path, - signed_headers=signed_headers, - expected_sig=signature[:16] + "...", - calculated_sig=calculated_sig[:16] + "...", - ) - return False, credentials, "Signature mismatch" + if hmac.compare_digest(calculated_sig, auth.signature): + return True, auth.credentials, "" - except (KeyError, ValueError, IndexError) as e: - return False, None, f"Invalid Authorization header: {e}" + logger.debug( + "Signature verification failed", + method=request.method, + path=path, + signed_headers=";".join(auth.signed_headers), + expected_sig=auth.signature[:16] + "...", + calculated_sig=calculated_sig[:16] + "...", + ) + return False, auth.credentials, "Signature mismatch" + + def _verify_header_signature( + self, request: ParsedRequest, path: str, auth_header: str + ) -> tuple[bool, S3Credentials | None, str]: + """Verify Authorization header signature.""" + auth = self._parse_header_auth(request, auth_header) + if auth.error: + return False, auth.credentials, auth.error + assert auth.credentials is not None + + time_error = self._check_request_time(auth.amz_date) + if time_error: + return False, auth.credentials, time_error + + canonical_request = self._build_canonical_request(request, path, auth.signed_headers) + calculated_sig = self._compute_v4_signature( + canonical_request, + auth.amz_date, + auth.date_stamp, + auth.region, + auth.service, + auth.credentials.secret_key, + ) + + if hmac.compare_digest(calculated_sig, auth.signature): + return True, auth.credentials, "" + + logger.debug( + "Signature verification failed", + method=request.method, + path=path, + signed_headers=";".join(auth.signed_headers), + expected_sig=auth.signature[:16] + "...", + calculated_sig=calculated_sig[:16] + "...", + ) + return False, auth.credentials, "Signature mismatch" def _verify_presigned_v4( self, request: ParsedRequest, path: str @@ -268,7 +351,12 @@ def _verify_presigned_v2( return False, None, f"Invalid V2 presigned URL: {e}" def _build_canonical_request( - self, request: ParsedRequest, path: str, signed_headers: list[str] + self, + request: ParsedRequest, + path: str, + signed_headers: list[str], + *, + payload_hash: str | None = None, ) -> str: """Build canonical request for signature verification.""" method = request.method.upper() @@ -285,9 +373,10 @@ def _build_canonical_request( signed_headers_str = ";".join(sorted(signed_headers)) - payload_hash = request.headers.get( - "x-amz-content-sha256", hashlib.sha256(request.body).hexdigest() - ) + if payload_hash is None: + payload_hash = request.headers.get( + "x-amz-content-sha256", hashlib.sha256(request.body).hexdigest() + ) return "\n".join( [ diff --git a/s3proxy/disconnect.py b/s3proxy/disconnect.py new file mode 100644 index 0000000..885bc87 --- /dev/null +++ b/s3proxy/disconnect.py @@ -0,0 +1,13 @@ +"""Client disconnect errors for streaming uploads.""" + +from __future__ import annotations + +from .errors import S3Error + + +class ClientDisconnectError(S3Error): + """Client closed the connection mid-upload.""" + + @classmethod + def raised(cls) -> ClientDisconnectError: + return cls(400, "BadRequest", "Client disconnected") diff --git a/s3proxy/handlers/multipart/upload_part.py b/s3proxy/handlers/multipart/upload_part.py index 5beae20..1f21977 100644 --- a/s3proxy/handlers/multipart/upload_part.py +++ b/s3proxy/handlers/multipart/upload_part.py @@ -17,6 +17,7 @@ from ... import crypto from ...client import S3Client, S3Credentials from ...errors import S3Error, raise_for_client_error, raise_for_exception +from ...signature import deferred_signature_required, verify_deferred_payload_hash from ...state import ( InternalPartMetadata, MultipartUploadState, @@ -188,7 +189,11 @@ async def handle_upload_part(self, request: Request, creds: S3Credentials) -> Re ) # Late signature verification for large signed uploads - if is_large_signed and content_sha and result["computed_sha256"] != content_sha: + if deferred_signature_required(request): + verify_deferred_payload_hash( + request, request.app.state.verifier, result["computed_sha256"] + ) + elif is_large_signed and content_sha and result["computed_sha256"] != content_sha: logger.warning( "UPLOAD_PART_SHA256_MISMATCH", bucket=bucket, diff --git a/s3proxy/handlers/objects/put.py b/s3proxy/handlers/objects/put.py index 213484c..ed22513 100644 --- a/s3proxy/handlers/objects/put.py +++ b/s3proxy/handlers/objects/put.py @@ -2,15 +2,19 @@ import base64 import hashlib +from collections.abc import AsyncIterator from typing import Any import structlog from fastapi import Request, Response +from starlette.requests import ClientDisconnect from structlog.stdlib import BoundLogger from ... import crypto from ...client import S3Client, S3Credentials +from ...disconnect import ClientDisconnectError from ...errors import S3Error +from ...signature import verify_deferred_payload_hash from ...state import ( MultipartMetadata, PartMetadata, @@ -22,6 +26,27 @@ logger: BoundLogger = structlog.get_logger(__name__) +_STREAM_CHUNK = 64 * 1024 + + +async def _iter_request_body(request: Request, decode_chunked: bool) -> AsyncIterator[bytes]: + """Single-pass body iterator: preloaded small body, else live stream. + + Avoids reading request.stream() after request.body() pinned the payload in + Starlette's cache (double buffer on large uploads). + """ + if decode_chunked: + async for chunk in decode_aws_chunked_stream(request): + yield chunk + return + preloaded = getattr(request.state, "s3proxy_preloaded_body", None) + if isinstance(preloaded, (bytes, bytearray)) and len(preloaded) > 0: + for offset in range(0, len(preloaded), _STREAM_CHUNK): + yield preloaded[offset : offset + _STREAM_CHUNK] + return + async for chunk in request.stream(): + yield chunk + class PutObjectMixin(BaseHandler): async def handle_put_object(self, request: Request, creds: S3Credentials) -> Response: @@ -126,7 +151,11 @@ async def _put_buffered( content_length_mb=round(content_length / 1024 / 1024, 2), ) - body = await request.body() + preloaded = getattr(request.state, "s3proxy_preloaded_body", None) + if isinstance(preloaded, (bytes, bytearray)) and len(preloaded) > 0: + body = bytes(preloaded) + else: + body = await request.body() if needs_chunked_decode: body = decode_aws_chunked(body) @@ -191,7 +220,8 @@ async def _put_streaming( total_plaintext_size = 0 part_num = 0 md5_hash = hashlib.md5(usedforsecurity=False) - sha256_hash = hashlib.sha256() if expected_sha256 else None + deferred_sig = getattr(request.state, "s3proxy_deferred_sig", False) is True + sha256_hash = hashlib.sha256() if (expected_sha256 or deferred_sig) else None buffer = bytearray() async def upload_part(data: bytes) -> None: @@ -230,10 +260,7 @@ async def upload_part(data: bytes) -> None: verify_sha256=expected_sha256 is not None, ) - if decode_chunked: - stream_source = decode_aws_chunked_stream(request) - else: - stream_source = request.stream() + stream_source = _iter_request_body(request, decode_chunked) async for chunk in stream_source: buffer.extend(chunk) @@ -254,8 +281,16 @@ async def upload_part(data: bytes) -> None: buffer.clear() await upload_part(part_data) - # Verify SHA256 if provided - if expected_sha256 is not None: + # Verify SHA256 if provided, or deferred SigV4 after streaming hash + if deferred_sig: + try: + verify_deferred_payload_hash( + request, request.app.state.verifier, sha256_hash.hexdigest() + ) + except S3Error: + await client.abort_multipart_upload(bucket, key, upload_id) + raise + elif expected_sha256 is not None: computed_sha256 = sha256_hash.hexdigest() if computed_sha256 != expected_sha256: logger.error( @@ -298,6 +333,9 @@ async def upload_part(data: bytes) -> None: ) return Response(headers={"ETag": f'"{etag}"'}) + except ClientDisconnectError, ClientDisconnect: + await self._safe_abort(client, bucket, key, upload_id) + raise ClientDisconnectError.raised() from None except S3Error: raise except Exception as e: diff --git a/s3proxy/request_handler.py b/s3proxy/request_handler.py index 9691a7f..d2e7577 100644 --- a/s3proxy/request_handler.py +++ b/s3proxy/request_handler.py @@ -12,7 +12,7 @@ from fastapi.responses import PlainTextResponse, StreamingResponse from structlog.stdlib import BoundLogger -from . import concurrency +from . import concurrency, crypto from .client import ParsedRequest, SigV4Verifier from .dashboard import record_request from .errors import S3Error, raise_for_client_error, raise_for_exception @@ -68,6 +68,28 @@ def _needs_body_for_signature(headers: dict[str, str]) -> bool: return headers.get("x-amz-content-sha256", "") == "" +def _parse_content_length(headers: dict[str, str]) -> int: + try: + return int(headers.get("content-length", "0")) + except ValueError: + return 0 + + +def _defer_signature_for_body(headers: dict[str, str], content_length: int) -> bool: + """Large bodies without x-amz-content-sha256 are hashed while streaming.""" + return _needs_body_for_signature(headers) and content_length > crypto.MAX_BUFFER_SIZE + + +def _signature_path(request: Request) -> str: + raw_path = request.scope.get("raw_path") + if raw_path: + sig_path = raw_path.decode("utf-8", errors="replace") + if "?" in sig_path: + sig_path = sig_path.split("?", 1)[0] + return sig_path + return request.url.path + + async def handle_proxy_request( request: Request, handler: S3ProxyHandler, @@ -198,17 +220,24 @@ async def _handle_proxy_request_impl( headers = {k.lower(): v for k, v in request.headers.items()} query = parse_qs(str(request.url.query), keep_blank_values=True) + content_length = _parse_content_length(headers) + defer_sig = request.method in ("PUT", "POST") and _defer_signature_for_body( + headers, content_length + ) + needs_body = request.method in ("PUT", "POST") and _needs_body_for_signature(headers) - content_length = headers.get("content-length", "0") - body = await request.body() if needs_body else b"" - if needs_body and len(body) > 0: - logger.debug( - "body_loaded", - content_length=content_length, - body_size=len(body), - method=request.method, - path=request.url.path, - ) + body = b"" + if needs_body and not defer_sig: + body = await request.body() + if body: + request.state.s3proxy_preloaded_body = body + logger.debug( + "body_loaded", + content_length=content_length, + body_size=len(body), + method=request.method, + path=request.url.path, + ) parsed = ParsedRequest( method=request.method, @@ -219,18 +248,30 @@ async def _handle_proxy_request_impl( body=body, ) - raw_path = request.scope.get("raw_path") - if raw_path: - sig_path = raw_path.decode("utf-8", errors="replace") - if "?" in sig_path: - sig_path = sig_path.split("?", 1)[0] + sig_path = _signature_path(request) + if defer_sig: + auth_header = headers.get("authorization", "") + if not auth_header.startswith("AWS4-HMAC-SHA256"): + raise S3Error.access_denied("No AWS signature found") + verified_creds, error = verifier.prepare_header_auth(parsed, auth_header) + if not verified_creds: + raise S3Error.access_denied(error or "Access Denied") + if error: + raise S3Error.access_denied(error) + request.state.s3proxy_deferred_sig = True + request.state.s3proxy_sig_path = sig_path + logger.debug( + "signature_deferred", + content_length=content_length, + method=request.method, + path=request.url.path, + ) else: - sig_path = request.url.path - valid, verified_creds, error = verifier.verify(parsed, sig_path) - if not valid or not verified_creds: - if error and "signature" in error.lower(): - raise S3Error.signature_does_not_match(error) - raise S3Error.access_denied(error or "No credentials") + valid, verified_creds, error = verifier.verify(parsed, sig_path) + if not valid or not verified_creds: + if error and "signature" in error.lower(): + raise S3Error.signature_does_not_match(error) + raise S3Error.access_denied(error or "No credentials") dispatcher = RequestDispatcher(handler) try: diff --git a/s3proxy/signature.py b/s3proxy/signature.py new file mode 100644 index 0000000..ca22d0f --- /dev/null +++ b/s3proxy/signature.py @@ -0,0 +1,40 @@ +"""Deferred SigV4 verification for large streaming uploads.""" + +from __future__ import annotations + +from urllib.parse import parse_qs + +from fastapi import Request + +from .client import ParsedRequest, SigV4Verifier +from .errors import S3Error + + +def deferred_signature_required(request: Request) -> bool: + return getattr(request.state, "s3proxy_deferred_sig", False) is True + + +def verify_deferred_payload_hash( + request: Request, verifier: SigV4Verifier, payload_hash: str +) -> None: + """Verify a request whose body was streamed after computing its SHA256.""" + if not deferred_signature_required(request): + return + + headers = {k.lower(): v for k, v in request.headers.items()} + query = parse_qs(str(request.url.query), keep_blank_values=True) + parsed = ParsedRequest( + method=request.method, + bucket="", + key="", + query_params=query, + headers=headers, + body=b"", + ) + sig_path = getattr(request.state, "s3proxy_sig_path", request.url.path) + auth_header = headers.get("authorization", "") + valid, _, error = verifier.verify_with_payload_hash(parsed, sig_path, auth_header, payload_hash) + if not valid: + if error and "signature" in error.lower(): + raise S3Error.signature_does_not_match(error) + raise S3Error.access_denied(error or "Access Denied") diff --git a/tests/unit/test_deferred_signature.py b/tests/unit/test_deferred_signature.py new file mode 100644 index 0000000..cfed82f --- /dev/null +++ b/tests/unit/test_deferred_signature.py @@ -0,0 +1,63 @@ +"""Deferred SigV4 verification after streaming body hash.""" + +import hashlib +from datetime import UTC, datetime +from unittest.mock import MagicMock + +import pytest + +from s3proxy.client import ParsedRequest, SigV4Verifier +from s3proxy.signature import verify_deferred_payload_hash + + +@pytest.mark.asyncio +async def test_verify_deferred_payload_hash_accepts_valid_signature(): + verifier = SigV4Verifier({"testkey": "testsecret"}) + payload = b"scylla-sst-bytes" + payload_hash = hashlib.sha256(payload).hexdigest() + amz_date = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + date_stamp = amz_date[:8] + parsed = ParsedRequest( + method="PUT", + bucket="", + key="", + headers={ + "host": "s3proxy.test", + "x-amz-date": amz_date, + "content-length": str(len(payload)), + }, + body=payload, + ) + path = "/bucket/key" + canonical = verifier._build_canonical_request(parsed, path, ["host", "x-amz-date"]) + signature = verifier._compute_v4_signature( + canonical, + amz_date, + date_stamp, + "us-east-1", + "s3", + "testsecret", + ) + auth_header = ( + "AWS4-HMAC-SHA256 " + f"Credential=testkey/{date_stamp}/us-east-1/s3/aws4_request, " + "SignedHeaders=host;x-amz-date, " + f"Signature={signature}" + ) + + request = MagicMock() + request.method = "PUT" + request.url = MagicMock() + request.url.path = path + request.url.query = "" + request.headers = { + "authorization": auth_header, + "x-amz-date": amz_date, + "host": "s3proxy.test", + "content-length": str(len(payload)), + } + request.state = MagicMock() + request.state.s3proxy_deferred_sig = True + request.state.s3proxy_sig_path = path + + verify_deferred_payload_hash(request, verifier, payload_hash) diff --git a/tests/unit/test_disconnect_upload.py b/tests/unit/test_disconnect_upload.py new file mode 100644 index 0000000..538116a --- /dev/null +++ b/tests/unit/test_disconnect_upload.py @@ -0,0 +1,66 @@ +"""Client disconnect during streaming upload must abort and stop reading.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from starlette.requests import ClientDisconnect + +from s3proxy.disconnect import ClientDisconnectError +from s3proxy.handlers.objects.put import PutObjectMixin + +MB = 1024 * 1024 + + +class _Client: + def __init__(self): + self.credentials = MagicMock(access_key="testkey") + self.abort_multipart_upload = AsyncMock() + + async def create_multipart_upload(self, *a, **k): + return {"UploadId": "upload-1"} + + async def upload_part(self, bucket, key, upload_id, part_number, body): + return {"ETag": '"etag"'} + + +class _DisconnectRequest: + def __init__(self, total: int): + self.headers = {"content-type": "application/octet-stream"} + self.state = MagicMock(spec=[]) + self.app = MagicMock() + self._total = total + self._sent = 0 + + async def stream(self): + chunk = 64 * 1024 + while self._sent < self._total: + self._sent += chunk + if self._sent >= 12 * MB: + raise ClientDisconnect() + yield b"x" * min(chunk, self._total - (self._sent - chunk)) + + +def _handler() -> PutObjectMixin: + h = PutObjectMixin.__new__(PutObjectMixin) + h.keyring = MagicMock() + h.keyring.key_for.return_value = ("kid1", b"0" * 32) + return h + + +@pytest.mark.asyncio +async def test_put_streaming_aborts_multipart_on_client_disconnect(): + request = _DisconnectRequest(12 * MB) + client = _Client() + + with pytest.raises(ClientDisconnectError): + await _handler()._put_streaming( + request, + client, + "bucket", + "key", + "application/octet-stream", + decode_chunked=False, + expected_sha256=None, + ) + + client.abort_multipart_upload.assert_awaited_once_with("bucket", "key", "upload-1") diff --git a/tests/unit/test_put_body_iterator.py b/tests/unit/test_put_body_iterator.py new file mode 100644 index 0000000..dc4ea99 --- /dev/null +++ b/tests/unit/test_put_body_iterator.py @@ -0,0 +1,38 @@ +"""PUT handler must not double-buffer a body already loaded for signature.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from s3proxy.handlers.objects.put import _iter_request_body + +MB = 1024 * 1024 + + +@pytest.mark.asyncio +async def test_iter_request_body_uses_preloaded_chunks(): + request = MagicMock() + request.state = MagicMock() + request.state.s3proxy_preloaded_body = b"x" * (2 * MB) + request.stream = AsyncMock() + + chunks = [c async for c in _iter_request_body(request, decode_chunked=False)] + + assert sum(len(c) for c in chunks) == 2 * MB + request.stream.assert_not_called() + + +@pytest.mark.asyncio +async def test_iter_request_body_streams_when_not_preloaded(): + request = MagicMock() + request.state = MagicMock(spec=[]) # no s3proxy_preloaded_body + + async def _stream(): + yield b"ab" + yield b"cd" + + request.stream = MagicMock(return_value=_stream()) + + chunks = [c async for c in _iter_request_body(request, decode_chunked=False)] + + assert b"".join(chunks) == b"abcd" diff --git a/tests/unit/test_request_handler_body.py b/tests/unit/test_request_handler_body.py new file mode 100644 index 0000000..3853eb3 --- /dev/null +++ b/tests/unit/test_request_handler_body.py @@ -0,0 +1,148 @@ +"""Request gate: large unsigned-header bodies stream once; sig verified after hash.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import Request + +from s3proxy.errors import S3Error +from s3proxy.request_handler import ( + _defer_signature_for_body, + _handle_proxy_request_impl, + handle_proxy_request, +) + +MB = 1024 * 1024 + + +class TestDeferSignatureForBody: + def test_defers_large_put_without_payload_hash(self): + assert _defer_signature_for_body({}, 179 * MB) is True + + def test_no_defer_with_unsigned_payload(self): + assert ( + _defer_signature_for_body({"x-amz-content-sha256": "UNSIGNED-PAYLOAD"}, 179 * MB) + is False + ) + + def test_no_defer_with_explicit_hash(self): + assert _defer_signature_for_body({"x-amz-content-sha256": "abc"}, 179 * MB) is False + + def test_no_defer_small_put_without_header(self): + assert _defer_signature_for_body({}, 4 * MB) is False + + +def _make_request( + *, + method: str = "PUT", + content_length: int, + content_sha: str | None = None, +) -> MagicMock: + request = MagicMock(spec=Request) + request.method = method + request.url = MagicMock() + request.url.path = "/bucket/scylla-backup/sst/big-Data.db" + request.url.query = "" + headers: dict[str, str] = { + "content-length": str(content_length), + "authorization": "AWS4-HMAC-SHA256 Credential=x, SignedHeaders=host, Signature=sig", + "x-amz-date": "20260708T080000Z", + } + if content_sha is not None: + headers["x-amz-content-sha256"] = content_sha + request.headers = headers + request.scope = {"raw_path": request.url.path.encode()} + request.state = MagicMock() + request.body = AsyncMock(return_value=b"") + request.stream = MagicMock() + return request + + +@pytest.mark.asyncio +async def test_large_put_without_header_defers_signature(): + request = _make_request(content_length=179 * MB) + verifier = MagicMock() + verifier.prepare_header_auth = MagicMock(return_value=(MagicMock(), "")) + + with patch("s3proxy.request_handler.RequestDispatcher") as dispatcher_cls: + dispatcher_cls.return_value.dispatch = AsyncMock(return_value=None) + await _handle_proxy_request_impl(request, MagicMock(), verifier) + + request.body.assert_not_awaited() + assert request.state.s3proxy_deferred_sig is True + verifier.prepare_header_auth.assert_called_once() + verifier.verify.assert_not_called() + + +@pytest.mark.asyncio +async def test_large_put_unsigned_payload_verifies_immediately(): + request = _make_request(content_length=179 * MB, content_sha="UNSIGNED-PAYLOAD") + verifier = MagicMock() + verifier.verify = MagicMock(return_value=(True, MagicMock(), "")) + + with patch("s3proxy.request_handler.RequestDispatcher") as dispatcher_cls: + dispatcher_cls.return_value.dispatch = AsyncMock(return_value=None) + await _handle_proxy_request_impl(request, MagicMock(), verifier) + + request.body.assert_not_awaited() + verifier.verify.assert_called_once() + + +@pytest.mark.asyncio +async def test_small_put_without_header_loads_body_once(): + payload = b"x" * (4 * MB) + request = _make_request(content_length=len(payload)) + request.body = AsyncMock(return_value=payload) + verifier = MagicMock() + verifier.verify = MagicMock(return_value=(True, MagicMock(), "")) + + with patch("s3proxy.request_handler.RequestDispatcher") as dispatcher_cls: + dispatcher_cls.return_value.dispatch = AsyncMock(return_value=None) + await _handle_proxy_request_impl(request, MagicMock(), verifier) + + request.body.assert_awaited_once() + assert request.state.s3proxy_preloaded_body == payload + verifier.verify.assert_called_once() + + +@pytest.mark.asyncio +async def test_deferred_prepare_failure_raises_access_denied(): + request = _make_request(content_length=179 * MB) + verifier = MagicMock() + verifier.prepare_header_auth = MagicMock(return_value=(None, "Unknown access key")) + + with pytest.raises(S3Error) as exc: + await _handle_proxy_request_impl(request, MagicMock(), verifier) + assert exc.value.status_code == 403 + request.body.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_impl_error_releases_memory_reservation(): + import s3proxy.concurrency as concurrency_module + import s3proxy.request_handler as request_handler_module + + request = MagicMock(spec=Request) + request.method = "PUT" + request.url = MagicMock() + request.url.path = "/bucket/key" + request.url.query = "" + request.headers = {"content-length": str(179 * MB)} + request.scope = {"raw_path": b"/bucket/key"} + request.client = MagicMock() + request.client.host = "10.0.0.1" + request.body = AsyncMock() + + with ( + patch.object(concurrency_module, "try_acquire_memory", new_callable=AsyncMock) as acquire, + patch.object(concurrency_module, "release_memory", new_callable=AsyncMock) as release, + patch.object( + request_handler_module, "_handle_proxy_request_impl", new_callable=AsyncMock + ) as impl, + ): + acquire.return_value = 32 * MB + impl.side_effect = S3Error.access_denied("bad sig") + with pytest.raises(S3Error): + await handle_proxy_request(request, MagicMock(), MagicMock()) + + release.assert_awaited_once_with(32 * MB) diff --git a/tests/unit/test_scylla_deferred_put_memory.py b/tests/unit/test_scylla_deferred_put_memory.py new file mode 100644 index 0000000..77841bd --- /dev/null +++ b/tests/unit/test_scylla_deferred_put_memory.py @@ -0,0 +1,138 @@ +"""Regression: Scylla SST PUT without x-amz-content-sha256 must not double-buffer. + +Production OOM (2026-07-08): request_handler loaded the full ~179MB body for SigV4, +then put.py streamed request.stream() from Starlette's cache — ~350MB per request +while the governor saw ~88MB. This test drives the real PUT streaming handler with +deferred signature (no preloaded body) and asserts peak Python heap stays bounded. +""" + +from __future__ import annotations + +import hashlib +import tracemalloc +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from s3proxy import crypto +from s3proxy.errors import S3Error +from s3proxy.handlers.objects.put import PutObjectMixin + +MB = 1024 * 1024 +SCYLLA_PUT_BYTES = 179 * MB + + +class _Client: + def __init__(self): + self.credentials = MagicMock(access_key="testkey") + self.abort_multipart_upload = AsyncMock() + + async def create_multipart_upload(self, *a, **k): + return {"UploadId": "upload-1"} + + async def upload_part(self, bucket, key, upload_id, part_number, body): + sent = bytes(body) + return {"ETag": hashlib.md5(sent).hexdigest()} + + async def complete_multipart_upload(self, *a, **k): + return {} + + +class _StreamRequest: + """Simulates a large PUT after deferred signature (body never preloaded).""" + + def __init__(self, total: int, chunk: int = 64 * 1024): + self._total = total + self._chunk = chunk + self.headers = { + "content-type": "application/octet-stream", + "content-length": str(total), + } + self.url = MagicMock() + self.url.path = "/bucket/backup/sst/me-big-Data.db" + self.state = MagicMock(spec=[]) + self.state.s3proxy_deferred_sig = True + self.app = MagicMock() + + async def stream(self): + for i in range(0, self._total, self._chunk): + yield b"s" * min(self._chunk, self._total - i) + + +def _handler() -> PutObjectMixin: + h = PutObjectMixin.__new__(PutObjectMixin) + h.keyring = MagicMock() + h.keyring.key_for.return_value = ("kid1", b"0" * 32) + return h + + +@pytest.mark.asyncio +async def test_scylla_sized_deferred_put_peak_is_bounded_not_double_body(): + """Peak must stay near streaming_upload_peak, not 2x Content-Length.""" + request = _StreamRequest(SCYLLA_PUT_BYTES) + client = _Client() + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "s3proxy.handlers.objects.put.verify_deferred_payload_hash", + lambda *a, **k: None, + ) + mp.setattr( + "s3proxy.handlers.objects.put.save_multipart_metadata", + AsyncMock(), + ) + tracemalloc.start() + base = tracemalloc.get_traced_memory()[0] + await _handler()._put_streaming( + request, + client, + "bucket", + "key", + "application/octet-stream", + decode_chunked=False, + expected_sha256=None, + ) + peak = tracemalloc.get_traced_memory()[1] + tracemalloc.stop() + + measured_mb = (peak - base) / MB + honest_peak_mb = crypto.streaming_upload_peak(SCYLLA_PUT_BYTES) / MB + double_body_mb = 2 * SCYLLA_PUT_BYTES / MB + + # Old bug held ~2x body (~358MB) for this payload; streaming peak is ~56MB. + assert measured_mb < honest_peak_mb * 1.5, ( + f"peak {measured_mb:.1f}MB exceeds 1.5x streaming_upload_peak " + f"({honest_peak_mb:.1f}MB) — likely double-buffer regression" + ) + assert measured_mb < double_body_mb / 4, ( + f"peak {measured_mb:.1f}MB is too close to old 2x-body failure ({double_body_mb:.0f}MB)" + ) + + +@pytest.mark.asyncio +async def test_deferred_sig_failure_aborts_multipart(): + """Invalid signature after streaming must not complete the upload.""" + request = _StreamRequest(4 * MB) + request.app.state.verifier = MagicMock() + client = _Client() + + def _bad_verify(*a, **k): + raise S3Error.signature_does_not_match("Signature mismatch") + + with pytest.MonkeyPatch.context() as mp: + mp.setattr( + "s3proxy.handlers.objects.put.verify_deferred_payload_hash", + _bad_verify, + ) + with pytest.raises(S3Error): + await _handler()._put_streaming( + request, + client, + "bucket", + "key", + "application/octet-stream", + decode_chunked=False, + expected_sha256=None, + ) + + client.abort_multipart_upload.assert_awaited_once()