diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py index 42f05a7ac..43c91dac7 100755 --- a/tools/releasetools/build_image.py +++ b/tools/releasetools/build_image.py @@ -29,18 +29,14 @@ from __future__ import print_function import os import os.path import re -import shlex import shutil import sys import common -import sparse_img - +import verity_utils OPTIONS = common.OPTIONS - -FIXED_SALT = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7" -BLOCK_SIZE = 4096 +BLOCK_SIZE = common.BLOCK_SIZE BYTES_IN_MB = 1024 * 1024 @@ -51,34 +47,6 @@ class BuildImageError(Exception): Exception.__init__(self, message) -def GetVerityFECSize(partition_size): - cmd = ["fec", "-s", str(partition_size)] - output = common.RunAndCheckOutput(cmd, verbose=False) - return int(output) - - -def GetVerityTreeSize(partition_size): - cmd = ["build_verity_tree", "-s", str(partition_size)] - output = common.RunAndCheckOutput(cmd, verbose=False) - return int(output) - - -def GetVerityMetadataSize(partition_size): - cmd = ["build_verity_metadata.py", "size", str(partition_size)] - output = common.RunAndCheckOutput(cmd, verbose=False) - return int(output) - - -def GetVeritySize(partition_size, fec_supported): - verity_tree_size = GetVerityTreeSize(partition_size) - verity_metadata_size = GetVerityMetadataSize(partition_size) - verity_size = verity_tree_size + verity_metadata_size - if fec_supported: - fec_size = GetVerityFECSize(partition_size + verity_size) - return verity_size + fec_size - return verity_size - - def GetDiskUsage(path): """Returns the number of bytes that "path" occupies on host. @@ -102,258 +70,6 @@ def GetDiskUsage(path): return int(output.split()[0]) * 512 -def GetSimgSize(image_file): - simg = sparse_img.SparseImage(image_file, build_map=False) - return simg.blocksize * simg.total_blocks - - -def ZeroPadSimg(image_file, pad_size): - blocks = pad_size // BLOCK_SIZE - print("Padding %d blocks (%d bytes)" % (blocks, pad_size)) - simg = sparse_img.SparseImage(image_file, mode="r+b", build_map=False) - simg.AppendFillChunk(0, blocks) - - -def AVBCalcMaxImageSize(avbtool, footer_type, partition_size, additional_args): - """Calculates max image size for a given partition size. - - Args: - avbtool: String with path to avbtool. - footer_type: 'hash' or 'hashtree' for generating footer. - partition_size: The size of the partition in question. - additional_args: Additional arguments to pass to "avbtool add_hash_footer" - or "avbtool add_hashtree_footer". - - Returns: - The maximum image size. - - Raises: - BuildImageError: On invalid image size. - """ - cmd = [avbtool, "add_%s_footer" % footer_type, - "--partition_size", str(partition_size), "--calc_max_image_size"] - cmd.extend(shlex.split(additional_args)) - - output = common.RunAndCheckOutput(cmd) - image_size = int(output) - if image_size <= 0: - raise BuildImageError( - "Invalid max image size: {}".format(output)) - return image_size - - -def AVBCalcMinPartitionSize(image_size, size_calculator): - """Calculates min partition size for a given image size. - - Args: - image_size: The size of the image in question. - size_calculator: The function to calculate max image size - for a given partition size. - - Returns: - The minimum partition size required to accommodate the image size. - """ - # Use image size as partition size to approximate final partition size. - image_ratio = size_calculator(image_size) / float(image_size) - - # Prepare a binary search for the optimal partition size. - lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - BLOCK_SIZE - - # Ensure lo is small enough: max_image_size should <= image_size. - delta = BLOCK_SIZE - max_image_size = size_calculator(lo) - while max_image_size > image_size: - image_ratio = max_image_size / float(lo) - lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - delta - delta *= 2 - max_image_size = size_calculator(lo) - - hi = lo + BLOCK_SIZE - - # Ensure hi is large enough: max_image_size should >= image_size. - delta = BLOCK_SIZE - max_image_size = size_calculator(hi) - while max_image_size < image_size: - image_ratio = max_image_size / float(hi) - hi = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE + delta - delta *= 2 - max_image_size = size_calculator(hi) - - partition_size = hi - - # Start to binary search. - while lo < hi: - mid = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE - max_image_size = size_calculator(mid) - if max_image_size >= image_size: # if mid can accommodate image_size - if mid < partition_size: # if a smaller partition size is found - partition_size = mid - hi = mid - else: - lo = mid + BLOCK_SIZE - - if OPTIONS.verbose: - print("AVBCalcMinPartitionSize({}): partition_size: {}.".format( - image_size, partition_size)) - - return partition_size - - -def AVBAddFooter(image_path, avbtool, footer_type, partition_size, - partition_name, key_path, algorithm, salt, - additional_args): - """Adds dm-verity hashtree and AVB metadata to an image. - - Args: - image_path: Path to image to modify. - avbtool: String with path to avbtool. - footer_type: 'hash' or 'hashtree' for generating footer. - partition_size: The size of the partition in question. - partition_name: The name of the partition - will be embedded in metadata. - key_path: Path to key to use or None. - algorithm: Name of algorithm to use or None. - salt: The salt to use (a hexadecimal string) or None. - additional_args: Additional arguments to pass to "avbtool add_hash_footer" - or "avbtool add_hashtree_footer". - """ - cmd = [avbtool, "add_%s_footer" % footer_type, - "--partition_size", partition_size, - "--partition_name", partition_name, - "--image", image_path] - - if key_path and algorithm: - cmd.extend(["--key", key_path, "--algorithm", algorithm]) - if salt: - cmd.extend(["--salt", salt]) - - cmd.extend(shlex.split(additional_args)) - - common.RunAndCheckOutput(cmd) - - -def AdjustPartitionSizeForVerity(partition_size, fec_supported): - """Modifies the provided partition size to account for the verity metadata. - - This information is used to size the created image appropriately. - - Args: - partition_size: the size of the partition to be verified. - - Returns: - A tuple of the size of the partition adjusted for verity metadata, and - the size of verity metadata. - """ - key = "%d %d" % (partition_size, fec_supported) - if key in AdjustPartitionSizeForVerity.results: - return AdjustPartitionSizeForVerity.results[key] - - hi = partition_size - if hi % BLOCK_SIZE != 0: - hi = (hi // BLOCK_SIZE) * BLOCK_SIZE - - # verity tree and fec sizes depend on the partition size, which - # means this estimate is always going to be unnecessarily small - verity_size = GetVeritySize(hi, fec_supported) - lo = partition_size - verity_size - result = lo - - # do a binary search for the optimal size - while lo < hi: - i = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE - v = GetVeritySize(i, fec_supported) - if i + v <= partition_size: - if result < i: - result = i - verity_size = v - lo = i + BLOCK_SIZE - else: - hi = i - - if OPTIONS.verbose: - print("Adjusted partition size for verity, partition_size: {}," - " verity_size: {}".format(result, verity_size)) - AdjustPartitionSizeForVerity.results[key] = (result, verity_size) - return (result, verity_size) - - -AdjustPartitionSizeForVerity.results = {} - - -def BuildVerityFEC(sparse_image_path, verity_path, verity_fec_path, - padding_size): - cmd = ["fec", "-e", "-p", str(padding_size), sparse_image_path, - verity_path, verity_fec_path] - common.RunAndCheckOutput(cmd) - - -def BuildVerityTree(sparse_image_path, verity_image_path): - cmd = ["build_verity_tree", "-A", FIXED_SALT, sparse_image_path, - verity_image_path] - output = common.RunAndCheckOutput(cmd) - root, salt = output.split() - return root, salt - - -def BuildVerityMetadata(image_size, verity_metadata_path, root_hash, salt, - block_device, signer_path, key, signer_args, - verity_disable): - cmd = ["build_verity_metadata.py", "build", str(image_size), - verity_metadata_path, root_hash, salt, block_device, signer_path, key] - if signer_args: - cmd.append("--signer_args=\"%s\"" % (' '.join(signer_args),)) - if verity_disable: - cmd.append("--verity_disable") - common.RunAndCheckOutput(cmd) - - -def Append2Simg(sparse_image_path, unsparse_image_path, error_message): - """Appends the unsparse image to the given sparse image. - - Args: - sparse_image_path: the path to the (sparse) image - unsparse_image_path: the path to the (unsparse) image - - Raises: - BuildImageError: On error. - """ - cmd = ["append2simg", sparse_image_path, unsparse_image_path] - try: - common.RunAndCheckOutput(cmd) - except: - raise BuildImageError(error_message) - - -def Append(target, file_to_append, error_message): - """Appends file_to_append to target. - - Raises: - BuildImageError: On error. - """ - try: - with open(target, "a") as out_file, open(file_to_append, "r") as input_file: - for line in input_file: - out_file.write(line) - except IOError: - raise BuildImageError(error_message) - - -def BuildVerifiedImage(data_image_path, verity_image_path, - verity_metadata_path, verity_fec_path, - padding_size, fec_supported): - Append( - verity_image_path, verity_metadata_path, - "Could not append verity metadata!") - - if fec_supported: - # Build FEC for the entire partition, including metadata. - BuildVerityFEC( - data_image_path, verity_image_path, verity_fec_path, padding_size) - Append(verity_image_path, verity_fec_path, "Could not append FEC!") - - Append2Simg( - data_image_path, verity_image_path, "Could not append verity data!") - - def UnsparseImage(sparse_image_path, replace=True): img_dir = os.path.dirname(sparse_image_path) unsparse_image_path = "unsparse_" + os.path.basename(sparse_image_path) @@ -372,56 +88,6 @@ def UnsparseImage(sparse_image_path, replace=True): return unsparse_image_path -def MakeVerityEnabledImage(out_file, fec_supported, prop_dict): - """Creates an image that is verifiable using dm-verity. - - Args: - out_file: the location to write the verifiable image at - prop_dict: a dictionary of properties required for image creation and - verification - - Raises: - AssertionError: On invalid partition sizes. - BuildImageError: On other errors. - """ - # get properties - image_size = int(prop_dict["image_size"]) - block_dev = prop_dict["verity_block_device"] - signer_key = prop_dict["verity_key"] + ".pk8" - if OPTIONS.verity_signer_path is not None: - signer_path = OPTIONS.verity_signer_path - else: - signer_path = prop_dict["verity_signer_cmd"] - signer_args = OPTIONS.verity_signer_args - - tempdir_name = common.MakeTempDir(suffix="_verity_images") - - # Get partial image paths. - verity_image_path = os.path.join(tempdir_name, "verity.img") - verity_metadata_path = os.path.join(tempdir_name, "verity_metadata.img") - verity_fec_path = os.path.join(tempdir_name, "verity_fec.img") - - # Build the verity tree and get the root hash and salt. - root_hash, salt = BuildVerityTree(out_file, verity_image_path) - - # Build the metadata blocks. - verity_disable = "verity_disable" in prop_dict - BuildVerityMetadata( - image_size, verity_metadata_path, root_hash, salt, block_dev, signer_path, - signer_key, signer_args, verity_disable) - - # Build the full verified image. - partition_size = int(prop_dict["partition_size"]) - verity_size = int(prop_dict["verity_size"]) - - padding_size = partition_size - image_size - verity_size - assert padding_size >= 0 - - BuildVerifiedImage( - out_file, verity_image_path, verity_metadata_path, verity_fec_path, - padding_size, fec_supported) - - def ConvertBlockMapToBaseFs(block_map_file): base_fs_file = common.MakeTempFile(prefix="script_gen_", suffix=".base_fs") convert_command = ["blk_alloc_to_base_fs", block_map_file, base_fs_file] @@ -570,9 +236,9 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None): # Adjust partition_size to add more space for AVB footer, to prevent # it from consuming partition_reserved_size. if avb_footer_type: - size = AVBCalcMinPartitionSize( + size = verity_utils.AVBCalcMinPartitionSize( size, - lambda x: AVBCalcMaxImageSize( + lambda x: verity_utils.AVBCalcMaxImageSize( avbtool, avb_footer_type, x, avb_signing_args)) prop_dict["partition_size"] = str(size) if OPTIONS.verbose: @@ -583,7 +249,7 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None): # Adjust the image size to make room for the hashes if this is to be verified. if verity_supported and is_verity_partition: partition_size = int(prop_dict.get("partition_size")) - image_size, verity_size = AdjustPartitionSizeForVerity( + image_size, verity_size = verity_utils.AdjustPartitionSizeForVerity( partition_size, verity_fec_supported) prop_dict["image_size"] = str(image_size) prop_dict["verity_size"] = str(verity_size) @@ -592,7 +258,7 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None): if avb_footer_type: partition_size = prop_dict["partition_size"] # avb_add_hash_footer_args or avb_add_hashtree_footer_args. - max_image_size = AVBCalcMaxImageSize( + max_image_size = verity_utils.AVBCalcMaxImageSize( avbtool, avb_footer_type, partition_size, avb_signing_args) prop_dict["image_size"] = str(max_image_size) @@ -709,17 +375,18 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None): if not fs_spans_partition: mount_point = prop_dict.get("mount_point") image_size = int(prop_dict["image_size"]) - sparse_image_size = GetSimgSize(out_file) + sparse_image_size = verity_utils.GetSimgSize(out_file) if sparse_image_size > image_size: raise BuildImageError( "Error: {} image size of {} is larger than partition size of " "{}".format(mount_point, sparse_image_size, image_size)) if verity_supported and is_verity_partition: - ZeroPadSimg(out_file, image_size - sparse_image_size) + verity_utils.ZeroPadSimg(out_file, image_size - sparse_image_size) # Create the verified image if this is to be verified. if verity_supported and is_verity_partition: - MakeVerityEnabledImage(out_file, verity_fec_supported, prop_dict) + verity_utils.MakeVerityEnabledImage( + out_file, verity_fec_supported, prop_dict) # Add AVB HASH or HASHTREE footer (metadata). if avb_footer_type: @@ -729,7 +396,7 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None): key_path = prop_dict.get("avb_key_path") algorithm = prop_dict.get("avb_algorithm") salt = prop_dict.get("avb_salt") - AVBAddFooter( + verity_utils.AVBAddFooter( out_file, avbtool, avb_footer_type, partition_size, partition_name, key_path, algorithm, salt, avb_signing_args) diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py index d1bfc8f2f..7cca7663f 100644 --- a/tools/releasetools/common.py +++ b/tools/releasetools/common.py @@ -73,6 +73,9 @@ class Options(object): OPTIONS = Options() +# The block size that's used across the releasetools scripts. +BLOCK_SIZE = 4096 + # Values for "certificate" in apkcerts that mean special things. SPECIAL_CERT_STRINGS = ("PRESIGNED", "EXTERNAL") diff --git a/tools/releasetools/test_build_image.py b/tools/releasetools/test_build_image.py index 6f853e859..634c6b191 100644 --- a/tools/releasetools/test_build_image.py +++ b/tools/releasetools/test_build_image.py @@ -15,14 +15,11 @@ # import filecmp -import math import os.path -import random import common from build_image import ( - AVBCalcMinPartitionSize, BLOCK_SIZE, BuildImageError, CheckHeadroom, - SetUpInDirAndFsConfig) + BuildImageError, CheckHeadroom, SetUpInDirAndFsConfig) from test_utils import ReleaseToolsTestCase @@ -32,13 +29,6 @@ class BuildImageTest(ReleaseToolsTestCase): EXT4FS_OUTPUT = ( "Created filesystem with 2777/129024 inodes and 515099/516099 blocks") - def setUp(self): - # To test AVBCalcMinPartitionSize(), by using 200MB to 2GB image size. - # - 51200 = 200MB * 1024 * 1024 / 4096 - # - 524288 = 2GB * 1024 * 1024 * 1024 / 4096 - self._image_sizes = [BLOCK_SIZE * random.randint(51200, 524288) + offset - for offset in range(BLOCK_SIZE)] - def test_CheckHeadroom_SizeUnderLimit(self): # Required headroom: 1000 blocks. prop_dict = { @@ -186,51 +176,3 @@ class BuildImageTest(ReleaseToolsTestCase): self.assertIn('fs-config-system\n', fs_config_data) self.assertIn('fs-config-root\n', fs_config_data) self.assertEqual('/', prop_dict['mount_point']) - - def test_AVBCalcMinPartitionSize_LinearFooterSize(self): - """Tests with footer size which is linear to partition size.""" - for image_size in self._image_sizes: - for ratio in 0.95, 0.56, 0.22: - expected_size = common.RoundUpTo4K(int(math.ceil(image_size / ratio))) - self.assertEqual( - expected_size, - AVBCalcMinPartitionSize(image_size, lambda x: int(x * ratio))) - - def test_AVBCalcMinPartitionSize_SlowerGrowthFooterSize(self): - """Tests with footer size which grows slower than partition size.""" - - def _SizeCalculator(partition_size): - """Footer size is the power of 0.95 of partition size.""" - # Minus footer size to return max image size. - return partition_size - int(math.pow(partition_size, 0.95)) - - for image_size in self._image_sizes: - min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator) - # Checks min_partition_size can accommodate image_size. - self.assertGreaterEqual( - _SizeCalculator(min_partition_size), - image_size) - # Checks min_partition_size (round to BLOCK_SIZE) is the minimum. - self.assertLess( - _SizeCalculator(min_partition_size - BLOCK_SIZE), - image_size) - - def test_AVBCalcMinPartitionSize_FasterGrowthFooterSize(self): - """Tests with footer size which grows faster than partition size.""" - - def _SizeCalculator(partition_size): - """Max image size is the power of 0.95 of partition size.""" - # Max image size grows less than partition size, which means - # footer size grows faster than partition size. - return int(math.pow(partition_size, 0.95)) - - for image_size in self._image_sizes: - min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator) - # Checks min_partition_size can accommodate image_size. - self.assertGreaterEqual( - _SizeCalculator(min_partition_size), - image_size) - # Checks min_partition_size (round to BLOCK_SIZE) is the minimum. - self.assertLess( - _SizeCalculator(min_partition_size - BLOCK_SIZE), - image_size) diff --git a/tools/releasetools/test_verity_utils.py b/tools/releasetools/test_verity_utils.py index a9cd17b1f..0988d8e70 100644 --- a/tools/releasetools/test_verity_utils.py +++ b/tools/releasetools/test_verity_utils.py @@ -16,15 +16,17 @@ """Unittests for verity_utils.py.""" +import math import os.path +import random -import build_image import common import sparse_img from rangelib import RangeSet from test_utils import get_testdata_dir, ReleaseToolsTestCase from verity_utils import ( - CreateHashtreeInfoGenerator, HashtreeInfo, + AdjustPartitionSizeForVerity, AVBCalcMinPartitionSize, BLOCK_SIZE, + CreateHashtreeInfoGenerator, HashtreeInfo, MakeVerityEnabledImage, VerifiedBootVersion1HashtreeInfoGenerator) @@ -62,7 +64,7 @@ class VerifiedBootVersion1HashtreeInfoGeneratorTest(ReleaseToolsTestCase): def _generate_image(self): partition_size = 1024 * 1024 - adjusted_size, verity_size = build_image.AdjustPartitionSizeForVerity( + adjusted_size, verity_size = AdjustPartitionSizeForVerity( partition_size, True) raw_image = "" @@ -80,7 +82,7 @@ class VerifiedBootVersion1HashtreeInfoGeneratorTest(ReleaseToolsTestCase): 'verity_signer_cmd': 'verity_signer', 'verity_size': str(verity_size), } - build_image.MakeVerityEnabledImage(output_file, True, prop_dict) + MakeVerityEnabledImage(output_file, True, prop_dict) return output_file @@ -159,3 +161,62 @@ class VerifiedBootVersion1HashtreeInfoGeneratorTest(ReleaseToolsTestCase): self.assertEqual(self.hash_algorithm, info.hash_algorithm) self.assertEqual(self.fixed_salt, info.salt) self.assertEqual(self.expected_root_hash, info.root_hash) + + +class VerityUtilsTest(ReleaseToolsTestCase): + + def setUp(self): + # To test AVBCalcMinPartitionSize(), by using 200MB to 2GB image size. + # - 51200 = 200MB * 1024 * 1024 / 4096 + # - 524288 = 2GB * 1024 * 1024 * 1024 / 4096 + self._image_sizes = [BLOCK_SIZE * random.randint(51200, 524288) + offset + for offset in range(BLOCK_SIZE)] + + def test_AVBCalcMinPartitionSize_LinearFooterSize(self): + """Tests with footer size which is linear to partition size.""" + for image_size in self._image_sizes: + for ratio in 0.95, 0.56, 0.22: + expected_size = common.RoundUpTo4K(int(math.ceil(image_size / ratio))) + self.assertEqual( + expected_size, + AVBCalcMinPartitionSize( + image_size, lambda x, ratio=ratio: int(x * ratio))) + + def test_AVBCalcMinPartitionSize_SlowerGrowthFooterSize(self): + """Tests with footer size which grows slower than partition size.""" + + def _SizeCalculator(partition_size): + """Footer size is the power of 0.95 of partition size.""" + # Minus footer size to return max image size. + return partition_size - int(math.pow(partition_size, 0.95)) + + for image_size in self._image_sizes: + min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator) + # Checks min_partition_size can accommodate image_size. + self.assertGreaterEqual( + _SizeCalculator(min_partition_size), + image_size) + # Checks min_partition_size (round to BLOCK_SIZE) is the minimum. + self.assertLess( + _SizeCalculator(min_partition_size - BLOCK_SIZE), + image_size) + + def test_AVBCalcMinPartitionSize_FasterGrowthFooterSize(self): + """Tests with footer size which grows faster than partition size.""" + + def _SizeCalculator(partition_size): + """Max image size is the power of 0.95 of partition size.""" + # Max image size grows less than partition size, which means + # footer size grows faster than partition size. + return int(math.pow(partition_size, 0.95)) + + for image_size in self._image_sizes: + min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator) + # Checks min_partition_size can accommodate image_size. + self.assertGreaterEqual( + _SizeCalculator(min_partition_size), + image_size) + # Checks min_partition_size (round to BLOCK_SIZE) is the minimum. + self.assertLess( + _SizeCalculator(min_partition_size - BLOCK_SIZE), + image_size) diff --git a/tools/releasetools/verity_utils.py b/tools/releasetools/verity_utils.py index c512ef3bb..626a1ddd7 100644 --- a/tools/releasetools/verity_utils.py +++ b/tools/releasetools/verity_utils.py @@ -16,13 +16,354 @@ from __future__ import print_function +import os.path +import shlex import struct import common -from build_image import (AdjustPartitionSizeForVerity, GetVerityTreeSize, - GetVerityMetadataSize, BuildVerityTree) +import sparse_img from rangelib import RangeSet +OPTIONS = common.OPTIONS +BLOCK_SIZE = common.BLOCK_SIZE +FIXED_SALT = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7" + + +class BuildVerityImageError(Exception): + """An Exception raised during verity image building.""" + + def __init__(self, message): + Exception.__init__(self, message) + + +def GetVerityFECSize(partition_size): + cmd = ["fec", "-s", str(partition_size)] + output = common.RunAndCheckOutput(cmd, verbose=False) + return int(output) + + +def GetVerityTreeSize(partition_size): + cmd = ["build_verity_tree", "-s", str(partition_size)] + output = common.RunAndCheckOutput(cmd, verbose=False) + return int(output) + + +def GetVerityMetadataSize(partition_size): + cmd = ["build_verity_metadata.py", "size", str(partition_size)] + output = common.RunAndCheckOutput(cmd, verbose=False) + return int(output) + + +def GetVeritySize(partition_size, fec_supported): + verity_tree_size = GetVerityTreeSize(partition_size) + verity_metadata_size = GetVerityMetadataSize(partition_size) + verity_size = verity_tree_size + verity_metadata_size + if fec_supported: + fec_size = GetVerityFECSize(partition_size + verity_size) + return verity_size + fec_size + return verity_size + + +def GetSimgSize(image_file): + simg = sparse_img.SparseImage(image_file, build_map=False) + return simg.blocksize * simg.total_blocks + + +def ZeroPadSimg(image_file, pad_size): + blocks = pad_size // BLOCK_SIZE + print("Padding %d blocks (%d bytes)" % (blocks, pad_size)) + simg = sparse_img.SparseImage(image_file, mode="r+b", build_map=False) + simg.AppendFillChunk(0, blocks) + + +def AdjustPartitionSizeForVerity(partition_size, fec_supported): + """Modifies the provided partition size to account for the verity metadata. + + This information is used to size the created image appropriately. + + Args: + partition_size: the size of the partition to be verified. + + Returns: + A tuple of the size of the partition adjusted for verity metadata, and + the size of verity metadata. + """ + key = "%d %d" % (partition_size, fec_supported) + if key in AdjustPartitionSizeForVerity.results: + return AdjustPartitionSizeForVerity.results[key] + + hi = partition_size + if hi % BLOCK_SIZE != 0: + hi = (hi // BLOCK_SIZE) * BLOCK_SIZE + + # verity tree and fec sizes depend on the partition size, which + # means this estimate is always going to be unnecessarily small + verity_size = GetVeritySize(hi, fec_supported) + lo = partition_size - verity_size + result = lo + + # do a binary search for the optimal size + while lo < hi: + i = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE + v = GetVeritySize(i, fec_supported) + if i + v <= partition_size: + if result < i: + result = i + verity_size = v + lo = i + BLOCK_SIZE + else: + hi = i + + if OPTIONS.verbose: + print("Adjusted partition size for verity, partition_size: {}," + " verity_size: {}".format(result, verity_size)) + AdjustPartitionSizeForVerity.results[key] = (result, verity_size) + return (result, verity_size) + + +AdjustPartitionSizeForVerity.results = {} + + +def BuildVerityFEC(sparse_image_path, verity_path, verity_fec_path, + padding_size): + cmd = ["fec", "-e", "-p", str(padding_size), sparse_image_path, + verity_path, verity_fec_path] + common.RunAndCheckOutput(cmd) + + +def BuildVerityTree(sparse_image_path, verity_image_path): + cmd = ["build_verity_tree", "-A", FIXED_SALT, sparse_image_path, + verity_image_path] + output = common.RunAndCheckOutput(cmd) + root, salt = output.split() + return root, salt + + +def BuildVerityMetadata(image_size, verity_metadata_path, root_hash, salt, + block_device, signer_path, key, signer_args, + verity_disable): + cmd = ["build_verity_metadata.py", "build", str(image_size), + verity_metadata_path, root_hash, salt, block_device, signer_path, key] + if signer_args: + cmd.append("--signer_args=\"%s\"" % (' '.join(signer_args),)) + if verity_disable: + cmd.append("--verity_disable") + common.RunAndCheckOutput(cmd) + + +def Append2Simg(sparse_image_path, unsparse_image_path, error_message): + """Appends the unsparse image to the given sparse image. + + Args: + sparse_image_path: the path to the (sparse) image + unsparse_image_path: the path to the (unsparse) image + + Raises: + BuildVerityImageError: On error. + """ + cmd = ["append2simg", sparse_image_path, unsparse_image_path] + try: + common.RunAndCheckOutput(cmd) + except: + raise BuildVerityImageError(error_message) + + +def Append(target, file_to_append, error_message): + """Appends file_to_append to target. + + Raises: + BuildVerityImageError: On error. + """ + try: + with open(target, "a") as out_file, open(file_to_append, "r") as input_file: + for line in input_file: + out_file.write(line) + except IOError: + raise BuildVerityImageError(error_message) + + +def BuildVerifiedImage(data_image_path, verity_image_path, + verity_metadata_path, verity_fec_path, + padding_size, fec_supported): + Append( + verity_image_path, verity_metadata_path, + "Could not append verity metadata!") + + if fec_supported: + # Build FEC for the entire partition, including metadata. + BuildVerityFEC( + data_image_path, verity_image_path, verity_fec_path, padding_size) + Append(verity_image_path, verity_fec_path, "Could not append FEC!") + + Append2Simg( + data_image_path, verity_image_path, "Could not append verity data!") + + +def MakeVerityEnabledImage(out_file, fec_supported, prop_dict): + """Creates an image that is verifiable using dm-verity. + + Args: + out_file: the location to write the verifiable image at + prop_dict: a dictionary of properties required for image creation and + verification + + Raises: + AssertionError: On invalid partition sizes. + """ + # get properties + image_size = int(prop_dict["image_size"]) + block_dev = prop_dict["verity_block_device"] + signer_key = prop_dict["verity_key"] + ".pk8" + if OPTIONS.verity_signer_path is not None: + signer_path = OPTIONS.verity_signer_path + else: + signer_path = prop_dict["verity_signer_cmd"] + signer_args = OPTIONS.verity_signer_args + + tempdir_name = common.MakeTempDir(suffix="_verity_images") + + # Get partial image paths. + verity_image_path = os.path.join(tempdir_name, "verity.img") + verity_metadata_path = os.path.join(tempdir_name, "verity_metadata.img") + verity_fec_path = os.path.join(tempdir_name, "verity_fec.img") + + # Build the verity tree and get the root hash and salt. + root_hash, salt = BuildVerityTree(out_file, verity_image_path) + + # Build the metadata blocks. + verity_disable = "verity_disable" in prop_dict + BuildVerityMetadata( + image_size, verity_metadata_path, root_hash, salt, block_dev, signer_path, + signer_key, signer_args, verity_disable) + + # Build the full verified image. + partition_size = int(prop_dict["partition_size"]) + verity_size = int(prop_dict["verity_size"]) + + padding_size = partition_size - image_size - verity_size + assert padding_size >= 0 + + BuildVerifiedImage( + out_file, verity_image_path, verity_metadata_path, verity_fec_path, + padding_size, fec_supported) + + +def AVBCalcMaxImageSize(avbtool, footer_type, partition_size, additional_args): + """Calculates max image size for a given partition size. + + Args: + avbtool: String with path to avbtool. + footer_type: 'hash' or 'hashtree' for generating footer. + partition_size: The size of the partition in question. + additional_args: Additional arguments to pass to "avbtool add_hash_footer" + or "avbtool add_hashtree_footer". + + Returns: + The maximum image size. + + Raises: + BuildVerityImageError: On invalid image size. + """ + cmd = [avbtool, "add_%s_footer" % footer_type, + "--partition_size", str(partition_size), "--calc_max_image_size"] + cmd.extend(shlex.split(additional_args)) + + output = common.RunAndCheckOutput(cmd) + image_size = int(output) + if image_size <= 0: + raise BuildVerityImageError( + "Invalid max image size: {}".format(output)) + return image_size + + +def AVBCalcMinPartitionSize(image_size, size_calculator): + """Calculates min partition size for a given image size. + + Args: + image_size: The size of the image in question. + size_calculator: The function to calculate max image size + for a given partition size. + + Returns: + The minimum partition size required to accommodate the image size. + """ + # Use image size as partition size to approximate final partition size. + image_ratio = size_calculator(image_size) / float(image_size) + + # Prepare a binary search for the optimal partition size. + lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - BLOCK_SIZE + + # Ensure lo is small enough: max_image_size should <= image_size. + delta = BLOCK_SIZE + max_image_size = size_calculator(lo) + while max_image_size > image_size: + image_ratio = max_image_size / float(lo) + lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - delta + delta *= 2 + max_image_size = size_calculator(lo) + + hi = lo + BLOCK_SIZE + + # Ensure hi is large enough: max_image_size should >= image_size. + delta = BLOCK_SIZE + max_image_size = size_calculator(hi) + while max_image_size < image_size: + image_ratio = max_image_size / float(hi) + hi = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE + delta + delta *= 2 + max_image_size = size_calculator(hi) + + partition_size = hi + + # Start to binary search. + while lo < hi: + mid = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE + max_image_size = size_calculator(mid) + if max_image_size >= image_size: # if mid can accommodate image_size + if mid < partition_size: # if a smaller partition size is found + partition_size = mid + hi = mid + else: + lo = mid + BLOCK_SIZE + + if OPTIONS.verbose: + print("AVBCalcMinPartitionSize({}): partition_size: {}.".format( + image_size, partition_size)) + + return partition_size + + +def AVBAddFooter(image_path, avbtool, footer_type, partition_size, + partition_name, key_path, algorithm, salt, + additional_args): + """Adds dm-verity hashtree and AVB metadata to an image. + + Args: + image_path: Path to image to modify. + avbtool: String with path to avbtool. + footer_type: 'hash' or 'hashtree' for generating footer. + partition_size: The size of the partition in question. + partition_name: The name of the partition - will be embedded in metadata. + key_path: Path to key to use or None. + algorithm: Name of algorithm to use or None. + salt: The salt to use (a hexadecimal string) or None. + additional_args: Additional arguments to pass to "avbtool add_hash_footer" + or "avbtool add_hashtree_footer". + """ + cmd = [avbtool, "add_%s_footer" % footer_type, + "--partition_size", partition_size, + "--partition_name", partition_name, + "--image", image_path] + + if key_path and algorithm: + cmd.extend(["--key", key_path, "--algorithm", algorithm]) + if salt: + cmd.extend(["--salt", salt]) + + cmd.extend(shlex.split(additional_args)) + + common.RunAndCheckOutput(cmd) + class HashtreeInfoGenerationError(Exception): """An Exception raised during hashtree info generation."""