Add options to sign the prebuilt custom images.

The custom images are any images owned by OEMs and SoCs, oem images
mounted on /oem is an example. The oem images can be used to customize
devices for different carriers, like wallpaper, ringtones, and
carrier-specific apks. OEMs can generate multiple oem images, like
oem.img, oem-carrier1.img and oem-carrier2.img and flash different oem
images for different carriers. The oem images are only one case, OEMs
and SoCs can add more custom images and mount them to custom partitions.

This change enables custom images to be vbmeta.img chained partitions.
The following configuration in BoardConfig.mk is an exmaple. It has two
custom partitions: oem and test. They will be signed by different keys.
And they will be chained by vbmeta.img. The custom images here are
prebuilts, which can be built by `make custom_images` separately.
BOARD_AVB_<CUSTOM_PARTITION>_IMAGE_LIST should include all custom images
to apply AVB signing. And to every custom partition, one image whose
name is partition name must be added in its
BOARD_AVB_<CUSTOM_PARTITION>_IMAGE_LIST.

BOARD_CUSTOMIMAGES_PARTITION_LIST := oem test

BOARD_AVB_OEM_KEY_PATH := external/avb/test/data/testkey_rsa4096.pem
BOARD_AVB_OEM_ALGORITHM := SHA256_RSA4096
BOARD_AVB_OEM_ADD_HASHTREE_FOOTER_ARGS :=
BOARD_AVB_OEM_ROLLBACK_INDEX_LOCATION := 1
BOARD_AVB_OEM_PARTITION_SIZE := 5242880
BOARD_AVB_OEM_IMAGE_LIST := \
    device/xxxx/yyyy/oem/oem.img \
    device/xxxx/yyyy/oem/oem1.img

BOARD_AVB_TEST_KEY_PATH := external/avb/test/data/testkey_rsa2048.pem
BOARD_AVB_TEST_ALGORITHM := SHA256_RSA2048
BOARD_AVB_TEST_ADD_HASHTREE_FOOTER_ARGS :=
BOARD_AVB_TEST_ROLLBACK_INDEX_LOCATION := 2
BOARD_AVB_TEST_PARTITION_SIZE := 10485760
BOARD_AVB_TEST_IMAGE_LIST := \
    device/xxxx/yyyy/test/test.img \
    device/xxxx/yyyy/test/test1.img

To resign the custom images in the target zip file, the
avb_extra_custom_image_key, avb_extra_custom_image_algorithms and
avb_extra_custom_image_extra_args options are added to the
sign_target_files_apks tool too. The following test cases list some
examples about how to use them.

BUG: 154171021
Test: 1) "atest --host releasetools_test releasetools_py3_test -c"
  2) Build images by 'make dist', sign and validate target files.
   a) Test on dist w/ chained vbmeta_system and ome custom images
   sign_target_files_apks -d certs \
     --avb_extra_custom_image_key oem=oem_rsa4096.pem \
     --avb_extra_custom_image_algorithm oem=SHA256_RSA4096 \
     xxx-target_xxx.zip signed.zip
   validate_target_files.py signed.zip
   Flash image and boot up.
   Verify the oem images and vbmeta images in OUT and target zips by
   avbtool.

   b) Test on dist w/ chained vbmeta_system and oem and test custom images
   sign_target_files_apks -d certs \
     --avb_extra_custom_image_key oem=oem_rsa4096.pem \
     --avb_extra_custom_image_algorithm oem=SHA256_RSA4096 \
     --avb_extra_custom_image_extra_args oem=--do_not_generate_fec \
     --avb_extra_custom_image_key test=test_rsa4096.pem \
     --avb_extra_custom_image_algorithm test=SHA256_RSA4096 \
     xxx-target_xxx.zip signed.zip
   validate_target_files.py signed.zip
   Verify the oem, test images and vbmeta images in OUT and target zips
   by avbtool.

   c) Test on dist w/o chained partition.
   sign_target_files_apks -d certs xxx-target_xxx.zip signed.zip
   validate_target_files.py signed.zip
   Flash image and boot up.
   Verify the vbmeta images in OUT and target zips by avbtool.

Change-Id: Ifccfee5e8909697eef6ccda0cc352fa16a9f6db6

core/Makefile

