improve _CombinedChunksReader verification calls

Author: Fallen-Breath

prime_backup/action/helpers/blob_exporter.py

@@ -1,4 +1,3 @@
-import contextlib
 import dataclasses
 import functools
 import threading
@@ -93,10 +92,12 @@ def __init__(self, chunks_gen: Generator[_OpenedChunk, None, None]):
 		self.idx = 0
 		self.chunks_gen = chunks_gen
 		self.current: Optional[_OpenedChunk] = None
+		self.current_read_len = 0
+		self.current_verified = False
 		self.reach_end = False
 
 	def read(self, length: int = -1) -> bytes:
-		if self.reach_end:
+		if self.reach_end or length == 0:
 			return b''
 		if self.current is None:
 			if not self.__switch_to_next():
@@ -116,16 +117,33 @@ def read(self, length: int = -1) -> bytes:
 				logger.get().error(f'Failed to read {to_read} bytes from chunk {self.current.offset_chunk}: {e}')
 				raise
 			total_read += len(buf)
+			self.current_read_len += len(buf)
 			results.append(buf)
-			if to_read == -1 or len(buf) < to_read:
-				self.current.verify_callback()
+			if self.current_read_len >= self.current.offset_chunk.chunk.raw_size:
+				self.__verify_current()
 				if not self.__switch_to_next():
 					break
+			elif to_read == -1 or len(buf) < to_read:
+				raise EOFError('chunk {} exhausted after reading {} bytes, expected {}'.format(
+					self.current.offset_chunk, self.current_read_len, self.current.offset_chunk.chunk.raw_size,
+				))
 		return b''.join(results)
 
+	def close(self):
+		self.chunks_gen.close()
+		self.current = None
+		self.reach_end = True
+
+	def __verify_current(self):
+		if self.current is not None and not self.current_verified:
+			self.current.verify_callback()
+			self.current_verified = True
+
 	def __switch_to_next(self) -> bool:
 		try:
 			self.current = next(self.chunks_gen)
+			self.current_read_len = 0
+			self.current_verified = False
 			return True
 		except StopIteration:
 			self.reach_end = True
