Add prebuilt_apex.

Bug: 127789981
Test: apex_test.go + com.android.tzdata.apex
Change-Id: I09bb0a4b2acf310c55c789569da3c9d755638232
diff --git a/android/prebuilt.go b/android/prebuilt.go
index 5bd0e2d..df25a89 100644
--- a/android/prebuilt.go
+++ b/android/prebuilt.go
@@ -18,6 +18,7 @@
 	"fmt"
 
 	"github.com/google/blueprint"
+	"github.com/google/blueprint/proptools"
 )
 
 // This file implements common functionality for handling modules that may exist as prebuilts,
@@ -42,6 +43,7 @@
 	properties PrebuiltProperties
 	module     Module
 	srcs       *[]string
+	src        *string
 }
 
 func (p *Prebuilt) Name(name string) string {
@@ -49,19 +51,27 @@
 }
 
 func (p *Prebuilt) SingleSourcePath(ctx ModuleContext) Path {
-	if len(*p.srcs) == 0 {
-		ctx.PropertyErrorf("srcs", "missing prebuilt source file")
-		return nil
-	}
+	if p.srcs != nil {
+		if len(*p.srcs) == 0 {
+			ctx.PropertyErrorf("srcs", "missing prebuilt source file")
+			return nil
+		}
 
-	if len(*p.srcs) > 1 {
-		ctx.PropertyErrorf("srcs", "multiple prebuilt source files")
-		return nil
-	}
+		if len(*p.srcs) > 1 {
+			ctx.PropertyErrorf("srcs", "multiple prebuilt source files")
+			return nil
+		}
 
-	// Return the singleton source after expanding any filegroup in the
-	// sources.
-	return PathForModuleSrc(ctx, (*p.srcs)[0])
+		// Return the singleton source after expanding any filegroup in the
+		// sources.
+		return PathForModuleSrc(ctx, (*p.srcs)[0])
+	} else {
+		if proptools.String(p.src) == "" {
+			ctx.PropertyErrorf("src", "missing prebuilt source file")
+			return nil
+		}
+		return PathForModuleSrc(ctx, *p.src)
+	}
 }
 
 func InitPrebuiltModule(module PrebuiltInterface, srcs *[]string) {
@@ -70,13 +80,19 @@
 	p.srcs = srcs
 }
 
+func InitSingleSourcePrebuiltModule(module PrebuiltInterface, src *string) {
+	p := module.Prebuilt()
+	module.AddProperties(&p.properties)
+	p.src = src
+}
+
 type PrebuiltInterface interface {
 	Module
 	Prebuilt() *Prebuilt
 }
 
 func RegisterPrebuiltsPreArchMutators(ctx RegisterMutatorsContext) {
-	ctx.BottomUp("prebuilts", prebuiltMutator).Parallel()
+	ctx.BottomUp("prebuilts", PrebuiltMutator).Parallel()
 }
 
 func RegisterPrebuiltsPostDepsMutators(ctx RegisterMutatorsContext) {
@@ -84,9 +100,9 @@
 	ctx.BottomUp("prebuilt_postdeps", PrebuiltPostDepsMutator).Parallel()
 }
 
-// prebuiltMutator ensures that there is always a module with an undecorated name, and marks
+// PrebuiltMutator ensures that there is always a module with an undecorated name, and marks
 // prebuilt modules that have both a prebuilt and a source module.