@@ -2829,6 +2829,41 @@ define images-for-partitions
 $(strip $(foreach item,$(1),$(INSTALLED_$(call to-upper,$(item))IMAGE_TARGET)))
 endef
 
+# -----------------------------------------------------------------
+# custom images
+INSTALLED_CUSTOMIMAGES_TARGET :=
+
+ifneq ($(strip $(BOARD_CUSTOMIMAGES_PARTITION_LIST)),)
+INTERNAL_AVB_CUSTOMIMAGES_SIGNING_ARGS :=
+
+# Sign custom image.
+# $(1): the prebuilt custom image.
+# $(2): the mount point of the prebuilt custom image.
+# $(3): the signed custom image target.
+define sign_custom_image
+$(3): $(1) $(INTERNAL_USERIMAGES_DEPS)
+	@echo Target custom image: $(3)
+	mkdir -p $(dir $(3))
+	cp $(1) $(3)
+ifeq ($(BOARD_AVB_ENABLE),true)
+	PATH=$(INTERNAL_USERIMAGES_BINARY_PATHS):$$$$PATH \
+          $(AVBTOOL) add_hashtree_footer \
+          --image $(3) \
+          --key $(BOARD_AVB_$(call to-upper,$(2))_KEY_PATH) \
+          --algorithm $(BOARD_AVB_$(call to-upper,$(2))_ALGORITHM) \
+          --partition_size $(BOARD_AVB_$(call to-upper,$(2))_PARTITION_SIZE) \
+          --partition_name $(2) \
+          $(INTERNAL_AVB_CUSTOMIMAGES_SIGNING_ARGS) \
+          $(BOARD_AVB_$(call to-upper,$(2))_ADD_HASHTREE_FOOTER_ARGS)
+endif
+INSTALLED_CUSTOMIMAGES_TARGET += $(3)
+endef
+
+$(foreach partition,$(BOARD_CUSTOMIMAGES_PARTITION_LIST), \
+  $(foreach image,$(BOARD_AVB_$(call to-upper,$(partition))_IMAGE_LIST), \
+     $(eval $(call sign_custom_image,$(image),$(partition),$(PRODUCT_OUT)/$(notdir $(image))))))
+endif
+
 # -----------------------------------------------------------------
 # vbmeta image
 ifeq ($(BOARD_AVB_ENABLE),true)
@@ -3000,6 +3035,18 @@ $(if $(BOARD_AVB_$(call to-upper,$(1))_KEY_PATH),\
             --include_descriptors_from_image $(call images-for-partitions,$(1)))))
 endef
 
+# Checks and sets build variables for a custom chained partition to include it into vbmeta.img.
+# $(1): the custom partition to enable AVB chain.
+define check-and-set-custom-avb-chain-args
+$(eval part := $(1))
+$(eval PART=$(call to-upper,$(part)))
+$(eval _rollback_index_location := BOARD_AVB_$(PART)_ROLLBACK_INDEX_LOCATION)
+$(if $($(_rollback_index_location)),,$(error $(_rollback_index_location) is not defined))
+
+INTERNAL_AVB_MAKE_VBMETA_IMAGE_ARGS += \
+    --chain_partition $(part):$($(_rollback_index_location)):$(AVB_CHAIN_KEY_DIR)/$(part).avbpubkey
+endef
+
 ifdef INSTALLED_BOOTIMAGE_TARGET
 $(eval $(call check-and-set-avb-args,boot))
 endif
@@ -3043,6 +3090,11 @@ ifdef BOARD_AVB_VBMETA_VENDOR
 $(eval $(call check-and-set-avb-args,vbmeta_vendor))
 endif
 
+ifneq ($(strip $(BOARD_CUSTOMIMAGES_PARTITION_LIST)),)
+$(foreach partition,$(BOARD_CUSTOMIMAGES_PARTITION_LIST), \
+    $(eval $(call check-and-set-custom-avb-chain-args,$(partition))))
+endif
+
 # Add kernel cmdline descriptor for kernel to mount system.img as root with
 # dm-verity. This works when system.img is either chained or not-chained:
 # - chained: The --setup_as_rootfs_from_kernel option will add dm-verity kernel
