Support building chained VBMeta images.

Bug: 112007947
Test: `m dist` with aosp_taimen-userdebug.
Test: Set up a target that uses chained VBMeta images of
      `vbmeta_mainline` and `vbmeta_vendor`. `m dist` and check the
      build log, as well as outputs from
      `avbtool info_image --image vbmeta.img`,
      `avbtool info_image --image vbmeta_mainline.img`,
      `avbtool info_image --image vbmeta_vendor.img`.
Change-Id: Ib1d4e97f583b65245703eae15d211adcd9e83741
diff --git a/core/Makefile b/core/Makefile
index a9c7abe..6aac748 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -2645,6 +2645,15 @@
 BOARD_AVB_KEY_PATH := external/avb/test/data/testkey_rsa4096.pem
 endif
 
+INTERNAL_AVB_PARTITIONS_IN_CHAINED_VBMETA_IMAGES := \
+    $(BOARD_AVB_VBMETA_MAINLINE) \
+    $(BOARD_AVB_VBMETA_VENDOR)
+
+# Not allowing the same partition to appear in multiple groups.
+ifneq ($(words $(sort $(INTERNAL_AVB_PARTITIONS_IN_CHAINED_VBMETA_IMAGES))),$(words $(INTERNAL_AVB_PARTITIONS_IN_CHAINED_VBMETA_IMAGES)))
+  $(error BOARD_AVB_VBMETA_MAINLINE and BOARD_AVB_VBMETA_VENDOR cannot have duplicates)
+endif
+
 BOOT_FOOTER_ARGS := BOARD_AVB_BOOT_ADD_HASH_FOOTER_ARGS
 DTBO_FOOTER_ARGS := BOARD_AVB_DTBO_ADD_HASH_FOOTER_ARGS
 SYSTEM_FOOTER_ARGS := BOARD_AVB_SYSTEM_ADD_HASHTREE_FOOTER_ARGS
@@ -2655,7 +2664,7 @@
 ODM_FOOTER_ARGS := BOARD_AVB_ODM_ADD_HASHTREE_FOOTER_ARGS
 
 # Helper function that checks and sets required build variables for an AVB chained partition.
-# $(1): the partition to enable AVB chain, e.g., boot or system.
+# $(1): the partition to enable AVB chain, e.g., boot or system or vbmeta_mainline.
 define _check-and-set-avb-chain-args
 $(eval part := $(1))
 $(eval PART=$(call to-upper,$(part)))
@@ -2677,20 +2686,27 @@
 $(eval INTERNAL_AVB_MAKE_VBMETA_IMAGE_ARGS += \
     --chain_partition $(part):$($(_rollback_index_location)):$(AVB_CHAIN_KEY_DIR)/$(part).avbpubkey)
 
-# Set rollback_index via footer args
-$(eval _footer_args := $(PART)_FOOTER_ARGS)
-$(eval $($(_footer_args)) += --rollback_index $($(_rollback_index)))
+# Set rollback_index via footer args for non-chained vbmeta image. Chained vbmeta image will pick up
+# the index via a separate flag (e.g. BOARD_AVB_VBMETA_MAINLINE_ROLLBACK_INDEX).
+$(if $(filter $(part),$(part:vbmeta_%=%)),\
+    $(eval _footer_args := $(PART)_FOOTER_ARGS) \
+    $(eval $($(_footer_args)) += --rollback_index $($(_rollback_index))))
 endef
 
 # Checks and sets the required build variables for an AVB partition. The partition will be
 # configured as a chained partition, if BOARD_AVB_<partition>_KEY_PATH is defined. Otherwise the
-# image descriptor will be included into vbmeta.img.
+# image descriptor will be included into vbmeta.img, unless it has been already added to any chained
+# VBMeta image.
 # $(1): Partition name, e.g. boot or system.
 define check-and-set-avb-args
+$(eval _in_chained_vbmeta := $(filter $(1),$(INTERNAL_AVB_PARTITIONS_IN_CHAINED_VBMETA_IMAGES)))
 $(if $(BOARD_AVB_$(call to-upper,$(1))_KEY_PATH),\
+    $(if $(_in_chained_vbmeta),\
+        $(error Chaining partition "$(1)" in chained VBMeta image is not supported)) \
     $(call _check-and-set-avb-chain-args,$(1)),\
-    $(eval INTERNAL_AVB_MAKE_VBMETA_IMAGE_ARGS += \
-        --include_descriptors_from_image $(call images-for-partitions,$(1))))
+    $(if $(_in_chained_vbmeta),,\
+        $(eval INTERNAL_AVB_MAKE_VBMETA_IMAGE_ARGS += \
+            --include_descriptors_from_image $(call images-for-partitions,$(1)))))
 endef
 
 ifdef INSTALLED_BOOTIMAGE_TARGET
