Build vbmeta partitions with soong

AVB is Android Verified Boot:
https://source.android.com/docs/security/features/verifiedboot
It works by signing all the partitions, but then also including an extra
metadata paritition called vbmeta that depends on all the other
signed partitions. This creates a requirement that you update all those
partitions and the vbmeta partition together, so in order to relax that
requirement products can set up "chained" vbmeta partitions, where a
chained partition 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.

Bug: 377563298
Test: m nothing (the diff test currently fails, but this is a start)
Change-Id: I397af3a418bd23bcbf8f620c436afcacf69974fd
diff --git a/fsgen/vbmeta_partitions.go b/fsgen/vbmeta_partitions.go
new file mode 100644
index 0000000..f5b5f1c
--- /dev/null
+++ b/fsgen/vbmeta_partitions.go
@@ -0,0 +1,174 @@
+// Copyright (C) 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package fsgen
+
+import (
+	"android/soong/android"
+	"android/soong/filesystem"
+	"slices"
+	"strconv"
+
+	"github.com/google/blueprint/proptools"
+)
+
+type vbmetaModuleInfo struct {
+	// The name of the generated vbmeta module
+	moduleName string
+	// The name of the module that avb understands. This is the name passed to --chain_partition,
+	// and also the basename of the output file. (the output file is called partitionName + ".img")
+	partitionName string
+}
+
+// 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.
+//
+// AVB is Android Verified Boot: https://source.android.com/docs/security/features/verifiedboot
+// It works by signing all the partitions, but then also including an extra metadata paritition
+// called vbmeta that depends on all the other signed partitions. This creates a requirement
+// that you update all those partitions and the vbmeta partition together, so in order to relax
+// that requirement products can set up "chained" vbmeta partitions, where a chained partition
+// 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 {
+	partitionVars := ctx.Config().ProductVariables().PartitionVarsForSoongMigrationOnlyDoNotUse
+	// Some products seem to have BuildingVbmetaImage as true even when BoardAvbEnable is false
+	if !partitionVars.BuildingVbmetaImage || !partitionVars.BoardAvbEnable {
+		return nil
+	}
+
+	var result []vbmetaModuleInfo
+
+	var chainedPartitions []filesystem.ChainedPartitionProperties
+	var partitionTypesHandledByChainedPartitions []string
+	for chainedName, props := range partitionVars.ChainedVbmetaPartitions {
+		chainedName = "vbmeta_" + chainedName
+		if len(props.Partitions) == 0 {
+			continue
+		}
+		if len(props.Key) == 0 {
+			ctx.ModuleErrorf("No key found for chained avb partition %s", chainedName)
+			continue
+		}
+		if len(props.Algorithm) == 0 {
+			ctx.ModuleErrorf("No algorithm found for chained avb partition %s", chainedName)
+			continue
+		}
+		if len(props.RollbackIndex) == 0 {
+			ctx.ModuleErrorf("No rollback index found for chained avb partition %s", chainedName)
+			continue
+		}
+		ril, err := strconv.ParseInt(props.RollbackIndexLocation, 10, 32)
+		if err != nil {
+			ctx.ModuleErrorf("Rollback index location must be an int, got %s", props.RollbackIndexLocation)
+			continue
+		}
+		// The default is to use the PlatformSecurityPatch, and a lot of product config files
+		// just set it to the platform security patch, so detect that and don't set the property
+		// in soong.
+		var rollbackIndex *int64
+		if props.RollbackIndex != ctx.Config().PlatformSecurityPatch() {
+			i, err := strconv.ParseInt(props.RollbackIndex, 10, 32)
+			if err != nil {
+				ctx.ModuleErrorf("Rollback index must be an int, got %s", props.RollbackIndex)
+				continue
+			}
+			rollbackIndex = &i
+		}
+
+		var partitionModules []string
+		for _, partition := range props.Partitions {
+			partitionTypesHandledByChainedPartitions = append(partitionTypesHandledByChainedPartitions, partition)
+			if !slices.Contains(generatedPartitionTypes, partition) {
+				// The partition is probably unsupported.
+				continue
+			}
+			partitionModules = append(partitionModules, generatedModuleNameForPartition(ctx.Config(), partition))
+		}
+
+		name := generatedModuleName(ctx.Config(), chainedName)
+		ctx.CreateModuleInDirectory(
+			filesystem.VbmetaFactory,
+			".", // Create in the root directory for now so its easy to get the key
+			&filesystem.VbmetaProperties{
+				Partition_name:          proptools.StringPtr(chainedName),
+				Stem:                    proptools.StringPtr(chainedName + ".img"),
+				Private_key:             proptools.StringPtr(props.Key),
+				Algorithm:               &props.Algorithm,
+				Rollback_index:          rollbackIndex,
+				Rollback_index_location: &ril,
+				Partitions:              proptools.NewSimpleConfigurable(partitionModules),
+			}, &struct {
+				Name *string
+			}{
+				Name: &name,
+			},
+		).HideFromMake()
+
+		chainedPartitions = append(chainedPartitions, filesystem.ChainedPartitionProperties{
+			Name:                    &chainedName,
+			Rollback_index_location: &ril,
+			Private_key:             &props.Key,
+		})
+
+		result = append(result, vbmetaModuleInfo{
+			moduleName:    name,
+			partitionName: chainedName,
+		})
+	}
+
+	vbmetaModuleName := generatedModuleName(ctx.Config(), "vbmeta")
+
+	var algorithm *string
+	if len(partitionVars.BoardAvbAlgorithm) > 0 {
+		algorithm = proptools.StringPtr(partitionVars.BoardAvbAlgorithm)
+	}
+	ril, err := strconv.ParseInt(partitionVars.BoardAvbRollbackIndex, 10, 32)
+	if err != nil {
+		ctx.ModuleErrorf("Rollback index location must be an int, got %s", partitionVars.BoardAvbRollbackIndex)
+	}
+
+	var partitionModules []string
+	for _, partitionType := range generatedPartitionTypes {
+		if slices.Contains(partitionTypesHandledByChainedPartitions, partitionType) {
+			// Already handled by a chained vbmeta partition
+			continue
+		}
+		partitionModules = append(partitionModules, generatedModuleNameForPartition(ctx.Config(), partitionType))
+	}
+
+	ctx.CreateModuleInDirectory(
+		filesystem.VbmetaFactory,
+		".", // Create in the root directory for now so its easy to get the key
+		&filesystem.VbmetaProperties{
+			Stem:                    proptools.StringPtr("vbmeta.img"),
+			Algorithm:               algorithm,
+			Private_key:             proptools.StringPtr(partitionVars.BoardAvbKeyPath),
+			Rollback_index_location: &ril,
+			Chained_partitions:      chainedPartitions,
+			Partitions:              proptools.NewSimpleConfigurable(partitionModules),
+		}, &struct {
+			Name *string
+		}{
+			Name: &vbmetaModuleName,
+		},
+	).HideFromMake()
+
+	result = append(result, vbmetaModuleInfo{
+		moduleName:    vbmetaModuleName,
+		partitionName: "vbmeta",
+	})
+	return result
+}