-func prebuiltMutator(ctx BottomUpMutatorContext) {
+func PrebuiltMutator(ctx BottomUpMutatorContext) {
 	if m, ok := ctx.Module().(PrebuiltInterface); ok && m.Prebuilt() != nil {
 		p := m.Prebuilt()
 		name := m.base().BaseModuleName()
@@ -104,7 +120,7 @@
 func PrebuiltSelectModuleMutator(ctx TopDownMutatorContext) {
 	if m, ok := ctx.Module().(PrebuiltInterface); ok && m.Prebuilt() != nil {
 		p := m.Prebuilt()
-		if p.srcs == nil {
+		if p.srcs == nil && p.src == nil {
 			panic(fmt.Errorf("prebuilt module did not have InitPrebuiltModule called on it"))
 		}
 		if !p.properties.SourceExists {
@@ -143,7 +159,11 @@
 // usePrebuilt returns true if a prebuilt should be used instead of the source module.  The prebuilt
 // will be used if it is marked "prefer" or if the source module is disabled.
 func (p *Prebuilt) usePrebuilt(ctx TopDownMutatorContext, source Module) bool {
-	if len(*p.srcs) == 0 {
+	if p.srcs != nil && len(*p.srcs) == 0 {
+		return false
+	}
+
+	if p.src != nil && *p.src == "" {
 		return false
 	}
 
diff --git a/apex/apex.go b/apex/apex.go
index c1f52a6..685a774 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -149,6 +149,7 @@
 	android.RegisterModuleType("apex", apexBundleFactory)
 	android.RegisterModuleType("apex_test", testApexBundleFactory)
 	android.RegisterModuleType("apex_defaults", defaultsFactory)
+	android.RegisterModuleType("prebuilt_apex", PrebuiltFactory)
 
 	android.PostDepsMutators(func(ctx android.RegisterMutatorsContext) {
 		ctx.TopDown("apex_deps", apexDepsMutator)
@@ -1283,3 +1284,71 @@
 	android.InitDefaultsModule(module)
 	return module
 }
+
+//
+// Prebuilt APEX
+//
+type Prebuilt struct {
+	android.ModuleBase
+	prebuilt android.Prebuilt
+
+	properties PrebuiltProperties
+
+	inputApex  android.Path
+	installDir android.OutputPath
+}
+
+type PrebuiltProperties struct {
+	// the path to the prebuilt .apex file to import.
+	Src string `android:"arch_variant"`
+
+	// the name of the apex_key module that contains the matching public key to be installed.
+	Key *string
+}
+
+func (p *Prebuilt) DepsMutator(ctx android.BottomUpMutatorContext) {
+	if String(p.properties.Key) == "" {
+		ctx.ModuleErrorf("key is missing")
+		return
+	}
+	ctx.AddDependency(ctx.Module(), keyTag, *p.properties.Key)
+}
+
+func (p *Prebuilt) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	// TODO(jungjw): Check the key validity.
+	p.inputApex = p.prebuilt.SingleSourcePath(ctx)
+	p.installDir = android.PathForModuleInstall(ctx, "apex")
+	ctx.InstallFile(p.installDir, ctx.ModuleName()+imageApexSuffix, p.inputApex)
+}
+
+func (p *Prebuilt) Prebuilt() *android.Prebuilt {
+	return &p.prebuilt
+}
+
+func (p *Prebuilt) Name() string {
+	return p.prebuilt.Name(p.ModuleBase.Name())
+}
+
+func (p *Prebuilt) AndroidMk() android.AndroidMkData {
+	return android.AndroidMkData{
+		Class:      "ETC",
+		OutputFile: android.OptionalPathForPath(p.inputApex),
+		Include:    "$(BUILD_PREBUILT)",
+		Extra: []android.AndroidMkExtraFunc{
+			func(w io.Writer, outputFile android.Path) {
+				fmt.Fprintln(w, "LOCAL_MODULE_PATH :=", filepath.Join("$(OUT_DIR)", p.installDir.RelPathString()))
+				fmt.Fprintln(w, "LOCAL_MODULE_STEM :=", p.BaseModuleName()+imageApexSuffix)
+				fmt.Fprintln(w, "LOCAL_REQUIRED_MODULES :=", String(p.properties.Key))
+			},
+		},
+	}
+}
+
+// prebuilt_apex imports an `.apex` file into the build graph as if it was built with apex.
+func PrebuiltFactory() android.Module {
+	module := &Prebuilt{}
+	module.AddProperties(&module.properties)
+	android.InitSingleSourcePrebuiltModule(module, &module.properties.Src)
+	android.InitAndroidArchModule(module, android.DeviceSupported, android.MultilibCommon)
+	return module
+}
diff --git a/apex/apex_test.go b/apex/apex_test.go
index 8a2e55a..2d9cca6 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -15,6 +15,8 @@
 package apex
 
 import (
+	"bufio"
+	"bytes"
 	"io/ioutil"
 	"os"
 	"strings"
@@ -36,11 +38,14 @@
 	ctx.RegisterModuleType("apex_test", android.ModuleFactoryAdaptor(testApexBundleFactory))
 	ctx.RegisterModuleType("apex_key", android.ModuleFactoryAdaptor(apexKeyFactory))
 	ctx.RegisterModuleType("apex_defaults", android.ModuleFactoryAdaptor(defaultsFactory))
+	ctx.RegisterModuleType("prebuilt_apex", android.ModuleFactoryAdaptor(PrebuiltFactory))
 	ctx.PreArchMutators(android.RegisterDefaultsPreArchMutators)
 
 	ctx.PostDepsMutators(func(ctx android.RegisterMutatorsContext) {
 		ctx.TopDown("apex_deps", apexDepsMutator)
 		ctx.BottomUp("apex", apexMutator)
+		ctx.TopDown("prebuilt_select", android.PrebuiltSelectModuleMutator).Parallel()
+		ctx.BottomUp("prebuilt_postdeps", android.PrebuiltPostDepsMutator).Parallel()
 	})
 
 	ctx.RegisterModuleType("cc_library", android.ModuleFactoryAdaptor(cc.LibraryFactory))
@@ -54,6 +59,9 @@
 	ctx.RegisterModuleType("sh_binary", android.ModuleFactoryAdaptor(android.ShBinaryFactory))
 	ctx.RegisterModuleType("android_app_certificate", android.ModuleFactoryAdaptor(java.AndroidAppCertificateFactory))
 	ctx.RegisterModuleType("filegroup", android.ModuleFactoryAdaptor(android.FileGroupFactory))
+	ctx.PreArchMutators(func(ctx android.RegisterMutatorsContext) {
+		ctx.BottomUp("prebuilts", android.PrebuiltMutator).Parallel()
+	})
 	ctx.PreDepsMutators(func(ctx android.RegisterMutatorsContext) {
 		ctx.BottomUp("image", cc.ImageMutator).Parallel()
 		ctx.BottomUp("link", cc.LinkageMutator).Parallel()
@@ -163,6 +171,7 @@
 		"custom_notice":                        nil,
 		"testkey2.avbpubkey":                   nil,
 		"testkey2.pem":                         nil,
+		"myapex.apex":                          nil,
 	})
 	_, errs := ctx.ParseFileList(".", []string{"Android.bp"})
 	android.FailIfErrored(t, errs)
@@ -1229,3 +1238,42 @@
 		t.Errorf("wrong private key path. expected %q. actual %q", expected_privkey, actual_privkey)
 	}
 }
+
+func TestPrebuilt(t *testing.T) {
+	ctx := testApex(t, `
+		prebuilt_apex {
+			name: "myapex",
+			src: "myapex.apex",
+			key: "myapex.key"
+		}
+
+		apex_key {
+			name: "myapex.key",
+			public_key: "testkey.avbpubkey",
+			private_key: "testkey.pem",
+			product_specific: true,
+		}
+	`)
+
+	prebuilt := ctx.ModuleForTests("myapex", "android_common").Module().(*Prebuilt)
+
+	// Check if the key module is added as a required module.
+	buf := &bytes.Buffer{}
+	prebuilt.AndroidMk().Extra[0](buf, nil)
+	found := false
+	scanner := bufio.NewScanner(bytes.NewReader(buf.Bytes()))
+	expected := "myapex.key"
+	for scanner.Scan() {
+		line := scanner.Text()
+		tok := strings.Split(line, " := ")
+		if tok[0] == "LOCAL_REQUIRED_MODULES" {
+			found = true
+			if tok[1] != "myapex.key" {
+				t.Errorf("Unexpected LOCAL_REQUIRED_MODULES '%s', expected '%s'", tok[1], expected)
+			}
+		}
+	}
+	if !found {
+		t.Errorf("Couldn't find a LOCAL_REQUIRED_MODULES entry")
+	}
+}
diff --git a/apex/key.go b/apex/key.go
index 848e8ce..fbd29bc 100644
--- a/apex/key.go
+++ b/apex/key.go
@@ -93,7 +93,7 @@
 	pubKeyName := m.public_key_file.Base()[0 : len(m.public_key_file.Base())-len(m.public_key_file.Ext())]
 	privKeyName := m.private_key_file.Base()[0 : len(m.private_key_file.Base())-len(m.private_key_file.Ext())]
 
-	if pubKeyName != privKeyName {
+	if m.properties.Public_key != nil && m.properties.Private_key != nil && pubKeyName != privKeyName {
 		ctx.ModuleErrorf("public_key %q (keyname:%q) and private_key %q (keyname:%q) do not have same keyname",
 			m.public_key_file.String(), pubKeyName, m.private_key_file, privKeyName)
 		return