Make chained_partitions accept other vbmeta modules

This is an attempt to make it clearer and less error-prone, by having
the vbmeta modules depend on their chained_partitions and get the
necessary information from them, instead of re-specifying the
information directly in the properties.

Bug: 377563298
Test: m nothing
Change-Id: I7d53651b8c48299b56d96cd8d184ae4b66ebeaf1
diff --git a/filesystem/vbmeta.go b/filesystem/vbmeta.go
index 6a3fc1f..ebb3ff9 100644
--- a/filesystem/vbmeta.go
+++ b/filesystem/vbmeta.go
@@ -26,8 +26,19 @@
 
 func init() {
 	android.RegisterModuleType("vbmeta", VbmetaFactory)
+	pctx.HostBinToolVariable("avbtool", "avbtool")
 }
 
+var (
+	extractPublicKeyRule = pctx.AndroidStaticRule("avb_extract_public_key",
+		blueprint.RuleParams{
+			Command: `${avbtool} extract_public_key --key $in --output $out`,
+			CommandDeps: []string{
+				"${avbtool}",
+			},
+		})
+)
+
 type vbmeta struct {
 	android.ModuleBase
 
@@ -60,8 +71,15 @@
 	// have to be signed (use_avb: true).
 	Partitions proptools.Configurable[[]string]
 
-	// List of chained partitions that this vbmeta deletages the verification.
-	Chained_partitions []ChainedPartitionProperties
+	// Metadata about the chained partitions that this vbmeta delegates the verification.
+	// This is an alternative to chained_partitions, using chained_partitions instead is simpler
+	// in most cases. However, this property allows building this vbmeta partition without
+	// its chained partitions existing in this build.
+	Chained_partition_metadata []ChainedPartitionProperties
+
+	// List of chained partitions that this vbmeta delegates the verification. They are the
+	// names of other vbmeta modules.
+	Chained_partitions []string
 
 	// List of key-value pair of avb properties
 	Avb_properties []avbProperty
@@ -93,6 +111,20 @@
 	Private_key *string `android:"path"`
 }
 
