diff --git a/README.md b/README.md index 8251c02..406d2db 100644 --- a/README.md +++ b/README.md @@ -323,7 +323,7 @@ python3 -m venv .venv Next install the project in editable mode with development dependencies: ```bash -pip install -m .[dev] +pip install -e .[dev] ``` Finally, to run tests use `unittest`: diff --git a/src/bmaptool/BmapCopy.py b/src/bmaptool/BmapCopy.py index 15113a8..1a006dd 100644 --- a/src/bmaptool/BmapCopy.py +++ b/src/bmaptool/BmapCopy.py @@ -59,6 +59,7 @@ import os import re import stat +import struct import sys import hashlib import logging @@ -792,6 +793,189 @@ def sync(self): ) +SPARSE_HEADER_MAGIC = 0xed26ff3a +CHUNK_TYPE_RAW = 0xCAC1 +CHUNK_TYPE_FILL = 0xCAC2 +CHUNK_TYPE_DONT_CARE = 0xCAC3 +CHUNK_TYPE_CRC32 = 0xCAC4 + +MAJOR_VERSION = 1 +MINOR_VERSION = 0 +FILE_HEADER_SIZE = 28 # always for version 1.0 +CHUNK_HEADER_SIZE = 12 # always for version 1.0 +CHECKSUM_DONT_CARE = 0 +CHUNK_HEADER_RESERVED1 = 0 + + +class BmapAndroidSparseImageCopy(BmapCopy): + """ + Copies the source image to an Android Sparse Image file + + See: https://android.googlesource.com/platform/system/core/+/refs/heads/main/libsparse/sparse_format.h + """ + def __init__(self, image, dest, bmap=None, image_size=None): + super().__init__(image, dest, bmap, image_size) + + self._sparse_image_chunk_cnt = 0 + + # Use bigger batch sizes to avoid so many raw chunks in the output + # Sparse Image file + self._batch_bytes = 100 * 2**20 + self._batch_blocks = self._batch_bytes // self.block_size + + + def _write_chunk(self, header, data=None): + assert(len(header) == CHUNK_HEADER_SIZE) + + try: + self._f_dest.write(header) + + if data is not None: + self._f_dest.write(data) + except IOError as err: + raise Error( + "error while writing blocks %d-%d of '%s': %s" + % (start, end, self._dest_path, err) + ) + + self._sparse_image_chunk_cnt += 1 + + + def _write_raw_chunk(self, buf): + assert(len(buf) % self.block_size == 0) + + block_count = len(buf) // self.block_size + + chunk_header = struct.pack("<2H2I", + CHUNK_TYPE_RAW, + CHUNK_HEADER_RESERVED1, + block_count, + CHUNK_HEADER_SIZE + len(buf), + ) + self._write_chunk(chunk_header, buf) + + + def _write_dont_care_chunk(self, block_count): + chunk_header = struct.pack("<2H2I", + CHUNK_TYPE_DONT_CARE, + CHUNK_HEADER_RESERVED1, + block_count, + CHUNK_HEADER_SIZE, + ) + self._write_chunk(chunk_header) + + + def copy(self, sync=True, verify=True): + # Create the queue for block batches and start the reader thread, which + # will read the image in batches and put the results to '_batch_queue'. + self._batch_queue = queue.Queue(self._batch_queue_len) + thread.start_new_thread(self._get_data, (verify,)) + + blocks_written = 0 + bytes_written = 0 + fsync_last = 0 + + self._sparse_image_chunk_cnt = 0 + self._progress_started = False + self._progress_index = 0 + self._progress_time = datetime.datetime.now() + + # skip past the file header. we'll seek back and write it later once we + # know all the information that goes in it + self._f_dest.seek(FILE_HEADER_SIZE) + + # Read the image in '_batch_blocks' chunks and write them to the + # destination file + prev_end = -1 + while True: + batch = self._batch_queue.get() + if batch is None: + # No more data, the image is written + break + elif batch[0] == "error": + # The reader thread encountered an error and passed us the + # exception. + exc_info = batch[1] + raise exc_info[1] + + (start, end, buf) = batch[1:4] + batch_block_count = end - start + 1 + + assert len(buf) <= (batch_block_count) * self.block_size + assert len(buf) > (end - start) * self.block_size + + # Synchronize the destination file if we reached the watermark + if self._dest_fsync_watermark: + if blocks_written >= fsync_last + self._dest_fsync_watermark: + fsync_last = blocks_written + self.sync() + + blocks_since_previous_chunk = start - prev_end - 1 + + if blocks_since_previous_chunk < 0: + raise Error("Out-of-order bmap blocks are not supported") + + if blocks_since_previous_chunk > 0: + self._write_dont_care_chunk(blocks_since_previous_chunk) + + self._write_raw_chunk(buf) + + self._batch_queue.task_done() + blocks_written += batch_block_count + bytes_written += len(buf) + + prev_end = end + + self._update_progress(blocks_written) + + if prev_end < self.blocks_cnt: + self._write_dont_care_chunk(self.blocks_cnt - prev_end - 1) + + if not self.image_size: + # The image size was unknown up until now, set it + self._set_image_size(bytes_written) + + # This is just a sanity check - we should have written exactly + # 'mapped_cnt' blocks. + if blocks_written != self.mapped_cnt: + raise Error( + "wrote %u blocks from image '%s' to '%s', but should " + "have %u - bmap file '%s' does not belong to this " + "image" + % ( + blocks_written, + self._image_path, + self._dest_path, + self.mapped_cnt, + self._bmap_path, + ) + ) + + file_header = struct.pack("