diff --git a/apex/apex.go b/apex/apex.go
index 3b06a99..408415e 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -379,6 +379,13 @@
 	outputFiles      map[apexPackaging]android.WritablePath
 	installDir       android.OutputPath
 
+	public_key_file   android.Path
+	private_key_file  android.Path
+	bundle_public_key bool
+
+	container_certificate_file android.Path
+	container_private_key_file android.Path
+
 	// list of files to be included in this apex
 	filesInfo []apexFile
 
@@ -635,10 +642,6 @@
 func (a *apexBundle) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	filesInfo := []apexFile{}
 
-	var keyFile android.Path
-	var pubKeyFile android.Path
-	var certificate java.Certificate
-
 	if a.properties.Payload_type == nil || *a.properties.Payload_type == "image" {
 		a.apexTypes = imageApex
 	} else if *a.properties.Payload_type == "zip" {
@@ -704,20 +707,20 @@
 				}
 			case keyTag:
 				if key, ok := child.(*apexKey); ok {
-					keyFile = key.private_key_file
-					if !key.installable() && ctx.Config().Debuggable() {
-						// If the key is not installed, bundled it with the APEX.
-						// Note: this bundled key is valid only for non-production builds
-						// (eng/userdebug).
-						pubKeyFile = key.public_key_file
-					}
+					a.private_key_file = key.private_key_file
+					a.public_key_file = key.public_key_file
+					// If the key is not installed, bundled it with the APEX.
+					// Note: this bundled key is valid only for non-production builds
+					// (eng/userdebug).
+					a.bundle_public_key = !key.installable() && ctx.Config().Debuggable()
 					return false
 				} else {
 					ctx.PropertyErrorf("key", "%q is not an apex_key module", depName)
 				}
 			case certificateTag:
 				if dep, ok := child.(*java.AndroidAppCertificate); ok {
-					certificate = dep.Certificate
+					a.container_certificate_file = dep.Certificate.Pem
+					a.container_private_key_file = dep.Certificate.Key
 					return false
 				} else {
 					ctx.ModuleErrorf("certificate dependency %q must be an android_app_certificate module", depName)
@@ -741,7 +744,7 @@
 	})
 
 	a.flattened = ctx.Config().FlattenApex() && !ctx.Config().UnbundledBuild()
-	if keyFile == nil {
+	if a.private_key_file == nil {
 		ctx.PropertyErrorf("key", "private_key for %q could not be found", String(a.properties.Key))
 		return
 	}
@@ -775,30 +778,28 @@
 	a.filesInfo = filesInfo
 
 	if a.apexTypes.zip() {
-		a.buildUnflattenedApex(ctx, keyFile, pubKeyFile, certificate, zipApex)
+		a.buildUnflattenedApex(ctx, zipApex)
 	}
 	if a.apexTypes.image() {
 		// Build rule for unflattened APEX is created even when ctx.Config().FlattenApex()
 		// is true. This is to support referencing APEX via ":<module_name" syntax
 		// in other modules. It is in AndroidMk where the selection of flattened
 		// or unflattened APEX is made.
-		a.buildUnflattenedApex(ctx, keyFile, pubKeyFile, certificate, imageApex)
+		a.buildUnflattenedApex(ctx, imageApex)
 		a.buildFlattenedApex(ctx)
 	}
 }
 
-func (a *apexBundle) buildUnflattenedApex(ctx android.ModuleContext, keyFile android.Path,
-	pubKeyFile android.Path, certificate java.Certificate, apexType apexPackaging) {
+func (a *apexBundle) buildUnflattenedApex(ctx android.ModuleContext, apexType apexPackaging) {
 	cert := String(a.properties.Certificate)
 	if cert != "" && android.SrcIsModule(cert) == "" {
 		defaultDir := ctx.Config().DefaultAppCertificateDir(ctx)
-		certificate = java.Certificate{
-			defaultDir.Join(ctx, cert+".x509.pem"),
-			defaultDir.Join(ctx, cert+".pk8"),
-		}
+		a.container_certificate_file = defaultDir.Join(ctx, cert+".x509.pem")
+		a.container_private_key_file = defaultDir.Join(ctx, cert+".pk8")
 	} else if cert == "" {
 		pem, key := ctx.Config().DefaultAppCertificate(ctx)
-		certificate = java.Certificate{pem, key}
+		a.container_certificate_file = pem
+		a.container_private_key_file = key
 	}
 
 	manifest := ctx.ExpandSource(proptools.StringDefault(a.properties.Manifest, "apex_manifest.json"), "manifest")
@@ -886,10 +887,10 @@
 		optFlags := []string{}
 
 		// Additional implicit inputs.
-		implicitInputs = append(implicitInputs, cannedFsConfig, fileContexts, keyFile)
-		if pubKeyFile != nil {
-			implicitInputs = append(implicitInputs, pubKeyFile)
-			optFlags = append(optFlags, "--pubkey "+pubKeyFile.String())
+		implicitInputs = append(implicitInputs, cannedFsConfig, fileContexts, a.private_key_file)
+		if a.bundle_public_key {
+			implicitInputs = append(implicitInputs, a.public_key_file)
+			optFlags = append(optFlags, "--pubkey "+a.public_key_file.String())
 		}
 
 		manifestPackageName, overridden := ctx.DeviceConfig().OverrideManifestPackageNameFor(ctx.ModuleName())
@@ -915,7 +916,7 @@
 				"manifest":         manifest.String(),
 				"file_contexts":    fileContexts.String(),
 				"canned_fs_config": cannedFsConfig.String(),
-				"key":              keyFile.String(),
+				"key":              a.private_key_file.String(),
 				"opt_flags":        strings.Join(optFlags, " "),
 			},
 		})
@@ -962,14 +963,14 @@
 		Output:      a.outputFiles[apexType],
 		Input:       unsignedOutputFile,
 		Args: map[string]string{
-			"certificates": strings.Join([]string{certificate.Pem.String(), certificate.Key.String()}, " "),
+			"certificates": a.container_certificate_file.String() + " " + a.container_private_key_file.String(),
 			"flags":        "-a 4096", //alignment
 		},
 	})
 
 	// Install to $OUT/soong/{target,host}/.../apex
 	if a.installable() && (!ctx.Config().FlattenApex() || apexType.zip()) {
-		ctx.InstallFile(android.PathForModuleInstall(ctx, "apex"), ctx.ModuleName()+suffix, a.outputFiles[apexType])
+		ctx.InstallFile(a.installDir, ctx.ModuleName()+suffix, a.outputFiles[apexType])
 	}
 }