+type vbmetaPartitionInfo struct {
+	// Name of the partition
+	Name string
+
+	// Rollback index location, non-negative int
+	RollbackIndexLocation int
+
+	// The path to the public key of the private key used to sign this partition. Derived from
+	// the private key.
+	PublicKey android.Path
+}
+
+var vbmetaPartitionProvider = blueprint.NewProvider[vbmetaPartitionInfo]()
+
 // vbmeta is the partition image that has the verification information for other partitions.
 func VbmetaFactory() android.Module {
 	module := &vbmeta{}
@@ -103,13 +135,18 @@
 
 type vbmetaDep struct {
 	blueprint.BaseDependencyTag
-	kind string
 }
 
-var vbmetaPartitionDep = vbmetaDep{kind: "partition"}
+type chainedPartitionDep struct {
+	blueprint.BaseDependencyTag
+}
+
+var vbmetaPartitionDep = vbmetaDep{}
+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...)
 }
 
 func (v *vbmeta) installFileName() string {
@@ -124,8 +161,6 @@
 const vbmetaMaxSize = 64 * 1024
 
 func (v *vbmeta) GenerateAndroidBuildActions(ctx android.ModuleContext) {
-	extractedPublicKeys := v.extractPublicKeys(ctx)
-
 	v.output = android.PathForModuleOut(ctx, v.installFileName()).OutputPath
 
 	builder := android.NewRuleBuilder(pctx, ctx)
@@ -175,25 +210,66 @@
 		cmd.FlagWithInput("--include_descriptors_from_image ", signedImage)
 	}
 
-	for i, cp := range v.properties.Chained_partitions {
-		name := proptools.String(cp.Name)
+	seenRils := make(map[int]bool)
+	for _, cp := range ctx.GetDirectDepsWithTag(vbmetaChainedPartitionDep) {
+		info, ok := android.OtherModuleProvider(ctx, cp, vbmetaPartitionProvider)
+		if !ok {
+			ctx.PropertyErrorf("chained_partitions", "Expected all modules in chained_partitions to provide vbmetaPartitionProvider, but %s did not", cp.Name())
+			continue
+		}
+		if info.Name == "" {
+			ctx.PropertyErrorf("chained_partitions", "name must be specified")
+			continue
+		}
+
+		ril := info.RollbackIndexLocation
+		if ril < 0 {
+			ctx.PropertyErrorf("chained_partitions", "rollback index location must be 0, 1, 2, ...")
+			continue
+		} else if seenRils[ril] {
+			ctx.PropertyErrorf("chained_partitions", "Multiple chained partitions with the same rollback index location %d", ril)
+			continue
+		}
+		seenRils[ril] = true
+
+		publicKey := info.PublicKey
+		cmd.FlagWithArg("--chain_partition ", fmt.Sprintf("%s:%d:%s", info.Name, ril, publicKey.String()))
+		cmd.Implicit(publicKey)
+	}
+	for _, cpm := range v.properties.Chained_partition_metadata {
+		name := proptools.String(cpm.Name)
 		if name == "" {
 			ctx.PropertyErrorf("chained_partitions", "name must be specified")
 			continue
 		}
 
-		ril := proptools.IntDefault(cp.Rollback_index_location, i+1)
+		ril := proptools.IntDefault(cpm.Rollback_index_location, -1)
 		if ril < 0 {
-			ctx.PropertyErrorf("chained_partitions", "must be 0, 1, 2, ...")
+			ctx.PropertyErrorf("chained_partition_metadata", "rollback index location must be 0, 1, 2, ...")
+			continue
+		} else if seenRils[ril] {
+			ctx.PropertyErrorf("chained_partition_metadata", "Multiple chained partitions with the same rollback index location %d", ril)
+			continue
+		}
+		seenRils[ril] = true
+
+		var publicKey android.Path
+		if cpm.Public_key != nil {
+			publicKey = android.PathForModuleSrc(ctx, *cpm.Public_key)
+		} else if cpm.Private_key != nil {
+			privateKey := android.PathForModuleSrc(ctx, *cpm.Private_key)
+			extractedPublicKey := android.PathForModuleOut(ctx, "chained_metadata", name+".avbpubkey")
+			ctx.Build(pctx, android.BuildParams{
+				Rule:   extractPublicKeyRule,
+				Input:  privateKey,
+				Output: extractedPublicKey,
+			})
+			publicKey = extractedPublicKey
+		} else {
+			ctx.PropertyErrorf("public_key", "Either public_key or private_key must be specified")
 			continue
 		}
 
-		var publicKey android.Path
-		if cp.Public_key != nil {
-			publicKey = android.PathForModuleSrc(ctx, proptools.String(cp.Public_key))
-		} else {
-			publicKey = extractedPublicKeys[name]
-		}
 		cmd.FlagWithArg("--chain_partition ", fmt.Sprintf("%s:%d:%s", name, ril, publicKey.String()))
 		cmd.Implicit(publicKey)
 	}
@@ -211,6 +287,19 @@
 	v.installDir = android.PathForModuleInstall(ctx, "etc")
 	ctx.InstallFile(v.installDir, v.installFileName(), v.output)
 
+	extractedPublicKey := android.PathForModuleOut(ctx, v.partitionName()+".avbpubkey")
+	ctx.Build(pctx, android.BuildParams{
+		Rule:   extractPublicKeyRule,
+		Input:  key,
+		Output: extractedPublicKey,
+	})
+
+	android.SetProvider(ctx, vbmetaPartitionProvider, vbmetaPartitionInfo{
+		Name:                  v.partitionName(),
+		RollbackIndexLocation: ril,
+		PublicKey:             extractedPublicKey,
+	})
+
 	ctx.SetOutputFiles([]android.Path{v.output}, "")
 }
 
@@ -224,43 +313,6 @@
 	}
 }
 
-// Extract public keys from chained_partitions.private_key. The keys are indexed with the partition
-// name.
-func (v *vbmeta) extractPublicKeys(ctx android.ModuleContext) map[string]android.OutputPath {
-	result := make(map[string]android.OutputPath)
-
-	builder := android.NewRuleBuilder(pctx, ctx)
-	for _, cp := range v.properties.Chained_partitions {
-		if cp.Private_key == nil {
-			continue
-		}
-
-		name := proptools.String(cp.Name)
-		if name == "" {
-			ctx.PropertyErrorf("chained_partitions", "name must be specified")
-			continue
-		}
-
-		if _, ok := result[name]; ok {
-			ctx.PropertyErrorf("chained_partitions", "name %q is duplicated", name)
-			continue
-		}
-
-		privateKeyFile := android.PathForModuleSrc(ctx, proptools.String(cp.Private_key))
-		publicKeyFile := android.PathForModuleOut(ctx, name+".avbpubkey").OutputPath
-
-		builder.Command().
-			BuiltTool("avbtool").
-			Text("extract_public_key").
-			FlagWithInput("--key ", privateKeyFile).
-			FlagWithOutput("--output ", publicKeyFile)
-
-		result[name] = publicKeyFile
-	}
-	builder.Build("vbmeta_extract_public_key", fmt.Sprintf("Extract public keys for %s", ctx.ModuleName()))
-	return result
-}
-
 var _ android.AndroidMkProviderInfoProducer = (*vbmeta)(nil)
 
 func (v *vbmeta) PrepareAndroidMKProviderInfo(config android.Config) *android.AndroidMkProviderInfo {