@@ -3113,6 +3165,10 @@ define extract-avb-chain-public-keys
   $(if $(BOARD_AVB_VBMETA_VENDOR_KEY_PATH),\
     $(hide) $(AVBTOOL) extract_public_key --key $(BOARD_AVB_VBMETA_VENDOR_KEY_PATH) \
         --output $(1)/vbmeta_vendor.avbpubkey)
+  $(if $(BOARD_CUSTOMIMAGES_PARTITION_LIST),\
+    $(hide) $(foreach partition,$(BOARD_CUSTOMIMAGES_PARTITION_LIST), \
+        $(AVBTOOL) extract_public_key --key $(BOARD_AVB_$(call to-upper,$(partition))_KEY_PATH) \
+            --output $(1)/$(partition).avbpubkey;))
 endef
 
 # Builds a chained VBMeta image. This VBMeta image will contain the descriptors for the partitions
@@ -3180,6 +3236,7 @@ $(INSTALLED_VBMETAIMAGE_TARGET): \
 	    $(INSTALLED_SYSTEM_EXTIMAGE_TARGET) \
 	    $(INSTALLED_ODMIMAGE_TARGET) \
 	    $(INSTALLED_DTBOIMAGE_TARGET) \
+	    $(INSTALLED_CUSTOMIMAGES_TARGET) \
 	    $(INSTALLED_RECOVERYIMAGE_TARGET) \
 	    $(INSTALLED_VBMETA_SYSTEMIMAGE_TARGET) \
 	    $(INSTALLED_VBMETA_VENDORIMAGE_TARGET) \
@@ -3739,6 +3796,16 @@ ifdef BOARD_AVB_RECOVERY_KEY_PATH
 	$(hide) echo "avb_recovery_algorithm=$(BOARD_AVB_RECOVERY_ALGORITHM)" >> $@
 	$(hide) echo "avb_recovery_rollback_index_location=$(BOARD_AVB_RECOVERY_ROLLBACK_INDEX_LOCATION)" >> $@
 endif # BOARD_AVB_RECOVERY_KEY_PATH
+ifneq (,$(strip $(BOARD_CUSTOMIMAGES_PARTITION_LIST)))
+	$(hide) echo "avb_custom_images_partition_list=$(BOARD_CUSTOMIMAGES_PARTITION_LIST)" >> $@
+	$(hide) $(foreach partition,$(BOARD_CUSTOMIMAGES_PARTITION_LIST), \
+	    echo "avb_$(partition)_key_path=$(BOARD_AVB_$(call to-upper,$(partition))_KEY_PATH)"  >> $@; \
+	    echo "avb_$(partition)_algorithm=$(BOARD_AVB_$(call to-upper,$(partition))_ALGORITHM)"  >> $@; \
+	    echo "avb_$(partition)_add_hashtree_footer_args=$(BOARD_AVB_$(call to-upper,$(partition))_ADD_HASHTREE_FOOTER_ARGS)"  >> $@; \
+	    echo "avb_$(partition)_rollback_index_location=$(BOARD_AVB_$(call to-upper,$(partition))_ROLLBACK_INDEX_LOCATION)"  >> $@; \
+	    echo "avb_$(partition)_partition_size=$(BOARD_AVB_$(call to-upper,$(partition))_PARTITION_SIZE)"  >> $@; \
+	    echo "avb_$(partition)_image_list=$(foreach image,$(BOARD_AVB_$(call to-upper,$(partition))_IMAGE_LIST),$(notdir $(image)))" >> $@;)
+endif # BOARD_CUSTOMIMAGES_PARTITION_LIST
 ifneq (,$(strip $(BOARD_AVB_VBMETA_SYSTEM)))
 	$(hide) echo "avb_vbmeta_system=$(BOARD_AVB_VBMETA_SYSTEM)" >> $@
 	$(hide) echo "avb_vbmeta_system_args=$(BOARD_AVB_MAKE_VBMETA_SYSTEM_IMAGE_ARGS)" >> $@
@@ -3972,6 +4039,7 @@ $(BUILT_TARGET_FILES_PACKAGE): \
 	    $(INSTALLED_VBMETAIMAGE_TARGET) \
 	    $(INSTALLED_ODMIMAGE_TARGET) \
 	    $(INSTALLED_DTBOIMAGE_TARGET) \