@@ -2723,6 +2739,15 @@
 $(eval $(call check-and-set-avb-args,recovery))
 endif
 
+# Not using INSTALLED_VBMETA_MAINLINEIMAGE_TARGET as it won't be set yet.
+ifdef BOARD_AVB_VBMETA_MAINLINE
+$(eval $(call check-and-set-avb-args,vbmeta_mainline))
+endif
+
+ifdef BOARD_AVB_VBMETA_VENDOR
+$(eval $(call check-and-set-avb-args,vbmeta_vendor))
+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
@@ -2734,13 +2759,27 @@
 endif
 
 BOARD_AVB_MAKE_VBMETA_IMAGE_ARGS += --padding_size 4096
+BOARD_AVB_MAKE_VBMETA_MAINLINE_IMAGE_ARGS += --padding_size 4096
+BOARD_AVB_MAKE_VBMETA_VENDOR_IMAGE_ARGS += --padding_size 4096
+
+ifeq (eng,$(filter eng, $(TARGET_BUILD_VARIANT)))
+BOARD_AVB_MAKE_VBMETA_IMAGE_ARGS += --set_hashtree_disabled_flag
+BOARD_AVB_MAKE_VBMETA_MAINLINE_IMAGE_ARGS += --set_hashtree_disabled_flag
+BOARD_AVB_MAKE_VBMETA_VENDOR_IMAGE_ARGS += --set_hashtree_disabled_flag
+endif
 
 ifdef BOARD_AVB_ROLLBACK_INDEX
 BOARD_AVB_MAKE_VBMETA_IMAGE_ARGS += --rollback_index $(BOARD_AVB_ROLLBACK_INDEX)
 endif
 
-ifeq (eng,$(filter eng, $(TARGET_BUILD_VARIANT)))
-BOARD_AVB_MAKE_VBMETA_IMAGE_ARGS += --set_hashtree_disabled_flag
+ifdef BOARD_AVB_VBMETA_MAINLINE_ROLLBACK_INDEX
+BOARD_AVB_MAKE_VBMETA_MAINLINE_IMAGE_ARGS += \
+    --rollback_index $(BOARD_AVB_VBMETA_MAINLINE_ROLLBACK_INDEX)
+endif
+
+ifdef BOARD_AVB_VBMETA_VENDOR_ROLLBACK_INDEX
+BOARD_AVB_MAKE_VBMETA_VENDOR_IMAGE_ARGS += \
+    --rollback_index $(BOARD_AVB_VBMETA_VENDOR_ROLLBACK_INDEX)
 endif
 
 # $(1): the directory to extract public keys to
@@ -2769,8 +2808,50 @@
   $(if $(BOARD_AVB_RECOVERY_KEY_PATH),\
     $(hide) $(AVBTOOL) extract_public_key --key $(BOARD_AVB_RECOVERY_KEY_PATH) \
       --output $(1)/recovery.avbpubkey)
+  $(if $(BOARD_AVB_VBMETA_MAINLINE_KEY_PATH),\
+    $(hide) $(AVBTOOL) extract_public_key --key $(BOARD_AVB_VBMETA_MAINLINE_KEY_PATH) \
+        --output $(1)/vbmeta_mainline.avbpubkey)
+  $(if $(BOARD_AVB_VBMETA_VENDOR_KEY_PATH),\
+    $(hide) $(AVBTOOL) extract_public_key --key $(BOARD_AVB_VBMETA_VENDOR_KEY_PATH) \
+        --output $(1)/vbmeta_vendor.avbpubkey)
 endef
 
