Merge "Modify autogen vbmeta chain and include partitions logic" into main
diff --git a/android/variable.go b/android/variable.go
index 14094e2..3e637fe 100644
--- a/android/variable.go
+++ b/android/variable.go
@@ -555,6 +555,7 @@
 
 type PartitionQualifiedVariablesType struct {
 	BuildingImage               bool   `json:",omitempty"`
+	PrebuiltImage               bool   `json:",omitempty"`
 	BoardErofsCompressor        string `json:",omitempty"`
 	BoardErofsCompressHints     string `json:",omitempty"`
 	BoardErofsPclusterSize      string `json:",omitempty"`
diff --git a/filesystem/bootimg.go b/filesystem/bootimg.go
index 98af479..6d6c15c 100644
--- a/filesystem/bootimg.go
+++ b/filesystem/bootimg.go
@@ -99,6 +99,9 @@
 	// The index used to prevent rollback of the image on device.
 	Avb_rollback_index *int64
 
+	// Rollback index location of this image. Must be 0, 1, 2, etc.
+	Avb_rollback_index_location *int64
+
 	// The security patch passed to as the com.android.build.<type>.security_patch avb property.
 	// Replacement for the make variables BOOT_SECURITY_PATCH / INIT_BOOT_SECURITY_PATCH.
 	Security_patch *string
@@ -237,6 +240,27 @@
 		Bootconfig: b.getBootconfigPath(ctx),
 		Output:     output,
 	})
+
+	extractedPublicKey := android.PathForModuleOut(ctx, b.partitionName()+".avbpubkey")
+	if b.properties.Avb_private_key != nil {
+		key := android.PathForModuleSrc(ctx, proptools.String(b.properties.Avb_private_key))
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   extractPublicKeyRule,
+			Input:  key,
+			Output: extractedPublicKey,
+		})
+	}
+	var ril int
+	if b.properties.Avb_rollback_index_location != nil {
+		ril = proptools.Int(b.properties.Avb_rollback_index_location)
+	}
+
+	android.SetProvider(ctx, vbmetaPartitionProvider, vbmetaPartitionInfo{
+		Name:                  b.bootImageType.String(),
+		RollbackIndexLocation: ril,
+		PublicKey:             extractedPublicKey,
+		Output:                output,
+	})
 }
 
 var BootimgInfoProvider = blueprint.NewProvider[BootimgInfo]()
diff --git a/filesystem/filesystem.go b/filesystem/filesystem.go
index 68cbee9..357ec32 100644
--- a/filesystem/filesystem.go
+++ b/filesystem/filesystem.go
@@ -131,6 +131,9 @@
 	// The index used to prevent rollback of the image. Only used if use_avb is true.
 	Rollback_index *int64
 
+	// Rollback index location of this image. Must be 0, 1, 2, etc.
+	Rollback_index_location *int64
+
 	// Name of the partition stored in vbmeta desc. Defaults to the name of this module.
 	Partition_name *string
 
@@ -537,6 +540,33 @@
 	if proptools.Bool(f.properties.Unchecked_module) {
 		ctx.UncheckedModule()
 	}
+
+	f.setVbmetaPartitionProvider(ctx)
+}
+
+func (f *filesystem) setVbmetaPartitionProvider(ctx android.ModuleContext) {
+	var extractedPublicKey android.ModuleOutPath
+	if f.properties.Avb_private_key != nil {
+		key := android.PathForModuleSrc(ctx, proptools.String(f.properties.Avb_private_key))
+		extractedPublicKey = android.PathForModuleOut(ctx, f.partitionName()+".avbpubkey")
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   extractPublicKeyRule,
+			Input:  key,
+			Output: extractedPublicKey,
+		})
+	}
+
+	var ril int
+	if f.properties.Rollback_index_location != nil {
+		ril = proptools.Int(f.properties.Rollback_index_location)
+	}
+
+	android.SetProvider(ctx, vbmetaPartitionProvider, vbmetaPartitionInfo{
+		Name:                  f.partitionName(),
+		RollbackIndexLocation: ril,
+		PublicKey:             extractedPublicKey,
+		Output:                f.output,
+	})
 }
 
 func (f *filesystem) getMapFile(ctx android.ModuleContext) android.WritablePath {
diff --git a/filesystem/vbmeta.go b/filesystem/vbmeta.go
index b7f312c..91826b2 100644
--- a/filesystem/vbmeta.go
+++ b/filesystem/vbmeta.go
@@ -148,8 +148,8 @@
 var vbmetaChainedPartitionDep = chainedPartitionDep{}
 
 func (v *vbmeta) DepsMutator(ctx android.BottomUpMutatorContext) {
-	ctx.AddDependency(ctx.Module(), vbmetaPartitionDep, v.properties.Partitions.GetOrDefault(ctx, nil)...)
-	ctx.AddDependency(ctx.Module(), vbmetaChainedPartitionDep, v.properties.Chained_partitions...)
+	ctx.AddVariationDependencies(ctx.Config().AndroidFirstDeviceTarget.Variations(), vbmetaPartitionDep, v.properties.Partitions.GetOrDefault(ctx, nil)...)
+	ctx.AddVariationDependencies(ctx.Config().AndroidFirstDeviceTarget.Variations(), vbmetaChainedPartitionDep, v.properties.Chained_partitions...)
 }
 
 func (v *vbmeta) installFileName() string {
diff --git a/fsgen/boot_imgs.go b/fsgen/boot_imgs.go
index 889a4c2..58ebcc4 100644
--- a/fsgen/boot_imgs.go
+++ b/fsgen/boot_imgs.go
@@ -69,23 +69,26 @@
 	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 @@
 	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 @@
 	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 d2f00cd..ebcc68b 100644
--- a/fsgen/filesystem_creator.go
+++ b/fsgen/filesystem_creator.go
@@ -178,7 +178,7 @@
 		)
 	}
 
-	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 @@
 	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 @@
 }
 
 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 @@
 			}
 			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 e3dc416..be738ea 100644
--- a/fsgen/vbmeta_partitions.go
+++ b/fsgen/vbmeta_partitions.go
@@ -19,7 +19,6 @@
 	"android/soong/filesystem"
 	"slices"
 	"strconv"
-	"strings"
 
 	"github.com/google/blueprint/proptools"
 )
@@ -32,6 +31,27 @@
 	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 @@
 // 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 @@
 
 	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 @@
 
 		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 @@
 			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 @@
 			},
 		).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 @@
 		ri = &parsedRi
 	}
 
-	var partitionModules []string
-	for _, partitionType := range generatedPartitionTypes {
-		if slices.Contains(partitionTypesHandledByChainedPartitions, partitionType) {
-			// Already handled by a chained vbmeta partition
+	// --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
+		}
+	}
+
+	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
 		}
-		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.
-			continue
+		if includeAsChainedPartitionInVbmeta(partitionType) {
+			chainedPartitionModules = append(chainedPartitionModules, generatedModuleNameForPartition(ctx.Config(), partitionType))
+		} else if includeAsIncludedPartitionInVbmeta(partitionType) {
+			includePartitionModules = append(includePartitionModules, generatedModuleNameForPartition(ctx.Config(), partitionType))
 		}
-		partitionModules = append(partitionModules, generatedModuleNameForPartition(ctx.Config(), partitionType))
 	}
 
 	ctx.CreateModuleInDirectory(
@@ -172,8 +252,8 @@
 			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