| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent e4a207d commit 2361ba6
2 files changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -36,14 +36,25 @@ class _DownloadState: | |||
| 36 | 36 | """A helper class to track the state of a single range download.""" | |
| 37 | 37 | ||
| 38 | 38 | def __init__( | |
| 39 | - self, initial_offset: int, initial_length: int, user_buffer: IO[bytes] | ||
| 39 | + self, | ||
| 40 | + initial_offset: int, | ||
| 41 | + initial_length: int, | ||
| 42 | + user_buffer: IO[bytes], | ||
| 43 | + is_full_object_read: bool = False, | ||
| 44 | + enable_checksum: bool = True, | ||
| 40 | 45 | ): | |
| 41 | 46 | self.initial_offset = initial_offset | |
| 42 | 47 | self.initial_length = initial_length | |
| 43 | 48 | self.user_buffer = user_buffer | |
| 44 | 49 | self.bytes_written = 0 | |
| 45 | 50 | self.next_expected_offset = initial_offset | |
| 46 | 51 | self.is_complete = False | |
| 52 | + self.is_full_object_read = is_full_object_read | ||
| 53 | + self.rolling_checksum = ( | ||
| 54 | + google_crc32c.Checksum() | ||
| 55 | + if (is_full_object_read and enable_checksum) | ||
| 56 | + else None | ||
| 57 | + ) | ||
| 47 | 58 | ||
| 48 | 59 | ||
| 49 | 60 | class _ReadResumptionStrategy(_BaseResumptionStrategy): | |
@@ -90,6 +101,7 @@ def update_state_from_response( | |||
| 90 | 101 | ) | |
| 91 | 102 | ||
| 92 | 103 | download_states = state["download_states"] | |
| 104 | + checksum_enabled = state.get("enable_checksum", True) | ||
| 93 | 105 | ||
| 94 | 106 | for object_data_range in proto.object_data_ranges: | |
| 95 | 107 | # Ignore empty ranges or ranges for IDs not in our state | |
@@ -125,7 +137,7 @@ def update_state_from_response( | |||
| 125 | 137 | checksummed_data = object_data_range.checksummed_data | |
| 126 | 138 | data = checksummed_data.content | |
| 127 | 139 | ||
| 128 | - if checksummed_data.HasField("crc32c"): | ||
| 140 | + if checksum_enabled and checksummed_data.HasField("crc32c"): | ||
| 129 | 141 | server_checksum = checksummed_data.crc32c | |
| 130 | 142 | client_checksum = google_crc32c.value(data) | |
| 131 | 143 | if server_checksum != client_checksum: | |
@@ -138,10 +150,14 @@ def update_state_from_response( | |||
| 138 | 150 | # Update State & Write Data | |
| 139 | 151 | chunk_size = len(data) | |
| 140 | 152 | read_state.user_buffer.write(data) | |
| 153 | + | ||
| 154 | + # Commit updates only after the write succeeds | ||
| 155 | + if checksum_enabled and read_state.rolling_checksum is not None: | ||
| 156 | + read_state.rolling_checksum.update(data) | ||
| 141 | 157 | read_state.bytes_written += chunk_size | |
| 142 | 158 | read_state.next_expected_offset += chunk_size | |
| 143 | 159 | ||
| 144 | - # Final Byte Count Verification | ||
| 160 | + # Final Byte Count & Full Object Checksum Verification | ||
| 145 | 161 | if object_data_range.range_end: | |
| 146 | 162 | read_state.is_complete = True | |
| 147 | 163 | if ( | |
@@ -154,6 +170,26 @@ def update_state_from_response( | |||
| 154 | 170 | f"Expected {read_state.initial_length}, got {read_state.bytes_written}", | |
| 155 | 171 | ) | |
| 156 | 172 | ||
| 173 | + # Perform full-object checksum verification once the stream finishes. | ||
| 174 | + if ( | ||
| 175 | + read_state.is_full_object_read | ||
| 176 | + and checksum_enabled | ||
| 177 | + and read_state.rolling_checksum is not None | ||
| 178 | + ): | ||
| 179 | + full_obj_server_crc32c = state.get("full_obj_server_crc32c") | ||
| 180 | + if full_obj_server_crc32c is not None: | ||
| 181 | + # Use standard big-endian byte conversion to retrieve the rolling checksum value. | ||
| 182 | + client_checksum = int.from_bytes( | ||
| 183 | + read_state.rolling_checksum.digest(), | ||
| 184 | + byteorder="big", | ||
| 185 | + ) | ||
| 186 | + if client_checksum != full_obj_server_crc32c: | ||
| 187 | + raise DataCorruption( | ||
| 188 | + response, | ||
| 189 | + f"Full object checksum mismatch for read_id {read_id}. " | ||
| 190 | + f"Server authoritative crc32c: {full_obj_server_crc32c}, client calculated rolling: {client_checksum}.", | ||
| 191 | + ) | ||
| 192 | + | ||
| 157 | 193 | async def recover_state_on_failure(self, error: Exception, state: Any) -> None: | |
| 158 | 194 | """Handles BidiReadObjectRedirectedError for reads.""" | |
| 159 | 195 | routing_token, read_handle = _handle_redirect(error) | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -45,6 +45,48 @@ def test_initialization(self): | |||
| 45 | 45 | self.assertEqual(state.bytes_written, 0) | |
| 46 | 46 | self.assertEqual(state.next_expected_offset, initial_offset) | |
| 47 | 47 | self.assertFalse(state.is_complete) | |
| 48 | + self.assertFalse(state.is_full_object_read) | ||
| 49 | + self.assertIsNone(state.rolling_checksum) | ||
| 50 | + | ||
| 51 | + def test_initialization_with_full_object_read(self): | ||
| 52 | + """Test that _DownloadState initializes correctly when is_full_object_read is True.""" | ||
| 53 | + initial_offset = 10 | ||
| 54 | + initial_length = 100 | ||
| 55 | + user_buffer = io.BytesIO() | ||
| 56 | + state_full = _DownloadState( | ||
| 57 | + initial_offset, initial_length, user_buffer, is_full_object_read=True | ||
| 58 | + ) | ||
| 59 | + | ||
| 60 | + self.assertEqual(state_full.initial_offset, initial_offset) | ||
| 61 | + self.assertEqual(state_full.initial_length, initial_length) | ||
| 62 | + self.assertEqual(state_full.user_buffer, user_buffer) | ||
| 63 | + self.assertEqual(state_full.bytes_written, 0) | ||
| 64 | + self.assertEqual(state_full.next_expected_offset, initial_offset) | ||
| 65 | + self.assertFalse(state_full.is_complete) | ||
| 66 | + self.assertTrue(state_full.is_full_object_read) | ||
| 67 | + self.assertIsNotNone(state_full.rolling_checksum) | ||
| 68 | + | ||
| 69 | + def test_initialization_with_full_object_read_and_checksum_disabled(self): | ||
| 70 | + """Test that _DownloadState does not initialize rolling_checksum when enable_checksum is False.""" | ||
| 71 | + initial_offset = 10 | ||
| 72 | + initial_length = 100 | ||
| 73 | + user_buffer = io.BytesIO() | ||
| 74 | + state_full = _DownloadState( | ||
| 75 | + initial_offset, | ||
| 76 | + initial_length, | ||
| 77 | + user_buffer, | ||
| 78 | + is_full_object_read=True, | ||
| 79 | + enable_checksum=False, | ||
| 80 | + ) | ||
| 81 | + | ||
| 82 | + self.assertEqual(state_full.initial_offset, initial_offset) | ||
| 83 | + self.assertEqual(state_full.initial_length, initial_length) | ||
| 84 | + self.assertEqual(state_full.user_buffer, user_buffer) | ||
| 85 | + self.assertEqual(state_full.bytes_written, 0) | ||
| 86 | + self.assertEqual(state_full.next_expected_offset, initial_offset) | ||
| 87 | + self.assertFalse(state_full.is_complete) | ||
| 88 | + self.assertTrue(state_full.is_full_object_read) | ||
| 89 | + self.assertIsNone(state_full.rolling_checksum) | ||
| 48 | 90 | ||
| 49 | 91 | ||
| 50 | 92 | class TestReadResumptionStrategy(unittest.TestCase): | |
@@ -53,12 +95,24 @@ def setUp(self): | |||
| 53 | 95 | ||
| 54 | 96 | self.state = {"download_states": {}, "read_handle": None, "routing_token": None} | |
| 55 | 97 | ||
| 56 | - def _add_download(self, read_id, offset=0, length=100, buffer=None): | ||
| 98 | + def _add_download( | ||
| 99 | + self, | ||
| 100 | + read_id, | ||
| 101 | + offset=0, | ||
| 102 | + length=100, | ||
| 103 | + buffer=None, | ||
| 104 | + is_full_object_read=False, | ||
| 105 | + enable_checksum=True, | ||
| 106 | + ): | ||
| 57 | 107 | """Helper to inject a download state into the correct nested location.""" | |
| 58 | 108 | if buffer is None: | |
| 59 | 109 | buffer = io.BytesIO() | |
| 60 | 110 | state = _DownloadState( | |
| 61 | - initial_offset=offset, initial_length=length, user_buffer=buffer | ||
| 111 | + initial_offset=offset, | ||
| 112 | + initial_length=length, | ||
| 113 | + user_buffer=buffer, | ||
| 114 | + is_full_object_read=is_full_object_read, | ||
| 115 | + enable_checksum=enable_checksum, | ||
| 62 | 116 | ) | |
| 63 | 117 | self.state["download_states"][read_id] = state | |
| 64 | 118 | return state | |
@@ -358,3 +412,61 @@ async def run(): | |||
| 358 | 412 | ||
| 359 | 413 | # Token should remain unchanged | |
| 360 | 414 | self.assertEqual(self.state["routing_token"], "existing-token") | |
| 415 | + | ||
| 416 | + def test_update_state_full_object_checksum_success(self): | ||
| 417 | + """Test that full object checksum verification succeeds on range_end.""" | ||
| 418 | + read_state = self._add_download( | ||
| 419 | + _READ_ID, offset=0, length=9, is_full_object_read=True | ||
| 420 | + ) | ||
| 421 | + self.state["enable_checksum"] = True | ||
| 422 | + self.state["full_obj_server_crc32c"] = google_crc32c.value(b"testdata1") | ||
| 423 | + | ||
| 424 | + resp1 = self._create_response(b"test", _READ_ID, offset=0) | ||
| 425 | + self.strategy.update_state_from_response(resp1, self.state) | ||
| 426 | + | ||
| 427 | + resp2 = self._create_response(b"data1", _READ_ID, offset=4, range_end=True) | ||
| 428 | + self.strategy.update_state_from_response(resp2, self.state) | ||
| 429 | + | ||
| 430 | + self.assertTrue(read_state.is_complete) | ||
| 431 | + self.assertEqual(read_state.bytes_written, 9) | ||
| 432 | + | ||
| 433 | + def test_update_state_full_object_checksum_failure(self): | ||
| 434 | + """Test that full object checksum verification raises DataCorruption on mismatch at range_end.""" | ||
| 435 | + self._add_download(_READ_ID, offset=0, length=9, is_full_object_read=True) | ||
| 436 | + self.state["enable_checksum"] = True | ||
| 437 | + self.state["full_obj_server_crc32c"] = 111111 # Wrong server checksum! | ||
| 438 | + | ||
| 439 | + resp1 = self._create_response(b"test", _READ_ID, offset=0) | ||
| 440 | + self.strategy.update_state_from_response(resp1, self.state) | ||
| 441 | + | ||
| 442 | + resp2 = self._create_response(b"data1", _READ_ID, offset=4, range_end=True) | ||
| 443 | + with self.assertRaisesRegex(DataCorruption, "Full object checksum mismatch"): | ||
| 444 | + self.strategy.update_state_from_response(resp2, self.state) | ||
| 445 | + | ||
| 446 | + def test_update_state_checksum_mismatch_ignored_when_disabled(self): | ||
| 447 | + """Test that a CRC32C mismatch is ignored when enable_checksum is False.""" | ||
| 448 | + self._add_download(_READ_ID) | ||
| 449 | + self.state["enable_checksum"] = False | ||
| 450 | + response = self._create_response(b"data", _READ_ID, offset=0, crc=999999) | ||
| 451 | + | ||
| 452 | + # Should NOT raise DataCorruption! | ||
| 453 | + self.strategy.update_state_from_response(response, self.state) | ||
| 454 | + | ||
| 455 | + def test_update_state_full_object_checksum_mismatch_ignored_when_disabled(self): | ||
| 456 | + """Test that a full-object CRC32C mismatch is ignored when enable_checksum is False.""" | ||
| 457 | + self._add_download( | ||
| 458 | + _READ_ID, | ||
| 459 | + offset=0, | ||
| 460 | + length=9, | ||
| 461 | + is_full_object_read=True, | ||
| 462 | + enable_checksum=False, | ||
| 463 | + ) | ||
| 464 | + self.state["enable_checksum"] = False | ||
| 465 | + self.state["full_obj_server_crc32c"] = 111111 # Wrong server checksum! | ||
| 466 | + | ||
| 467 | + resp1 = self._create_response(b"test", _READ_ID, offset=0) | ||
| 468 | + self.strategy.update_state_from_response(resp1, self.state) | ||
| 469 | + | ||
| 470 | + resp2 = self._create_response(b"data1", _READ_ID, offset=4, range_end=True) | ||
| 471 | + # Should NOT raise DataCorruption! | ||
| 472 | + self.strategy.update_state_from_response(resp2, self.state) | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments