Support signing custom images with AVB HASH or AVB HASHTREE

`make custom_images` supports to build different kinds of *non-droid* images,
e.g., odm.img. Adding the support of signing them with either AVB HASH footer
or AVB HASHTREE footer. The user can use HASH for small images and
HASHTREE for large images.

Sample signing configurations:
 * AVB HASH footer:
   - CUSTOM_IMAGE_AVB_HASH_ENABLE := true
   - CUSTOM_IMAGE_AVB_ADD_HASH_FOOTER_ARGS := --append_to_release_string my_odm_image

 * AVB HASHTREE footer:
   - CUSTOM_IMAGE_AVB_HASHTREE_ENABLE := true
   - CUSTOM_IMAGE_AVB_ADD_HASHTREE_FOOTER_ARGS := --fec_num_roots 8

 * Using custom signing key:
   - CUSTOM_IMAGE_AVB_ALGORITHM := SHA256_RSA2048
   - CUSTOM_IMAGE_AVB_KEY_PATH := external/avb/test/data/testkey_rsa2048.pem

Bug: 36701014
Test: `make custom_images` with AVB HASH footer
Test: `make custom_images` with AVB HASHTREE footer
Test: `make droid` to check system.img is still properly signed with AVB HASHTREE
Test: `make droid` to check vendor.img is still properly signed with AVB HASHTREE

Change-Id: I8dc420e12e37e9a631345c0cd883339db05d489f
diff --git a/core/Makefile b/core/Makefile
index 0b63915..fd65013 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -938,9 +938,9 @@
 $(if $(PRODUCTS.$(INTERNAL_PRODUCT).PRODUCT_SUPPORTS_VBOOT),$(hide) echo "vboot_signer_cmd=$(VBOOT_SIGNER)" >> $(1))
 $(if $(BOARD_AVB_ENABLE),$(hide) echo "avb_signing_args=$(INTERNAL_AVB_SIGNING_ARGS)" >> $(1))
 $(if $(BOARD_AVB_ENABLE),$(hide) echo "avb_avbtool=$(AVBTOOL)" >> $(1))
-$(if $(BOARD_AVB_ENABLE),$(hide) echo "system_avb_enable=$(BOARD_AVB_ENABLE)" >> $(1))
+$(if $(BOARD_AVB_ENABLE),$(hide) echo "system_avb_hashtree_enable=$(BOARD_AVB_ENABLE)" >> $(1))
 $(if $(BOARD_AVB_ENABLE),$(hide) echo "system_avb_add_hashtree_footer_args=$(BOARD_AVB_SYSTEM_ADD_HASHTREE_FOOTER_ARGS)" >> $(1))
-$(if $(BOARD_AVB_ENABLE),$(hide) echo "vendor_avb_enable=$(BOARD_AVB_ENABLE)" >> $(1))
+$(if $(BOARD_AVB_ENABLE),$(hide) echo "vendor_avb_hashtree_enable=$(BOARD_AVB_ENABLE)" >> $(1))
 $(if $(BOARD_AVB_ENABLE),$(hide) echo "vendor_avb_add_hashtree_footer_args=$(BOARD_AVB_VENDOR_ADD_HASHTREE_FOOTER_ARGS)" >> $(1))
 $(if $(filter true,$(BOARD_USES_RECOVERY_AS_BOOT)),\
     $(hide) echo "recovery_as_boot=true" >> $(1))
diff --git a/core/tasks/build_custom_images.mk b/core/tasks/build_custom_images.mk
index c5f2a96..0750217 100644
--- a/core/tasks/build_custom_images.mk
+++ b/core/tasks/build_custom_images.mk
@@ -37,6 +37,14 @@
 #   - CUSTOM_IMAGE_SELINUX, set to "true" if the image supports selinux.
 #   - CUSTOM_IMAGE_SUPPORT_VERITY, set to "true" if the product supports verity.
 #   - CUSTOM_IMAGE_VERITY_BLOCK_DEVICE
+#   - CUSTOM_IMAGE_AVB_HASH_ENABLE, set to "true" to add AVB HASH footer.
+#   - CUSTOM_IMAGE_AVB_ADD_HASH_FOOTER_ARGS, additional args of AVB HASH footer.
+#   - CUSTOM_IMAGE_AVB_HASHTREE_ENABLE, set to "true" to add AVB HASHTREE
+#     footer.
+#   - CUSTOM_IMAGE_AVB_ADD_HASHTREE_FOOTER_ARGS, additional args of AVB
+#     HASHTREE footer.
+#   - CUSTOM_IMAGE_AVB_KEY_PATH, custom AVB signing key.
+#   - CUSTOM_IMAGE_AVB_ALGORITHM, custom AVB signing algorithm.
 #
 # To build all those images, run "make custom_images".
 
@@ -54,8 +62,12 @@
   CUSTOM_IMAGE_SELINUX \
   CUSTOM_IMAGE_SUPPORT_VERITY \
   CUSTOM_IMAGE_VERITY_BLOCK_DEVICE \
-  CUSTOM_IMAGE_AVB_ENABLE \
+  CUSTOM_IMAGE_AVB_HASH_ENABLE \
+  CUSTOM_IMAGE_AVB_ADD_HASH_FOOTER_ARGS \
+  CUSTOM_IMAGE_AVB_HASHTREE_ENABLE \
   CUSTOM_IMAGE_AVB_ADD_HASHTREE_FOOTER_ARGS \
+  CUSTOM_IMAGE_AVB_KEY_PATH \
+  CUSTOM_IMAGE_AVB_ALGORITHM \
 
 # We don't expect product makefile to inherit/override PRODUCT_CUSTOM_IMAGE_MAKEFILES,
 # so we don't put it in the _product_var_list.
diff --git a/core/tasks/tools/build_custom_image.mk b/core/tasks/tools/build_custom_image.mk
index f916e86..8c098d6 100644
--- a/core/tasks/tools/build_custom_image.mk
+++ b/core/tasks/tools/build_custom_image.mk
@@ -62,6 +62,14 @@
   $(eval my_image_copy_files += $(src))\
   $(eval my_copy_pairs += $(src):$(my_staging_dir)/$(word 2,$(pair))))
 
+ifndef CUSTOM_IMAGE_AVB_KEY_PATH
+# If key path isn't specified, use the default signing args.
+my_avb_signing_args := $(INTERNAL_AVB_SIGNING_ARGS)
+else
+my_avb_signing_args := \
+  --algorithm $(CUSTOM_IMAGE_AVB_ALGORITHM) --key $(CUSTOM_IMAGE_AVB_KEY_PATH)
+endif
+
 $(my_built_custom_image): PRIVATE_INTERMEDIATES := $(intermediates)
 $(my_built_custom_image): PRIVATE_MOUNT_POINT := $(CUSTOM_IMAGE_MOUNT_POINT)
 $(my_built_custom_image): PRIVATE_PARTITION_SIZE := $(CUSTOM_IMAGE_PARTITION_SIZE)
@@ -74,10 +82,16 @@
 $(my_built_custom_image): PRIVATE_VERITY_KEY := $(PRODUCTS.$(INTERNAL_PRODUCT).PRODUCT_VERITY_SIGNING_KEY)
 $(my_built_custom_image): PRIVATE_VERITY_BLOCK_DEVICE := $(CUSTOM_IMAGE_VERITY_BLOCK_DEVICE)
 $(my_built_custom_image): PRIVATE_DICT_FILE := $(CUSTOM_IMAGE_DICT_FILE)
-$(my_built_custom_image): PRIVATE_AVB_ENABLE := $(CUSTOM_IMAGE_AVB_ENABLE)
+$(my_built_custom_image): PRIVATE_AVB_AVBTOOL := $(AVBTOOL)
+$(my_built_custom_image): PRIVATE_AVB_SIGNING_ARGS := $(my_avb_signing_args)
+$(my_built_custom_image): PRIVATE_AVB_HASH_ENABLE := $(CUSTOM_IMAGE_AVB_HASH_ENABLE)
+$(my_built_custom_image): PRIVATE_AVB_ADD_HASH_FOOTER_ARGS := $(CUSTOM_IMAGE_AVB_ADD_HASH_FOOTER_ARGS)
+$(my_built_custom_image): PRIVATE_AVB_HASHTREE_ENABLE := $(CUSTOM_IMAGE_AVB_HASHTREE_ENABLE)
 $(my_built_custom_image): PRIVATE_AVB_ADD_HASHTREE_FOOTER_ARGS := $(CUSTOM_IMAGE_AVB_ADD_HASHTREE_FOOTER_ARGS)
-ifeq (true,$(CUSTOM_IMAGE_AVB_ENABLE))
+ifeq (true,$(filter true, $(CUSTOM_IMAGE_AVB_HASH_ENABLE) $(CUSTOM_IMAGE_AVB_HASHTREE_ENABLE)))
   $(my_built_custom_image): $(AVBTOOL)
+else ifneq (,$(filter true, $(CUSTOM_IMAGE_AVB_HASH_ENABLE) $(CUSTOM_IMAGE_AVB_HASHTREE_ENABLE)))
+  $(error Cannot set both CUSTOM_IMAGE_AVB_HASH_ENABLE and CUSTOM_IMAGE_AVB_HASHTREE_ENABLE to true)
 endif
 $(my_built_custom_image): $(INTERNAL_USERIMAGES_DEPS) $(my_built_modules) $(my_image_copy_files) \
   $(CUSTOM_IMAGE_DICT_FILE)
@@ -93,6 +107,7 @@
 	# Generate the dict.
 	$(hide) echo "# For all accepted properties, see BuildImage() in tools/releasetools/build_image.py" > $(PRIVATE_INTERMEDIATES)/image_info.txt
 	$(hide) echo "mount_point=$(PRIVATE_MOUNT_POINT)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt
+	$(hide) echo "partition_name=$(PRIVATE_MOUNT_POINT)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt
 	$(hide) echo "fs_type=$(PRIVATE_FILE_SYSTEM_TYPE)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt
 	$(hide) echo "partition_size=$(PRIVATE_PARTITION_SIZE)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt
 	$(hide) echo "ext_mkuserimg=$(notdir $(MKEXTUSERIMG))" >> $(PRIVATE_INTERMEDIATES)/image_info.txt
@@ -102,10 +117,13 @@
 	    echo "verity_key=$(PRIVATE_VERITY_KEY)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
 	    echo "verity_signer_cmd=$(VERITY_SIGNER)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
 	    echo "verity_block_device=$(PRIVATE_VERITY_BLOCK_DEVICE)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt)