+# Builds a chained VBMeta image. This VBMeta image will contain the descriptors for the partitions
+# specified in BOARD_AVB_VBMETA_<NAME>. The built VBMeta image will be included into the top-level
+# vbmeta image as a chained partition. For example, if a target defines `BOARD_AVB_VBMETA_MAINLINE
+# := system product_services`, `vbmeta_mainline.img` will be created that includes the descriptors
+# for `system.img` and `product_services.img`. `vbmeta_mainline.img` itself will be included into
+# `vbmeta.img` as a chained partition.
+# $(1): VBMeta image name, such as "vbmeta_mainline", "vbmeta_vendor" etc.
+# $(2): Output filename.
+define build-chained-vbmeta-image
+  $(call pretty,"Target chained vbmeta image: $@")
+  $(hide) $(AVBTOOL) make_vbmeta_image \
+      $(INTERNAL_AVB_$(call to-upper,$(1))_SIGNING_ARGS) \
+      $(BOARD_AVB_MAKE_$(call to-upper,$(1))_IMAGE_ARGS) \
+      $(foreach image,$(BOARD_AVB_$(call to-upper,$(1))), \
+          --include_descriptors_from_image $(call images-for-partitions,$(image))) \
+      --output $@
+endef
+
+ifdef BOARD_AVB_VBMETA_MAINLINE
+INSTALLED_VBMETA_MAINLINEIMAGE_TARGET := $(PRODUCT_OUT)/vbmeta_mainline.img
+$(INSTALLED_VBMETA_MAINLINEIMAGE_TARGET): \
+		$(AVBTOOL) \
+		$(call images-for-partitions,$(BOARD_AVB_VBMETA_MAINLINE)) \
+		$(BOARD_AVB_VBMETA_MAINLINE_KEY_PATH)
+	$(call build-chained-vbmeta-image,vbmeta_mainline)
+endif
+
+ifdef BOARD_AVB_VBMETA_VENDOR
+INSTALLED_VBMETA_VENDORIMAGE_TARGET := $(PRODUCT_OUT)/vbmeta_vendor.img
+$(INSTALLED_VBMETA_VENDORIMAGE_TARGET): \
+		$(AVBTOOL) \
+		$(call images-for-partitions,$(BOARD_AVB_VBMETA_VENDOR)) \
+		$(BOARD_AVB_VBMETA_VENDOR_KEY_PATH)
+	$(call build-chained-vbmeta-image,vbmeta_vendor)
+endif
+
 define build-vbmetaimage-target
   $(call pretty,"Target vbmeta image: $(INSTALLED_VBMETAIMAGE_TARGET)")
   $(hide) mkdir -p $(AVB_CHAIN_KEY_DIR)
@@ -2797,6 +2878,10 @@
 		$(INSTALLED_ODMIMAGE_TARGET) \
 		$(INSTALLED_DTBOIMAGE_TARGET) \
 		$(INSTALLED_RECOVERYIMAGE_TARGET) \
+		$(INSTALLED_VBMETA_MAINLINEIMAGE_TARGET) \
+		$(INSTALLED_VBMETA_VENDORIMAGE_TARGET) \
+		$(BOARD_AVB_VBMETA_MAINLINE_KEY_PATH) \
+		$(BOARD_AVB_VBMETA_VENDOR_KEY_PATH) \
 		$(BOARD_AVB_KEY_PATH)
 	$(build-vbmetaimage-target)
 
@@ -3380,6 +3465,20 @@
 	$(hide) echo "avb_recovery_algorithm=$(BOARD_AVB_RECOVERY_ALGORITHM)" >> $(zip_root)/META/misc_info.txt
 	$(hide) echo "avb_recovery_rollback_index_location=$(BOARD_AVB_RECOVERY_ROLLBACK_INDEX_LOCATION)" >> $(zip_root)/META/misc_info.txt
 endif # BOARD_AVB_RECOVERY_KEY_PATH
+ifneq (,$(strip $(BOARD_AVB_VBMETA_MAINLINE)))
+	$(hide) echo "avb_vbmeta_mainline=$(BOARD_AVB_VBMETA_MAINLINE)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_mainline_args=$(BOARD_AVB_MAKE_VBMETA_MAINLINE_IMAGE_ARGS)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_mainline_key_path=$(BOARD_AVB_VBMETA_MAINLINE_KEY_PATH)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_mainline_algorithm=$(BOARD_AVB_VBMETA_MAINLINE_ALGORITHM)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_mainline_rollback_index_location=$(BOARD_AVB_VBMETA_MAINLINE_ROLLBACK_INDEX_LOCATION)" >> $(zip_root)/META/misc_info.txt
+endif # BOARD_AVB_VBMETA_MAINLINE
+ifneq (,$(strip $(BOARD_AVB_VBMETA_VENDOR)))
+	$(hide) echo "avb_vbmeta_vendor=$(BOARD_AVB_VBMETA_VENDOR)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_vendor_args=$(BOARD_AVB_MAKE_VBMETA_MAINLINE_IMAGE_ARGS)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_vendor_key_path=$(BOARD_AVB_VBMETA_VENDOR_KEY_PATH)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_vendor_algorithm=$(BOARD_AVB_VBMETA_VENDOR_ALGORITHM)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "avb_vbmeta_vendor_rollback_index_location=$(BOARD_AVB_VBMETA_VENDOR_ROLLBACK_INDEX_LOCATION)" >> $(zip_root)/META/misc_info.txt
+endif # BOARD_AVB_VBMETA_VENDOR_KEY_PATH
 endif # BOARD_AVB_ENABLE
 ifdef BOARD_BPT_INPUT_FILES
 	$(hide) echo "board_bpt_enable=true" >> $(zip_root)/META/misc_info.txt