Modify autogen vbmeta chain and include partitions logic

This change modifies the logic in selecting the chain and include
partitions for the top level vbmeta partition generation to follow that
in make.

Implementation details:
- Introduce separate functions for selecting the include and chain
  partitions
- Set vbmetaPartitionProvider in boot images and filesystem modules
  given that boot images can be selected as chained partitions.

Test: compare soong and make generated vbmeta partitions command
Bug: 390204058
Change-Id: I7940629212c4624e88d25b162755205503cfa469

3 files changed, 170 insertions, 71 deletions

diff --git a/fsgen/boot_imgs.go b/fsgen/boot_imgs.go
index 889a4c271..58ebcc4ba 100644
--- a/fsgen/boot_imgs.go
+++ b/fsgen/boot_imgs.go
@@ -69,23 +69,26 @@ func createBootImage(ctx android.LoadHookContext, dtbImg dtbImg) bool {
 	ctx.CreateModule(
 		filesystem.BootimgFactory,
 		&filesystem.BootimgProperties{
-			Kernel_prebuilt:    proptools.StringPtr(":" + kernelFilegroupName),
-			Header_version:     proptools.StringPtr(partitionVariables.BoardBootHeaderVersion),
-			Partition_size:     partitionSize,
-			Use_avb:            avbInfo.avbEnable,
-			Avb_mode:           avbInfo.avbMode,
-			Avb_private_key:    avbInfo.avbkeyFilegroup,
-			Avb_rollback_index: avbInfo.avbRollbackIndex,
-			Avb_algorithm:      avbInfo.avbAlgorithm,
-			Security_patch:     securityPatch,
-			Dtb_prebuilt:       dtbPrebuilt,
-			Cmdline:            cmdline,
-			Stem:               proptools.StringPtr("boot.img"),
+			Kernel_prebuilt:             proptools.StringPtr(":" + kernelFilegroupName),
+			Header_version:              proptools.StringPtr(partitionVariables.BoardBootHeaderVersion),
+			Partition_size:              partitionSize,
+			Use_avb:                     avbInfo.avbEnable,
+			Avb_mode:                    avbInfo.avbMode,
+			Avb_private_key:             avbInfo.avbkeyFilegroup,
+			Avb_rollback_index:          avbInfo.avbRollbackIndex,
+			Avb_rollback_index_location: avbInfo.avbRollbackIndexLocation,
+			Avb_algorithm:               avbInfo.avbAlgorithm,
+			Security_patch:              securityPatch,
+			Dtb_prebuilt:                dtbPrebuilt,
+			Cmdline:                     cmdline,
+			Stem:                        proptools.StringPtr("boot.img"),
 		},
 		&struct {
-			Name *string
+			Name       *string
+			Visibility []string
 		}{
-			Name: proptools.StringPtr(bootImageName),
+			Name:       proptools.StringPtr(bootImageName),
+			Visibility: []string{"//visibility:public"},
 		},
 	)
 	return true
@@ -123,23 +126,26 @@ func createVendorBootImage(ctx android.LoadHookContext, dtbImg dtbImg) bool {
 	ctx.CreateModule(
 		filesystem.BootimgFactory,
 		&filesystem.BootimgProperties{
-			Boot_image_type:    proptools.StringPtr("vendor_boot"),
-			Ramdisk_module:     proptools.StringPtr(generatedModuleNameForPartition(ctx.Config(), "vendor_ramdisk")),
-			Header_version:     proptools.StringPtr(partitionVariables.BoardBootHeaderVersion),
-			Partition_size:     partitionSize,
-			Use_avb:            avbInfo.avbEnable,
-			Avb_mode:           avbInfo.avbMode,
-			Avb_private_key:    avbInfo.avbkeyFilegroup,
-			Avb_rollback_index: avbInfo.avbRollbackIndex,
-			Dtb_prebuilt:       dtbPrebuilt,
-			Cmdline:            cmdline,
-			Bootconfig:         vendorBootConfigImg,
-			Stem:               proptools.StringPtr("vendor_boot.img"),
+			Boot_image_type:             proptools.StringPtr("vendor_boot"),
+			Ramdisk_module:              proptools.StringPtr(generatedModuleNameForPartition(ctx.Config(), "vendor_ramdisk")),
+			Header_version:              proptools.StringPtr(partitionVariables.BoardBootHeaderVersion),
+			Partition_size:              partitionSize,
+			Use_avb:                     avbInfo.avbEnable,
+			Avb_mode:                    avbInfo.avbMode,
+			Avb_private_key:             avbInfo.avbkeyFilegroup,
+			Avb_rollback_index:          avbInfo.avbRollbackIndex,
+			Avb_rollback_index_location: avbInfo.avbRollbackIndexLocation,
+			Dtb_prebuilt:                dtbPrebuilt,
+			Cmdline:                     cmdline,
+			Bootconfig:                  vendorBootConfigImg,
+			Stem:                        proptools.StringPtr("vendor_boot.img"),
 		},
 		&struct {
-			Name *string
+			Name       *string
+			Visibility []string
 		}{
-			Name: proptools.StringPtr(bootImageName),
+			Name:       proptools.StringPtr(bootImageName),
+			Visibility: []string{"//visibility:public"},
 		},
 	)
 	return true
@@ -172,22 +178,25 @@ func createInitBootImage(ctx android.LoadHookContext) bool {
 	ctx.CreateModule(
 		filesystem.BootimgFactory,
 		&filesystem.BootimgProperties{
-			Boot_image_type:    proptools.StringPtr("init_boot"),
-			Ramdisk_module:     proptools.StringPtr(generatedModuleNameForPartition(ctx.Config(), "ramdisk")),
-			Header_version:     proptools.StringPtr(partitionVariables.BoardBootHeaderVersion),
-			Security_patch:     securityPatch,
-			Partition_size:     partitionSize,
-			Use_avb:            avbInfo.avbEnable,
-			Avb_mode:           avbInfo.avbMode,
-			Avb_private_key:    avbInfo.avbkeyFilegroup,
-			Avb_rollback_index: avbInfo.avbRollbackIndex,
-			Avb_algorithm:      avbInfo.avbAlgorithm,
-			Stem:               proptools.StringPtr("init_boot.img"),
+			Boot_image_type:             proptools.StringPtr("init_boot"),
+			Ramdisk_module:              proptools.StringPtr(generatedModuleNameForPartition(ctx.Config(), "ramdisk")),
+			Header_version:              proptools.StringPtr(partitionVariables.BoardBootHeaderVersion),
+			Security_patch:              securityPatch,
+			Partition_size:              partitionSize,
+			Use_avb:                     avbInfo.avbEnable,
+			Avb_mode:                    avbInfo.avbMode,
+			Avb_private_key:             avbInfo.avbkeyFilegroup,
+			Avb_rollback_index:          avbInfo.avbRollbackIndex,
+			Avb_rollback_index_location: avbInfo.avbRollbackIndexLocation,
+			Avb_algorithm:               avbInfo.avbAlgorithm,
+			Stem:                        proptools.StringPtr("init_boot.img"),
 		},
 		&struct {
-			Name *string
+			Name       *string
+			Visibility []string
 		}{
-			Name: proptools.StringPtr(bootImageName),
+			Name:       proptools.StringPtr(bootImageName),
+			Visibility: []string{"//visibility:public"},
 		},
 	)
 	return true
diff --git a/fsgen/filesystem_creator.go b/fsgen/filesystem_creator.go
index d2f00cd26..ebcc68b2b 100644
--- a/fsgen/filesystem_creator.go
+++ b/fsgen/filesystem_creator.go
@@ -178,7 +178,7 @@ func (f *filesystemCreator) createInternalModules(ctx android.LoadHookContext) {
 		)
 	}
 
-	for _, x := range createVbmetaPartitions(ctx, finalSoongGeneratedPartitions) {
+	for _, x := range f.createVbmetaPartitions(ctx, finalSoongGeneratedPartitions) {
 		f.properties.Vbmeta_module_names = append(f.properties.Vbmeta_module_names, x.moduleName)
 		f.properties.Vbmeta_partition_names = append(f.properties.Vbmeta_partition_names, x.partitionName)
 	}
@@ -865,6 +865,8 @@ func generateFsProps(ctx android.EarlyModuleContext, partitionType string) (*fil
 	fsProps.Avb_algorithm = avbInfo.avbAlgorithm
 	// BOARD_AVB_SYSTEM_ROLLBACK_INDEX
 	fsProps.Rollback_index = avbInfo.avbRollbackIndex
+	// BOARD_AVB_SYSTEM_ROLLBACK_INDEX_LOCATION
+	fsProps.Rollback_index_location = avbInfo.avbRollbackIndexLocation
 	fsProps.Avb_hash_algorithm = avbInfo.avbHashAlgorithm
 
 	fsProps.Partition_name = proptools.StringPtr(partitionType)
@@ -893,13 +895,14 @@ func generateFsProps(ctx android.EarlyModuleContext, partitionType string) (*fil
 }
 
 type avbInfo struct {
-	avbEnable        *bool
-	avbKeyPath       *string
-	avbkeyFilegroup  *string
-	avbAlgorithm     *string
-	avbRollbackIndex *int64
-	avbMode          *string
-	avbHashAlgorithm *string
+	avbEnable                *bool
+	avbKeyPath               *string
+	avbkeyFilegroup          *string
+	avbAlgorithm             *string
+	avbRollbackIndex         *int64
+	avbRollbackIndexLocation *int64
+	avbMode                  *string
+	avbHashAlgorithm         *string
 }
 
 func getAvbInfo(config android.Config, partitionType string) avbInfo {
@@ -946,6 +949,13 @@ func getAvbInfo(config android.Config, partitionType string) avbInfo {
 			}
 			result.avbRollbackIndex = &parsed
 		}
+		if specificPartitionVars.BoardAvbRollbackIndexLocation != "" {
+			parsed, err := strconv.ParseInt(specificPartitionVars.BoardAvbRollbackIndexLocation, 10, 64)
+			if err != nil {
+				panic(fmt.Sprintf("Rollback index location must be an int, got %s", specificPartitionVars.BoardAvbRollbackIndexLocation))
+			}
+			result.avbRollbackIndexLocation = &parsed
+		}
 
 		// Make allows you to pass arbitrary arguments to avbtool via this variable, but in practice
 		// it's only used for --hash_algorithm. The soong module has a dedicated property for the
diff --git a/fsgen/vbmeta_partitions.go b/fsgen/vbmeta_partitions.go
index e3dc41679..be738ea86 100644
--- a/fsgen/vbmeta_partitions.go
+++ b/fsgen/vbmeta_partitions.go
@@ -19,7 +19,6 @@ import (
 	"android/soong/filesystem"
 	"slices"
 	"strconv"
-	"strings"
 
 	"github.com/google/blueprint/proptools"
 )
@@ -32,6 +31,27 @@ type vbmetaModuleInfo struct {
 	partitionName string
 }
 
+// https://cs.android.com/android/platform/superproject/main/+/main:build/make/core/Makefile;l=4849;drc=62e20f0d218f60bae563b4ee742d88cca1fc1901
+var avbPartitions = []string{
+	"boot",
+	"init_boot",
+	"vendor_boot",
+	"vendor_kernel_boot",
+	"system",
+	"vendor",
+	"product",
+	"system_ext",
+	"odm",
+	"vendor_dlkm",
+	"odm_dlkm",
+	"system_dlkm",
+	"dtbo",
+	"pvmfw",
+	"recovery",
+	"vbmeta_system",
+	"vbmeta_vendor",
+}
+
 // Creates the vbmeta partition and the chained vbmeta partitions. Returns the list of module names
 // that the function created. May return nil if the product isn't using avb.
 //
@@ -43,7 +63,7 @@ type vbmetaModuleInfo struct {
 // like vbmeta_system might contain the avb metadata for just a few products. In cuttlefish
 // vbmeta_system contains metadata about product, system, and system_ext. Using chained partitions,
 // that group of partitions can be updated independently from the other signed partitions.
-func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes []string) []vbmetaModuleInfo {
+func (f *filesystemCreator) createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes []string) []vbmetaModuleInfo {
 	partitionVars := ctx.Config().ProductVariables().PartitionVarsForSoongMigrationOnlyDoNotUse
 	// Some products seem to have BuildingVbmetaImage as true even when BoardAvbEnable is false
 	if !partitionVars.BuildingVbmetaImage || !partitionVars.BoardAvbEnable {
@@ -52,14 +72,16 @@ func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes
 
 	var result []vbmetaModuleInfo
 
-	var chainedPartitions []string
-	var partitionTypesHandledByChainedPartitions []string
+	// https://cs.android.com/android/platform/superproject/main/+/main:build/make/core/Makefile;l=4593;drc=62e20f0d218f60bae563b4ee742d88cca1fc1901
+	var internalAvbPartitionsInChainedVbmetaImages []string
+	var chainedPartitionTypes []string
 	for _, chainedName := range android.SortedKeys(partitionVars.ChainedVbmetaPartitions) {
 		props := partitionVars.ChainedVbmetaPartitions[chainedName]
 		chainedName = "vbmeta_" + chainedName
 		if len(props.Partitions) == 0 {
 			continue
 		}
+		internalAvbPartitionsInChainedVbmetaImages = append(internalAvbPartitionsInChainedVbmetaImages, props.Partitions...)
 		if len(props.Key) == 0 {
 			ctx.ModuleErrorf("No key found for chained avb partition %q", chainedName)
 			continue
@@ -92,7 +114,6 @@ func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes
 
 		var partitionModules []string
 		for _, partition := range props.Partitions {
-			partitionTypesHandledByChainedPartitions = append(partitionTypesHandledByChainedPartitions, partition)
 			if !slices.Contains(generatedPartitionTypes, partition) {
 				// The partition is probably unsupported.
 				continue
@@ -100,7 +121,7 @@ func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes
 			partitionModules = append(partitionModules, generatedModuleNameForPartition(ctx.Config(), partition))
 		}
 
-		name := generatedModuleName(ctx.Config(), chainedName)
+		name := generatedModuleNameForPartition(ctx.Config(), chainedName)
 		ctx.CreateModuleInDirectory(
 			filesystem.VbmetaFactory,
 			".", // Create in the root directory for now so its easy to get the key
@@ -119,15 +140,15 @@ func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes
 			},
 		).HideFromMake()
 
-		chainedPartitions = append(chainedPartitions, name)
-
 		result = append(result, vbmetaModuleInfo{
 			moduleName:    name,
 			partitionName: chainedName,
 		})
+
+		chainedPartitionTypes = append(chainedPartitionTypes, chainedName)
 	}
 
-	vbmetaModuleName := generatedModuleName(ctx.Config(), "vbmeta")
+	vbmetaModuleName := generatedModuleNameForPartition(ctx.Config(), "vbmeta")
 
 	var algorithm *string
 	var ri *int64
@@ -148,20 +169,79 @@ func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes
 		ri = &parsedRi
 	}
 
-	var partitionModules []string
-	for _, partitionType := range generatedPartitionTypes {
-		if slices.Contains(partitionTypesHandledByChainedPartitions, partitionType) {
-			// Already handled by a chained vbmeta partition
-			continue
+	// --chain_partition argument is only set for partitions that set
+	// `BOARD_AVB_<partition name>_KEY_PATH` value and is not "recovery"
+	// https://cs.android.com/android/platform/superproject/main/+/main:build/make/core/Makefile;l=4823;drc=62e20f0d218f60bae563b4ee742d88cca1fc1901
+	includeAsChainedPartitionInVbmeta := func(partition string) bool {
+		val, ok := partitionVars.PartitionQualifiedVariables[partition]
+		return ok && len(val.BoardAvbKeyPath) > 0 && partition != "recovery"
+	}
+
+	// --include_descriptors_from_image is passed if both conditions are met:
+	// - `BOARD_AVB_<partition name>_KEY_PATH` value is not set
+	// - not included in INTERNAL_AVB_PARTITIONS_IN_CHAINED_VBMETA_IMAGES
+	// for partitions that set INSTALLED_<partition name>IMAGE_TARGET
+	// https://cs.android.com/android/platform/superproject/main/+/main:build/make/core/Makefile;l=4827;drc=62e20f0d218f60bae563b4ee742d88cca1fc1901
+	includeAsIncludedPartitionInVbmeta := func(partition string) bool {
+		if android.InList(partition, internalAvbPartitionsInChainedVbmetaImages) {
+			// Already handled by chained vbmeta partitions
+			return false
+		}
+		partitionQualifiedVars := partitionVars.PartitionQualifiedVariables[partition]
+
+		// The return logic in the switch cases below are identical to
+		// ifdef INSTALLED_<partition name>IMAGE_TARGET
+		switch partition {
+		case "boot":
+			return partitionQualifiedVars.BuildingImage || partitionQualifiedVars.PrebuiltImage || partitionVars.BoardUsesRecoveryAsBoot
+		case "vendor_kernel_boot", "dtbo":
+			return partitionQualifiedVars.PrebuiltImage
+		case "system":
+			return partitionQualifiedVars.BuildingImage
+		case "init_boot", "vendor_boot", "vendor", "product", "system_ext", "odm", "vendor_dlkm", "odm_dlkm", "system_dlkm":
+			return partitionQualifiedVars.BuildingImage || partitionQualifiedVars.PrebuiltImage
+		// TODO: Import BOARD_USES_PVMFWIMAGE
+		// ifeq ($(BOARD_USES_PVMFWIMAGE),true)
+		// case "pvmfw":
+		case "recovery":
+			// ifdef INSTALLED_RECOVERYIMAGE_TARGET
+			return !ctx.DeviceConfig().BoardUsesRecoveryAsBoot() && !ctx.DeviceConfig().BoardMoveRecoveryResourcesToVendorBoot()
+		// Technically these partitions are determined based on len(BOARD_AVB_VBMETA_SYSTEM) and
+		// len(BOARD_AVB_VBMETA_VENDOR) but if these are non empty these partitions are
+		// already included in the chained partitions.
+		case "vbmeta_system", "vbmeta_vendor":
+			return false
+		default:
+			return false
 		}
-		if strings.Contains(partitionType, "ramdisk") || strings.Contains(partitionType, "boot") || partitionType == "userdata" || partitionType == "recovery" {
-			// ramdisk and userdata are never signed with avb information
-			// recovery partition is skipped in adding the partition descriptor into vbmeta.img.
-			// boot partitions just have the avb footer, and don't have a corresponding vbmeta
-			// partition.
+	}
+
+	var chainedPartitionModules []string
+	var includePartitionModules []string
+	allGeneratedPartitionTypes := append(generatedPartitionTypes,
+		chainedPartitionTypes...,
+	)
+	if len(f.properties.Boot_image) > 0 {
+		allGeneratedPartitionTypes = append(allGeneratedPartitionTypes, "boot")
+	}
+	if len(f.properties.Init_boot_image) > 0 {
+		allGeneratedPartitionTypes = append(allGeneratedPartitionTypes, "init_boot")
+	}
+	if len(f.properties.Vendor_boot_image) > 0 {
+		allGeneratedPartitionTypes = append(allGeneratedPartitionTypes, "vendor_boot")
+	}
+
+	// https://cs.android.com/android/platform/superproject/main/+/main:build/make/core/Makefile;l=4919;drc=62e20f0d218f60bae563b4ee742d88cca1fc1901
+	for _, partitionType := range android.SortedUniqueStrings(append(avbPartitions, chainedPartitionTypes...)) {
+		if !android.InList(partitionType, allGeneratedPartitionTypes) {
+			// Skip if the partition is not auto generated
 			continue
 		}
-		partitionModules = append(partitionModules, generatedModuleNameForPartition(ctx.Config(), partitionType))
+		if includeAsChainedPartitionInVbmeta(partitionType) {
+			chainedPartitionModules = append(chainedPartitionModules, generatedModuleNameForPartition(ctx.Config(), partitionType))
+		} else if includeAsIncludedPartitionInVbmeta(partitionType) {
+			includePartitionModules = append(includePartitionModules, generatedModuleNameForPartition(ctx.Config(), partitionType))
+		}
 	}
 
 	ctx.CreateModuleInDirectory(
@@ -172,8 +252,8 @@ func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes
 			Algorithm:          algorithm,
 			Private_key:        key,
 			Rollback_index:     ri,
-			Chained_partitions: chainedPartitions,
-			Partitions:         proptools.NewSimpleConfigurable(partitionModules),
+			Chained_partitions: chainedPartitionModules,
+			Partitions:         proptools.NewSimpleConfigurable(includePartitionModules),
 			Partition_name:     proptools.StringPtr("vbmeta"),
 		}, &struct {
 			Name *string