-	$(if $(PRIVATE_AVB_ENABLE),\
-	  $(hide) echo "avb_enable=$(PRIVATE_AVB_ENABLE)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
-	    echo "avb_avbtool=$(AVBTOOL)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
-	    echo "avb_signing_args=$(INTERNAL_AVB_SIGNING_ARGS)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
+	$(hide) echo "avb_avbtool=$(PRIVATE_AVB_AVBTOOL)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt
+	$(hide) echo "avb_signing_args=$(PRIVATE_AVB_SIGNING_ARGS)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt
+	$(if $(PRIVATE_AVB_HASH_ENABLE),\
+	  $(hide) echo "avb_hash_enable=$(PRIVATE_AVB_HASH_ENABLE)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
+	    echo "avb_add_hash_footer_args=$(PRIVATE_AVB_ADD_HASH_FOOTER_ARGS)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt)
+	$(if $(PRIVATE_AVB_HASHTREE_ENABLE),\
+	  $(hide) echo "avb_hashtree_enable=$(PRIVATE_AVB_HASHTREE_ENABLE)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
 	    echo "avb_add_hashtree_footer_args=$(PRIVATE_AVB_ADD_HASHTREE_FOOTER_ARGS)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt)
 	$(if $(PRIVATE_DICT_FILE),\
 	  $(hide) echo "# Properties from $(PRIVATE_DICT_FILE)" >> $(PRIVATE_INTERMEDIATES)/image_info.txt;\
diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 3094dca..05c895a 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -97,18 +97,19 @@
   simg = sparse_img.SparseImage(image_file, mode="r+b", build_map=False)
   simg.AppendFillChunk(0, blocks)
 
-def AVBCalcMaxImageSize(avbtool, partition_size, additional_args):
+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_hashtree_image'.
   Returns:
     The maximum image size or 0 if an error occurred.
   """
-  cmdline = "%s add_hashtree_footer " % avbtool
+  cmdline = "%s add_%s_footer " % (avbtool, footer_type)
   cmdline += "--partition_size %d " % partition_size
   cmdline += "--calc_max_image_size "
   cmdline += additional_args
@@ -118,13 +119,14 @@
   else:
     return int(output)
 
-def AVBAddHashtree(image_path, avbtool, partition_size, partition_name,
-                   signing_args, additional_args):
+def AVBAddFooter(image_path, avbtool, footer_type, partition_size,
+                 partition_name, signing_args, 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.
     signing_args: Arguments for signing the image.
@@ -133,7 +135,7 @@
   Returns:
     True if the operation succeeded.
   """
-  cmdline = "%s add_hashtree_footer " % avbtool
+  cmdline = "%s add_%s_footer " % (avbtool, footer_type)
   cmdline += "--partition_size %d " % partition_size
   cmdline += "--partition_name %s " % partition_name
   cmdline += "--image %s " % image_path
@@ -410,12 +412,19 @@
     prop_dict["original_partition_size"] = str(partition_size)
     prop_dict["verity_size"] = str(verity_size)
 
-  # Adjust partition size for AVB.
-  if prop_dict.get("avb_enable") == "true":
+  # Adjust partition size for AVB hash footer or AVB hashtree footer.
+  avb_footer_type = ''
+  if prop_dict.get("avb_hash_enable") == "true":
+    avb_footer_type = 'hash'
+  elif prop_dict.get("avb_hashtree_enable") == "true":
+    avb_footer_type = 'hashtree'
+
+  if avb_footer_type:
     avbtool = prop_dict.get("avb_avbtool")
     partition_size = int(prop_dict.get("partition_size"))
-    additional_args = prop_dict["avb_add_hashtree_footer_args"]
-    max_image_size = AVBCalcMaxImageSize(avbtool, partition_size,
+    # avb_add_hash_footer_args or avb_add_hashtree_footer_args.
+    additional_args = prop_dict["avb_add_" + avb_footer_type + "_footer_args"]
+    max_image_size = AVBCalcMaxImageSize(avbtool, avb_footer_type, partition_size,
                                          additional_args)
     if max_image_size == 0:
       return False
@@ -561,15 +570,16 @@
     if not MakeVerityEnabledImage(out_file, verity_fec_supported, prop_dict):
       return False
 
-  # Add AVB hashtree and metadata.
-  if "avb_enable" in prop_dict:
+  # Add AVB HASH or HASHTREE footer (metadata).
+  if avb_footer_type:
     avbtool = prop_dict.get("avb_avbtool")
     original_partition_size = int(prop_dict.get("original_partition_size"))
     partition_name = prop_dict["partition_name"]
     signing_args = prop_dict["avb_signing_args"]
-    additional_args = prop_dict["avb_add_hashtree_footer_args"]
-    if not AVBAddHashtree(out_file, avbtool, original_partition_size,
-                          partition_name, signing_args, additional_args):
+    # avb_add_hash_footer_args or avb_add_hashtree_footer_args
+    additional_args = prop_dict["avb_add_" + avb_footer_type + "_footer_args"]
+    if not AVBAddFooter(out_file, avbtool, avb_footer_type, original_partition_size,
+                      partition_name, signing_args, additional_args):
       return False
 
   if run_fsck and prop_dict.get("skip_fsck") != "true":
@@ -639,7 +649,7 @@
     copy_prop("system_squashfs_block_size", "squashfs_block_size")
     copy_prop("system_squashfs_disable_4k_align", "squashfs_disable_4k_align")
     copy_prop("system_base_fs_file", "base_fs_file")
-    copy_prop("system_avb_enable", "avb_enable")
+    copy_prop("system_avb_hashtree_enable", "avb_hashtree_enable")
     copy_prop("system_avb_add_hashtree_footer_args",
               "avb_add_hashtree_footer_args")
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
@@ -656,7 +666,7 @@
     copy_prop("system_squashfs_compressor_opt", "squashfs_compressor_opt")
     copy_prop("system_squashfs_block_size", "squashfs_block_size")
     copy_prop("system_base_fs_file", "base_fs_file")
-    copy_prop("system_avb_enable", "avb_enable")
+    copy_prop("system_avb_hashtree_enable", "avb_hashtree_enable")
     copy_prop("system_avb_add_hashtree_footer_args",
               "avb_add_hashtree_footer_args")
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
@@ -681,7 +691,7 @@
     copy_prop("vendor_squashfs_block_size", "squashfs_block_size")
     copy_prop("vendor_squashfs_disable_4k_align", "squashfs_disable_4k_align")
     copy_prop("vendor_base_fs_file", "base_fs_file")
-    copy_prop("vendor_avb_enable", "avb_enable")
+    copy_prop("vendor_avb_hashtree_enable", "avb_hashtree_enable")
     copy_prop("vendor_avb_add_hashtree_footer_args",
               "avb_add_hashtree_footer_args")
     copy_prop("vendor_extfs_inode_count", "extfs_inode_count")