@@ -217,8 +235,6 @@ def __export_as_reader_direct(self, reader_csm: Callable[[SupportsReadBytes], An
 
 	def __export_as_reader_chunked(self, reader_csm: Callable[[SupportsReadBytes], Any]):
 		blob_chunks = self.blob_chunks_getter.get(self.blob.id)
-		exit_flag = False
-		gen_yielded = False
 
 		def open_chunk_gen() -> Generator[_OpenedChunk, None, None]:
 			hash_method = chunk_utils.get_hash_method()
@@ -231,8 +247,6 @@ def open_chunk_gen() -> Generator[_OpenedChunk, None, None]:
 						self.logger.error(f'Failed to peek-read chunk {oc.chunk.hash} in pack entry {oc.chunk.pack_entry.pack_id}@{oc.chunk.pack_entry.offset} for {oc}: {e}')
 						raise
 
-					nonlocal gen_yielded
-					gen_yielded = True
 					if self.verify_blob:
 						def verify_callback(ck: ChunkInfo, br: BypassReader):
 							self.__verify_exported_chunk(ck, br.get_read_len(), br.get_hash())
@@ -242,20 +256,13 @@ def verify_callback(ck: ChunkInfo, br: BypassReader):
 					else:
 						yield _OpenedChunk(oc, peek_reader, lambda: None)
 
-				nonlocal exit_flag
-				if exit_flag:
-					break
-
 		with PackFileObjectPool() as pack_file_obj_pool:
 			chunk_gen = open_chunk_gen()
+			reader = _CombinedChunksReader(chunk_gen)
 			try:
-				reader_csm(_CombinedChunksReader(chunk_gen))
+				reader_csm(reader)
 			finally:
-				# ensure possible opened ChunkIO is closed
-				exit_flag = True
-				if gen_yielded:
-					with contextlib.suppress(StopIteration):
-						next(chunk_gen)
+				reader.close()
 
 	def __verify_exported_blob(self, written_size: int, written_hash: str):
 		self.__verify_exported_data(lambda: 'blob', self.blob.raw_size, self.blob.hash, written_size, written_hash)

tests/test_pack_storage.py

@@ -1,7 +1,7 @@
 import dataclasses
 from io import BytesIO
 from pathlib import Path
-from typing import Dict, Generator, Optional
+from typing import Dict, Generator, List, Optional
 
 import pytest
 
@@ -11,6 +11,7 @@
 from prime_backup.action.delete_backup_file_action import DeleteBackupFileAction
 from prime_backup.action.export_backup_action_tar import ExportBackupToTarAction
 from prime_backup.action.get_pack_action import GetPackByFileNamePrefixAction, GetPackByIdAction
+from prime_backup.action.helpers.blob_exporter import _CombinedChunksReader, _OpenedChunk
 from prime_backup.action.helpers.chunk_io import ChunkIO
 from prime_backup.action.helpers.pack_reader import PackEntryReader
 from prime_backup.action.helpers.pack_writer import PackWriter
@@ -28,7 +29,7 @@
 from prime_backup.db.session import DbSession
 from prime_backup.exceptions import PackFileNameNotUnique
 from prime_backup.types.backup_info import BackupInfo
-from prime_backup.types.chunk_info import ChunkInfo
+from prime_backup.types.chunk_info import ChunkInfo, OffsetChunkInfo
 from prime_backup.types.chunk_method import ChunkMethod
 from prime_backup.types.hash_method import HashMethod
 from prime_backup.types.operator import Operator
@@ -317,6 +318,66 @@ def test_pack_entry_reader_seek_uses_entry_relative_offsets() -> None:
 		reader.seek(0, 12345)
 
 
+def test_combined_chunks_reader_verifies_exactly_consumed_final_chunk() -> None:
+	verified_chunks: List[int] = []
+	generator_closed = False
+
+	def make_opened_chunk(chunk_id: int, data: bytes) -> _OpenedChunk:
+		chunk = ChunkInfo(
+			id=chunk_id,
+			hash=str(chunk_id),
+			compress=CompressMethod.plain,
+			raw_size=len(data),
+			stored_size=len(data),
+			pack_entry=PackEntryLocation(chunk_id, 0),
+		)
+		return _OpenedChunk(
+			OffsetChunkInfo(0, chunk),
+			BytesIO(data),
+			lambda: verified_chunks.append(chunk_id),
+		)
+
+	def opened_chunks() -> Generator[_OpenedChunk, None, None]:
+		nonlocal generator_closed
+		try:
+			yield make_opened_chunk(1, b'ab')
+			yield make_opened_chunk(2, b'cd')
+		finally:
+			generator_closed = True
+
+	reader = _CombinedChunksReader(opened_chunks())
+	assert reader.read(4) == b'abcd'
+	assert verified_chunks == [1, 2]
+
+	reader.close()
+	assert verified_chunks == [1, 2]
+	assert generator_closed
+
+
+def test_combined_chunks_reader_raises_eof_error_on_short_chunk() -> None:
+	chunk = ChunkInfo(
+		id=1,
+		hash='short',
+		compress=CompressMethod.plain,
+		raw_size=4,
+		stored_size=2,
+		pack_entry=PackEntryLocation(1, 0),
+	)
+	def opened_chunks() -> Generator[_OpenedChunk, None, None]:
+		yield _OpenedChunk(
+			OffsetChunkInfo(0, chunk),
+			BytesIO(b'ab'),
+			lambda: None,
+		)
+
+	reader = _CombinedChunksReader(opened_chunks())
+
+	with pytest.raises(EOFError, match='exhausted after reading 2 bytes, expected 4'):
+		reader.read()
+
+	reader.close()
+
+
 def test_scan_unknown_pack_files_keeps_known_derived_pack_file_names(env: PackStorageEnv) -> None:
 	__create_backup()