+	    $(INSTALLED_CUSTOMIMAGES_TARGET) \
 	    $(INTERNAL_SYSTEMOTHERIMAGE_FILES) \
 	    $(INSTALLED_ANDROID_INFO_TXT_TARGET) \
 	    $(INSTALLED_KERNEL_TARGET) \
@@ -4216,6 +4284,11 @@ ifdef BOARD_PREBUILT_DTBOIMAGE
 	$(hide) mkdir -p $(zip_root)/PREBUILT_IMAGES
 	$(hide) cp $(INSTALLED_DTBOIMAGE_TARGET) $(zip_root)/PREBUILT_IMAGES/
 endif # BOARD_PREBUILT_DTBOIMAGE
+ifneq ($(strip $(BOARD_CUSTOMIMAGES_PARTITION_LIST)),)
+	$(hide) mkdir -p $(zip_root)/PREBUILT_IMAGES
+	$(hide) $(foreach partition,$(BOARD_CUSTOMIMAGES_PARTITION_LIST), \
+	    $(foreach image,$(BOARD_AVB_$(call to-upper,$(partition))_IMAGE_LIST),cp $(image) $(zip_root)/PREBUILT_IMAGES/;))
+endif # BOARD_CUSTOMIMAGES_PARTITION_LIST
 	@# The radio images in BOARD_PACK_RADIOIMAGES will be additionally copied from RADIO/ into
 	@# IMAGES/, which then will be added into <product>-img.zip. Such images must be listed in
 	@# INSTALLED_RADIOIMAGE_TARGET.

tools/releasetools/add_img_to_target_files.py

@@ -60,6 +60,7 @@ import build_super_image
 import common
 import rangelib
 import sparse_img
+import verity_utils
 
 if sys.hexversion < 0x02070000:
   print("Python 2.7 or newer is required.", file=sys.stderr)
@@ -312,6 +313,56 @@ def AddDtbo(output_zip):
   img.Write()
   return img.name
 
+def AddCustomImages(output_zip, partition_name):
+  """Adds and signs custom images in IMAGES/.
+
+  Args:
+    output_zip: The output zip file (needs to be already open), or None to
+        write images to OPTIONS.input_tmp/.
+
+  Uses the image under IMAGES/ if it already exists. Otherwise looks for the
+  image under PREBUILT_IMAGES/, signs it as needed, and returns the image name.
+
+  Raises:
+    AssertionError: If image can't be found.
+  """
+
+  partition_size = OPTIONS.info_dict.get(
+      "avb_{}_partition_size".format(partition_name))
+  key_path = OPTIONS.info_dict.get("avb_{}_key_path".format(partition_name))
+  algorithm = OPTIONS.info_dict.get("avb_{}_algorithm".format(partition_name))
+  extra_args = OPTIONS.info_dict.get(
+      "avb_{}_add_hashtree_footer_args".format(partition_name))
+  partition_size = OPTIONS.info_dict.get(
+      "avb_{}_partition_size".format(partition_name))
+
+  builder = verity_utils.CreateCustomImageBuilder(
+      OPTIONS.info_dict, partition_name, partition_size,
+      key_path, algorithm, extra_args)
+
+  for img_name in OPTIONS.info_dict.get(
+      "avb_{}_image_list".format(partition_name)).split():
+    custom_image = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", img_name)
+    if os.path.exists(custom_image.name):
+      continue
+
+    custom_image_prebuilt_path = os.path.join(
+        OPTIONS.input_tmp, "PREBUILT_IMAGES", img_name)
+    assert os.path.exists(custom_image_prebuilt_path), \
+      "Failed to find %s at %s" % (img_name, custom_image_prebuilt_path)
+
+    shutil.copy(custom_image_prebuilt_path, custom_image.name)
+
+    if builder is not None:
+      builder.Build(custom_image.name)
+
+    custom_image.Write()
+
+  default = os.path.join(OPTIONS.input_tmp, "IMAGES", partition_name + ".img")
+  assert os.path.exists(default), \
+      "There should be one %s.img" % (partition_name)
+  return default
+
 
 def CreateImage(input_dir, info_dict, what, output_file, block_list=None):
   logger.info("creating %s.img...", what)
