diff --git a/packages/google-cloud-storage/google/cloud/storage/_media/_upload.py b/packages/google-cloud-storage/google/cloud/storage/_media/_upload.py index acbc04e715c9..819e496b868f 100644 --- a/packages/google-cloud-storage/google/cloud/storage/_media/_upload.py +++ b/packages/google-cloud-storage/google/cloud/storage/_media/_upload.py @@ -676,6 +676,11 @@ def _prepare_request(self): msg = _STREAM_ERROR_TEMPLATE.format(start_byte, self.bytes_uploaded) raise ValueError(msg) + if self._total_bytes is None and not content_range.endswith("/*"): + total_str = content_range.split("/")[-1] + if total_str.isdigit(): + self._total_bytes = int(total_str) + self._update_checksum(start_byte, payload) headers = { diff --git a/packages/google-cloud-storage/tests/resumable_media/unit/requests/test_upload.py b/packages/google-cloud-storage/tests/resumable_media/unit/requests/test_upload.py index e99855570694..fcdbe38c1b4c 100644 --- a/packages/google-cloud-storage/tests/resumable_media/unit/requests/test_upload.py +++ b/packages/google-cloud-storage/tests/resumable_media/unit/requests/test_upload.py @@ -18,9 +18,8 @@ import tempfile from unittest import mock -import pytest # type: ignore - import google.cloud.storage._media.requests.upload as upload_mod +import pytest # type: ignore URL_PREFIX = "https://www.googleapis.com/upload/storage/v1/b/{BUCKET}/o" SIMPLE_URL = URL_PREFIX + "?uploadType=media&name={OBJECT}" @@ -361,6 +360,53 @@ def test_recover(self): timeout=EXPECTED_TIMEOUT, ) + def test_transmit_next_chunk_streaming_final_chunk_attaches_checksum(self): + import base64 + + import google_crc32c + + data = b"Streaming payload that finishes in one chunk." + upload = upload_mod.ResumableUpload(RESUMABLE_URL, ONE_MB, checksum="crc32c") + upload._stream = io.BytesIO(data) + upload._content_type = BASIC_CONTENT + upload._total_bytes = None # Unknown initial size (streaming write) + upload._resumable_url = "http://test.invalid?upload_id=not-none" + + crc32c_int = google_crc32c.value(data) + crc32c_bytes = crc32c_int.to_bytes(4, "big") + crc32c_b64 = base64.b64encode(crc32c_bytes).decode("utf-8") + + transport = mock.Mock(spec=["request"]) + put_response = mock.Mock( + status_code=http.client.OK, + headers={}, + json=mock.Mock(return_value={"crc32c": crc32c_b64}), + ) + transport.request.return_value = put_response + upload.transmit_next_chunk(transport) + + assert upload._total_bytes == len(data) + called_args, called_kwargs = transport.request.call_args + headers = called_kwargs["headers"] + assert "x-goog-hash" in headers + assert headers["x-goog-hash"] == f"crc32c={crc32c_b64}" + + def test_transmit_next_chunk_streaming_no_checksum_requested(self): + data = b"Streaming payload without checksum." + upload = upload_mod.ResumableUpload(RESUMABLE_URL, ONE_MB, checksum=None) + upload._stream = io.BytesIO(data) + upload._content_type = BASIC_CONTENT + upload._total_bytes = None + upload._resumable_url = "http://test.invalid?upload_id=not-none" + + transport = self._chunk_mock(http.client.OK, {}) + upload.transmit_next_chunk(transport) + + assert upload._total_bytes == len(data) + called_args, called_kwargs = transport.request.call_args + headers = called_kwargs["headers"] + assert "x-goog-hash" not in headers + def test_mpu_container(): container = upload_mod.XMLMPUContainer(EXAMPLE_XML_UPLOAD_URL, filename) diff --git a/packages/google-cloud-storage/tests/system/test_blob.py b/packages/google-cloud-storage/tests/system/test_blob.py index 8cea2add2d4a..76b93d2789b0 100644 --- a/packages/google-cloud-storage/tests/system/test_blob.py +++ b/packages/google-cloud-storage/tests/system/test_blob.py @@ -23,7 +23,6 @@ import mock import pytest from google.api_core import exceptions - from google.cloud.storage._helpers import _base64_md5hash from google.cloud.storage.exceptions import DataCorruption @@ -134,7 +133,7 @@ def test_large_file_write_from_stream_w_failed_checksum( "google.cloud.storage._media._helpers.prepare_checksum_digest", return_value="FFFFFF==", ): - with pytest.raises(DataCorruption): + with pytest.raises((DataCorruption, exceptions.BadRequest)): blob.upload_from_file(file_obj, checksum="crc32c") assert not blob.exists() @@ -1345,3 +1344,62 @@ def test_blob_contexts_custom_setter(shared_bucket, blobs_to_delete): blob.reload() assert blob.contexts.custom["k1"].value == "v1-updated" assert blob.contexts.custom["k2"].value == "v2" + + +def test_upload_from_file_streaming_with_trailing_checksum_validation( + shared_bucket, blobs_to_delete +): + import base64 + import io + import os + + import google_crc32c + + blob_name = f"StreamingTrailingChecksum-{uuid.uuid4().hex}" + blob = shared_bucket.blob(blob_name) + blobs_to_delete.append(blob) + + # Payload > 8 MiB to trigger chunked resumable upload + payload = os.urandom(8 * 1024 * 1024 + 1024) + io_stream = io.BytesIO(payload) + + # Calculate expected CRC32C base64 hash + expected_crc32c_int = google_crc32c.value(payload) + expected_crc32c_b64 = base64.b64encode( + expected_crc32c_int.to_bytes(4, "big") + ).decode("utf-8") + + # Upload from stream WITHOUT passing size parameter + # (routes to resumable upload with total_bytes=None) + blob.upload_from_file(io_stream, checksum="crc32c") + + blob.reload() + assert blob.size == len(payload) + assert blob.crc32c == expected_crc32c_b64 + + +def test_upload_from_file_streaming_corrupted_checksum_rejection( + shared_bucket, blobs_to_delete +): + import io + import os + + import pytest + from google.api_core.exceptions import BadRequest + + blob_name = f"StreamingCorruptedChecksum-{uuid.uuid4().hex}" + blob = shared_bucket.blob(blob_name) + blobs_to_delete.append(blob) + + payload = os.urandom(8 * 1024 * 1024 + 1024) + io_stream = io.BytesIO(payload) + + # Supply an intentionally corrupted/mismatched checksum + bad_crc32c_b64 = "AAAAAA==" + + with pytest.raises(BadRequest): + blob.upload_from_file( + io_stream, + checksum="crc32c", + crc32c_checksum_value=bad_crc32c_b64, + )