@@ -411,8 +462,9 @@ def AddVBMeta(output_zip, partitions, name, needed_partitions):
   Args:
     output_zip: The output zip file, which needs to be already open.
     partitions: A dict that's keyed by partition names with image paths as
-        values. Only valid partition names are accepted, as listed in
+        values. Only valid partition names are accepted, as partitions listed
-        common.AVB_PARTITIONS.
+        in common.AVB_PARTITIONS and custom partitions listed in
+        OPTIONS.info_dict.get("avb_custom_images_partition_list")
     name: Name of the VBMeta partition, e.g. 'vbmeta', 'vbmeta_system'.
     needed_partitions: Partitions whose descriptors should be included into the
         generated VBMeta image.
@@ -831,11 +883,20 @@ def AddImagesToTargetFiles(filename):
     banner("dtbo")
     partitions['dtbo'] = AddDtbo(output_zip)
 
+  # Custom images.
+  custom_partitions = OPTIONS.info_dict.get(
+      "avb_custom_images_partition_list", "").strip().split()
+  for partition_name in custom_partitions:
+    partition_name = partition_name.strip()
+    banner("custom images for " + partition_name)
+    partitions[partition_name] = AddCustomImages(output_zip, partition_name)
+
   if OPTIONS.info_dict.get("avb_enable") == "true":
     # vbmeta_partitions includes the partitions that should be included into
     # top-level vbmeta.img, which are the ones that are not included in any
     # chained VBMeta image plus the chained VBMeta images themselves.
-    vbmeta_partitions = common.AVB_PARTITIONS[:]
+    # Currently custom_partitions are all chained to VBMeta image.
+    vbmeta_partitions = common.AVB_PARTITIONS[:] + tuple(custom_partitions)
 
     vbmeta_system = OPTIONS.info_dict.get("avb_vbmeta_system", "").strip()
     if vbmeta_system:

tools/releasetools/common.py

@@ -1132,8 +1132,9 @@ def BuildVBMeta(image_path, partitions, name, needed_partitions):
   Args:
     image_path: The output path for the new VBMeta image.
     partitions: A dict that's keyed by partition names with image paths as
-        values. Only valid partition names are accepted, as listed in
+        values. Only valid partition names are accepted, as partitions listed
-        common.AVB_PARTITIONS.
+        in common.AVB_PARTITIONS and custom partitions listed in
+        OPTIONS.info_dict.get("avb_custom_images_partition_list")
     name: Name of the VBMeta partition, e.g. 'vbmeta', 'vbmeta_system'.
     needed_partitions: Partitions whose descriptors should be included into the
         generated VBMeta image.
@@ -1145,11 +1146,15 @@ def BuildVBMeta(image_path, partitions, name, needed_partitions):
   cmd = [avbtool, "make_vbmeta_image", "--output", image_path]
   AppendAVBSigningArgs(cmd, name)
 
+  custom_partitions = OPTIONS.info_dict.get(
+      "avb_custom_images_partition_list", "").strip().split()
+
   for partition, path in partitions.items():
     if partition not in needed_partitions:
       continue
     assert (partition in AVB_PARTITIONS or
-            partition in AVB_VBMETA_PARTITIONS), \
+            partition in AVB_VBMETA_PARTITIONS or
+            partition in custom_partitions), \
         'Unknown partition: {}'.format(partition)
     assert os.path.exists(path), \
         'Failed to find {} for {}'.format(path, partition)

tools/releasetools/sign_target_files_apks.py

@@ -112,6 +112,17 @@ Usage:  sign_target_files_apks [flags] input_target_files output_target_files
       (e.g. "--signing_helper /path/to/helper"). The args will be appended to
       the existing ones in info dict.
 
+  --avb_extra_custom_image_key <partition=key>
+  --avb_extra_custom_image_algorithm <partition=algorithm>
+      Use the specified algorithm (e.g. SHA256_RSA4096) and the key to AVB-sign
+      the specified custom images mounted on the partition. Otherwise it uses
+      the existing values in info dict.
+
+  --avb_extra_custom_image_extra_args <partition=extra_args>
+      Specify any additional args that are needed to AVB-sign the custom images
+      mounted on the partition (e.g. "--signing_helper /path/to/helper"). The
+      args will be appended to the existing ones in info dict.
+
   --android_jar_path <path>
       Path to the android.jar to repack the apex file.
 """
@@ -956,12 +967,20 @@ def ReplaceAvbSigningKeys(misc_info):
     if extra_args:
       print('Setting extra AVB signing args for %s to "%s"' % (
           partition, extra_args))
-      args_key = AVB_FOOTER_ARGS_BY_PARTITION[partition]
+      if partition in AVB_FOOTER_ARGS_BY_PARTITION:
+        args_key = AVB_FOOTER_ARGS_BY_PARTITION[partition]
+      else:
+        # custom partition
+        args_key = "avb_{}_add_hashtree_footer_args".format(partition)
       misc_info[args_key] = (misc_info.get(args_key, '') + ' ' + extra_args)
 
   for partition in AVB_FOOTER_ARGS_BY_PARTITION:
     ReplaceAvbPartitionSigningKey(partition)
 
+  for custom_partition in misc_info.get(
+      "avb_custom_images_partition_list", "").strip().split():
+    ReplaceAvbPartitionSigningKey(custom_partition)
+
 
 def RewriteAvbProps(misc_info):
   """Rewrites the props in AVB signing args."""
@@ -1208,6 +1227,15 @@ def main(argv):
       OPTIONS.avb_extra_args['vbmeta_vendor'] = a
     elif o == "--avb_apex_extra_args":
       OPTIONS.avb_extra_args['apex'] = a
+    elif o == "--avb_extra_custom_image_key":
+      partition, key = a.split("=")
+      OPTIONS.avb_keys[partition] = key
+    elif o == "--avb_extra_custom_image_algorithm":
+      partition, algorithm = a.split("=")
+      OPTIONS.avb_algorithms[partition] = algorithm
+    elif o == "--avb_extra_custom_image_extra_args":
+      partition, extra_args = a.split("=")
+      OPTIONS.avb_extra_args[partition] = extra_args
     else:
       return False
     return True
@@ -1252,6 +1280,9 @@ def main(argv):
           "avb_vbmeta_vendor_algorithm=",
           "avb_vbmeta_vendor_key=",
           "avb_vbmeta_vendor_extra_args=",
+          "avb_extra_custom_image_key=",
+          "avb_extra_custom_image_algorithm=",
+          "avb_extra_custom_image_extra_args=",
       ],
       extra_option_handler=option_handler)
 

tools/releasetools/validate_target_files.py

@@ -352,8 +352,13 @@ def ValidateVerifiedBootImages(input_tmp, info_dict, options):
     cmd = [info_dict['avb_avbtool'], 'verify_image', '--image', image,
            '--follow_chain_partitions']
 
+    # Custom images.
+    custom_partitions = info_dict.get(
+        "avb_custom_images_partition_list", "").strip().split()
+
     # Append the args for chained partitions if any.
-    for partition in common.AVB_PARTITIONS + common.AVB_VBMETA_PARTITIONS:
+    for partition in (common.AVB_PARTITIONS + common.AVB_VBMETA_PARTITIONS +
+                      tuple(custom_partitions)):
       key_name = 'avb_' + partition + '_key_path'
       if info_dict.get(key_name) is not None:
         if info_dict.get('ab_update') != 'true' and partition == 'recovery':

tools/releasetools/verity_utils.py

@@ -695,3 +695,22 @@ class VerifiedBootVersion1HashtreeInfoGenerator(HashtreeInfoGenerator):
       raise HashtreeInfoGenerationError("Failed to reconstruct the verity tree")
 
     return self.hashtree_info
+
+
+def CreateCustomImageBuilder(info_dict, partition_name, partition_size,
+                            key_path, algorithm, signing_args):
+  builder = None
+  if info_dict.get("avb_enable") == "true":
+    builder = VerifiedBootVersion2VerityImageBuilder(
+        partition_name,
+        partition_size,
+        VerifiedBootVersion2VerityImageBuilder.AVB_HASHTREE_FOOTER,
+        info_dict.get("avb_avbtool"),
+        key_path,
+        algorithm,
+        # Salt is None because custom images have no fingerprint property to be
+        # used as the salt.
+        None,
+        signing_args)
+
+  return builder