copy JNI from AAR files to android_app APK

JNI libs present in AAR files are ignored in the build. If those
libraries are necessary, they must be manually extracted and a
prebuilt_cc_library module must be created for them. Instead, we can
extract these libraries and copy them to an output APK as part of the build.

Bug: 228848794
Test: created android_app that depends on an android_library_import with
  JNI embedded, and verified that .so libs are present in APK
Test: created android_app that depends on an android_library which
  depends on android_library_import with JNI embedded, and verified
  that .so libs are present in APK
Test: verified that multiple .so's from a single architecture are copied
  to APK
Test: verified that the absence of a jni/<arch> folder or files in the
  folder will cause the Ninja action to error
Change-Id: I242a862d7bd881f7a4a8c23493924d8fe441ea25
diff --git a/java/Android.bp b/java/Android.bp
index e25accf..9df4ab4 100644
--- a/java/Android.bp
+++ b/java/Android.bp
@@ -76,6 +76,7 @@
         "tradefed.go",
     ],
     testSrcs: [
+        "aar_test.go",
         "androidmk_test.go",
         "app_import_test.go",
         "app_set_test.go",
diff --git a/java/aar.go b/java/aar.go
index dadfc6a..5c87b20 100644
--- a/java/aar.go
+++ b/java/aar.go
@@ -570,12 +570,24 @@
 
 	a.exportedProguardFlagFiles = android.FirstUniquePaths(a.exportedProguardFlagFiles)
 	a.exportedStaticPackages = android.FirstUniquePaths(a.exportedStaticPackages)
+
+	prebuiltJniPackages := android.Paths{}
+	ctx.VisitDirectDeps(func(module android.Module) {
+		if info, ok := ctx.OtherModuleProvider(module, JniPackageProvider).(JniPackageInfo); ok {
+			prebuiltJniPackages = append(prebuiltJniPackages, info.JniPackages...)
+		}
+	})
+	if len(prebuiltJniPackages) > 0 {
+		ctx.SetProvider(JniPackageProvider, JniPackageInfo{
+			JniPackages: prebuiltJniPackages,
+		})
+	}
 }
 
 // android_library builds and links sources into a `.jar` file for the device along with Android resources.
 //
 // An android_library has a single variant that produces a `.jar` file containing `.class` files that were
-// compiled against the device bootclasspath, along with a `package-res.apk` file containing  Android resources compiled
+// compiled against the device bootclasspath, along with a `package-res.apk` file containing Android resources compiled
 // with aapt2.  This module is not suitable for installing on a device, but can be used as a `static_libs` dependency of
 // an android_app module.
 func AndroidLibraryFactory() android.Module {
@@ -619,6 +631,10 @@
 	Libs []string
 	// If set to true, run Jetifier against .aar file. Defaults to false.
 	Jetifier *bool
+	// If true, extract JNI libs from AAR archive. These libs will be accessible to android_app modules and
+	// will be passed transitively through android_libraries to an android_app.
+	//TODO(b/241138093) evaluate whether we can have this flag default to true for Bazel conversion
+	Extract_jni *bool
 }
 
 type AARImport struct {
@@ -643,7 +659,8 @@
 
 	hideApexVariantFromMake bool
 
-	aarPath android.Path
+	aarPath     android.Path
+	jniPackages android.Paths
 
 	sdkVersion    android.SdkSpec
 	minSdkVersion android.SdkSpec
@@ -751,6 +768,28 @@
 	ctx.AddVariationDependencies(nil, staticLibTag, a.properties.Static_libs...)
 }
 
+type JniPackageInfo struct {
+	// List of zip files containing JNI libraries
+	// Zip files should have directory structure jni/<arch>/*.so
+	JniPackages android.Paths
+}
+
+var JniPackageProvider = blueprint.NewProvider(JniPackageInfo{})
+
+// Unzip an AAR and extract the JNI libs for $archString.
+var extractJNI = pctx.AndroidStaticRule("extractJNI",
+	blueprint.RuleParams{
+		Command: `rm -rf $out $outDir && touch $out && ` +
+			`unzip -qoDD -d $outDir $in "jni/${archString}/*" && ` +
+			`jni_files=$$(find $outDir/jni -type f) && ` +
+			// print error message if there are no JNI libs for this arch
+			`[ -n "$$jni_files" ] || (echo "ERROR: no JNI libs found for arch ${archString}" && exit 1) && ` +
+			`${config.SoongZipCmd} -o $out -P 'lib/${archString}' ` +
+			`-C $outDir/jni/${archString} $$(echo $$jni_files | xargs -n1 printf " -f %s")`,
+		CommandDeps: []string{"${config.SoongZipCmd}"},
+	},
+	"outDir", "archString")
+
 // Unzip an AAR into its constituent files and directories.  Any files in Outputs that don't exist in the AAR will be
 // touched to create an empty file. The res directory is not extracted, as it will be extracted in its own rule.
 var unzipAAR = pctx.AndroidStaticRule("unzipAAR",
@@ -858,6 +897,31 @@
 		ImplementationAndResourcesJars: android.PathsIfNonNil(a.classpathFile),
 		ImplementationJars:             android.PathsIfNonNil(a.classpathFile),
 	})
+
+	if proptools.Bool(a.properties.Extract_jni) {
+		for _, t := range ctx.MultiTargets() {
+			arch := t.Arch.Abi[0]
+			path := android.PathForModuleOut(ctx, arch+"_jni.zip")
+			a.jniPackages = append(a.jniPackages, path)
+
+			outDir := android.PathForModuleOut(ctx, "aarForJni")
+			aarPath := android.PathForModuleSrc(ctx, a.properties.Aars[0])
+			ctx.Build(pctx, android.BuildParams{
+				Rule:        extractJNI,
+				Input:       aarPath,
+				Outputs:     android.WritablePaths{path},
+				Description: "extract JNI from AAR",
+				Args: map[string]string{
+					"outDir":     outDir.String(),
+					"archString": arch,
+				},
+			})
+		}
+
+		ctx.SetProvider(JniPackageProvider, JniPackageInfo{
+			JniPackages: a.jniPackages,
+		})
+	}
 }
 
 func (a *AARImport) HeaderJars() android.Paths {
@@ -906,6 +970,6 @@
 
 	android.InitPrebuiltModule(module, &module.properties.Aars)
 	android.InitApexModule(module)
-	InitJavaModule(module, android.DeviceSupported)
+	InitJavaModuleMultiTargets(module, android.DeviceSupported)
 	return module
 }
diff --git a/java/aar_test.go b/java/aar_test.go
new file mode 100644
index 0000000..8afa039
--- /dev/null
+++ b/java/aar_test.go
@@ -0,0 +1,83 @@
+// Copyright 2022 Google Inc. All rights reserved.
+//
+// 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 java
+
+import (
+	"android/soong/android"
+	"testing"
+)
+
+func TestAarImportProducesJniPackages(t *testing.T) {
+	ctx := android.GroupFixturePreparers(
+		PrepareForTestWithJavaDefaultModules,
+	).RunTestWithBp(t, `
+		android_library_import {
+			name: "aar-no-jni",
+			aars: ["aary.aar"],
+		}
+		android_library_import {
+			name: "aar-jni",
+			aars: ["aary.aar"],
+			extract_jni: true,
+		}`)
+
+	testCases := []struct {
+		name       string
+		hasPackage bool
+	}{
+		{
+			name:       "aar-no-jni",
+			hasPackage: false,
+		},
+		{
+			name:       "aar-jni",
+			hasPackage: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			appMod := ctx.Module(tc.name, "android_common")
+			appTestMod := ctx.ModuleForTests(tc.name, "android_common")
+
+			info, ok := ctx.ModuleProvider(appMod, JniPackageProvider).(JniPackageInfo)
+			if !ok {
+				t.Errorf("expected android_library_import to have JniPackageProvider")
+			}
+
+			if !tc.hasPackage {
+				if len(info.JniPackages) != 0 {
+					t.Errorf("expected JniPackages to be empty, but got %v", info.JniPackages)
+				}
+				outputFile := "arm64-v8a_jni.zip"
+				jniOutputLibZip := appTestMod.MaybeOutput(outputFile)
+				if jniOutputLibZip.Rule != nil {
+					t.Errorf("did not expect an output file, but found %v", outputFile)
+				}
+				return
+			}
+
+			if len(info.JniPackages) != 1 {
+				t.Errorf("expected a single JniPackage, but got %v", info.JniPackages)
+			}
+
+			outputFile := info.JniPackages[0].String()
+			jniOutputLibZip := appTestMod.Output(outputFile)
+			if jniOutputLibZip.Rule == nil {
+				t.Errorf("did not find output file %v", outputFile)
+			}
+		})
+	}
+}
diff --git a/java/app.go b/java/app.go
index 23a9816..afef334 100755
--- a/java/app.go
+++ b/java/app.go
@@ -472,14 +472,14 @@
 	return a.dexJarFile.PathOrNil()
 }
 
-func (a *AndroidApp) jniBuildActions(jniLibs []jniLib, ctx android.ModuleContext) android.WritablePath {
+func (a *AndroidApp) jniBuildActions(jniLibs []jniLib, prebuiltJniPackages android.Paths, ctx android.ModuleContext) android.WritablePath {
 	var jniJarFile android.WritablePath
-	if len(jniLibs) > 0 {
+	if len(jniLibs) > 0 || len(prebuiltJniPackages) > 0 {
 		a.jniLibs = jniLibs
 		if a.shouldEmbedJnis(ctx) {
 			jniJarFile = android.PathForModuleOut(ctx, "jnilibs.zip")
 			a.installPathForJNISymbols = a.installPath(ctx)
-			TransformJniLibsToJar(ctx, jniJarFile, jniLibs, a.useEmbeddedNativeLibs(ctx))
+			TransformJniLibsToJar(ctx, jniJarFile, jniLibs, prebuiltJniPackages, a.useEmbeddedNativeLibs(ctx))
 			for _, jni := range jniLibs {
 				if jni.coverageFile.Valid() {
 					// Only collect coverage for the first target arch if this is a multilib target.
@@ -623,8 +623,8 @@
 
 	dexJarFile := a.dexBuildActions(ctx)
 
-	jniLibs, certificateDeps := collectAppDeps(ctx, a, a.shouldEmbedJnis(ctx), !Bool(a.appProperties.Jni_uses_platform_apis))
-	jniJarFile := a.jniBuildActions(jniLibs, ctx)
+	jniLibs, prebuiltJniPackages, certificateDeps := collectAppDeps(ctx, a, a.shouldEmbedJnis(ctx), !Bool(a.appProperties.Jni_uses_platform_apis))
+	jniJarFile := a.jniBuildActions(jniLibs, prebuiltJniPackages, ctx)
 
 	if ctx.Failed() {
 		return
@@ -724,9 +724,10 @@
 
 func collectAppDeps(ctx android.ModuleContext, app appDepsInterface,
 	shouldCollectRecursiveNativeDeps bool,
-	checkNativeSdkVersion bool) ([]jniLib, []Certificate) {
+	checkNativeSdkVersion bool) ([]jniLib, android.Paths, []Certificate) {
 
 	var jniLibs []jniLib
+	var prebuiltJniPackages android.Paths
 	var certificates []Certificate
 	seenModulePaths := make(map[string]bool)
 
@@ -775,6 +776,10 @@
 			return shouldCollectRecursiveNativeDeps
 		}
 
+		if info, ok := ctx.OtherModuleProvider(module, JniPackageProvider).(JniPackageInfo); ok {
+			prebuiltJniPackages = append(prebuiltJniPackages, info.JniPackages...)
+		}
+
 		if tag == certificateTag {
 			if dep, ok := module.(*AndroidAppCertificate); ok {
 				certificates = append(certificates, dep.Certificate)
@@ -786,7 +791,7 @@
 		return false
 	})
 
-	return jniLibs, certificates
+	return jniLibs, prebuiltJniPackages, certificates
 }
 
 func (a *AndroidApp) WalkPayloadDeps(ctx android.ModuleContext, do android.PayloadDepsCallback) {
diff --git a/java/app_builder.go b/java/app_builder.go
index 4a18dca..4348644 100644
--- a/java/app_builder.go
+++ b/java/app_builder.go
@@ -218,8 +218,14 @@
 	})
 }
 
-func TransformJniLibsToJar(ctx android.ModuleContext, outputFile android.WritablePath,
-	jniLibs []jniLib, uncompressJNI bool) {
+const jniJarOutputPathString = "jniJarOutput.zip"
+
+func TransformJniLibsToJar(
+	ctx android.ModuleContext,
+	outputFile android.WritablePath,
+	jniLibs []jniLib,
+	prebuiltJniPackages android.Paths,
+	uncompressJNI bool) {
 
 	var deps android.Paths
 	jarArgs := []string{
@@ -245,13 +251,20 @@
 		rule = zipRE
 		args["implicits"] = strings.Join(deps.Strings(), ",")
 	}
+	jniJarPath := android.PathForModuleOut(ctx, jniJarOutputPathString)
 	ctx.Build(pctx, android.BuildParams{
 		Rule:        rule,
 		Description: "zip jni libs",
-		Output:      outputFile,
+		Output:      jniJarPath,
 		Implicits:   deps,
 		Args:        args,
 	})
+	ctx.Build(pctx, android.BuildParams{
+		Rule:        mergeAssetsRule,
+		Description: "merge prebuilt JNI packages",
+		Inputs:      append(prebuiltJniPackages, jniJarPath),
+		Output:      outputFile,
+	})
 }
 
 func targetToJniDir(target android.Target) string {
diff --git a/java/app_import.go b/java/app_import.go
index b017eca..9d199d6 100644
--- a/java/app_import.go
+++ b/java/app_import.go
@@ -256,7 +256,7 @@
 		ctx.ModuleErrorf("One and only one of certficate, presigned, and default_dev_cert properties must be set")
 	}
 
-	_, certificates := collectAppDeps(ctx, a, false, false)
+	_, _, certificates := collectAppDeps(ctx, a, false, false)
 
 	// TODO: LOCAL_EXTRACT_APK/LOCAL_EXTRACT_DPI_APK
 	// TODO: LOCAL_PACKAGE_SPLITS
diff --git a/java/app_test.go b/java/app_test.go
index c4ac4df..bb44803 100644
--- a/java/app_test.go
+++ b/java/app_test.go
@@ -1218,7 +1218,7 @@
 	for _, test := range testCases {
 		t.Run(test.name, func(t *testing.T) {
 			app := ctx.ModuleForTests(test.name, "android_common")
-			jniLibZip := app.Output("jnilibs.zip")
+			jniLibZip := app.Output(jniJarOutputPathString)
 			var abis []string
 			args := strings.Fields(jniLibZip.Args["jarArgs"])
 			for i := 0; i < len(args); i++ {
@@ -1351,7 +1351,7 @@
 	for _, test := range testCases {
 		t.Run(test.name, func(t *testing.T) {
 			app := ctx.ModuleForTests(test.name, "android_common")
-			jniLibZip := app.MaybeOutput("jnilibs.zip")
+			jniLibZip := app.MaybeOutput(jniJarOutputPathString)
 			if g, w := (jniLibZip.Rule != nil), test.packaged; g != w {
 				t.Errorf("expected jni packaged %v, got %v", w, g)
 			}
@@ -1442,7 +1442,7 @@
 		t.Run(test.name, func(t *testing.T) {
 			app := ctx.ModuleForTests(test.name, "android_common")
 
-			jniLibZip := app.MaybeOutput("jnilibs.zip")
+			jniLibZip := app.MaybeOutput(jniJarOutputPathString)
 			if len(jniLibZip.Implicits) != 1 {
 				t.Fatalf("expected exactly one jni library, got %q", jniLibZip.Implicits.Strings())
 			}
@@ -2425,7 +2425,7 @@
 	for _, test := range testCases {
 		t.Run(test.name, func(t *testing.T) {
 			app := ctx.ModuleForTests(test.name, "android_common")
-			jniLibZip := app.Output("jnilibs.zip")
+			jniLibZip := app.Output(jniJarOutputPathString)
 			var jnis []string
 			args := strings.Fields(jniLibZip.Args["jarArgs"])
 			for i := 0; i < len(args); i++ {
@@ -3074,3 +3074,89 @@
 	}
 	android.AssertStringDoesContain(t, "expected error rule message", fooApk.Args["error"], "missing dependencies: missing_certificate\n")
 }
+
+func TestAppIncludesJniPackages(t *testing.T) {
+	ctx := android.GroupFixturePreparers(
+		PrepareForTestWithJavaDefaultModules,
+	).RunTestWithBp(t, `
+		android_library_import {
+			name: "aary-nodeps",
+			aars: ["aary.aar"],
+			extract_jni: true,
+		}
+
+		android_library {
+			name: "aary-lib",
+			sdk_version: "current",
+			min_sdk_version: "21",
+			static_libs: ["aary-nodeps"],
+		}
+
+		android_app {
+			name: "aary-lib-dep",
+			sdk_version: "current",
+			min_sdk_version: "21",
+			manifest: "AndroidManifest.xml",
+			static_libs: ["aary-lib"],
+			use_embedded_native_libs: true,
+		}
+
+		android_app {
+			name: "aary-import-dep",
+			sdk_version: "current",
+			min_sdk_version: "21",
+			manifest: "AndroidManifest.xml",
+			static_libs: ["aary-nodeps"],
+			use_embedded_native_libs: true,
+		}
+
+		android_app {
+			name: "aary-no-use-embedded",
+			sdk_version: "current",
+			min_sdk_version: "21",
+			manifest: "AndroidManifest.xml",
+			static_libs: ["aary-nodeps"],
+		}`)
+
+	testCases := []struct {
+		name       string
+		hasPackage bool
+	}{
+		{
+			name:       "aary-import-dep",
+			hasPackage: true,
+		},
+		{
+			name:       "aary-lib-dep",
+			hasPackage: true,
+		},
+		{
+			name:       "aary-no-use-embedded",
+			hasPackage: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			app := ctx.ModuleForTests(tc.name, "android_common")
+
+			outputFile := "jnilibs.zip"
+			jniOutputLibZip := app.MaybeOutput(outputFile)
+			if jniOutputLibZip.Rule == nil && !tc.hasPackage {
+				return
+			}
+
+			jniPackage := "arm64-v8a_jni.zip"
+			inputs := jniOutputLibZip.Inputs
+			foundPackage := false
+			for i := 0; i < len(inputs); i++ {
+				if strings.Contains(inputs[i].String(), jniPackage) {
+					foundPackage = true
+				}
+			}
+			if foundPackage != tc.hasPackage {
+				t.Errorf("expected to find %v in %v inputs; inputs = %v", jniPackage, outputFile, inputs)
+			}
+		})
+	}
+}
diff --git a/java/rro.go b/java/rro.go
index 7952c2c..9c8c53b 100644
--- a/java/rro.go
+++ b/java/rro.go
@@ -142,7 +142,7 @@
 	r.aapt.buildActions(ctx, r, nil, nil, aaptLinkFlags...)
 
 	// Sign the built package
-	_, certificates := collectAppDeps(ctx, r, false, false)
+	_, _, certificates := collectAppDeps(ctx, r, false, false)
 	certificates = processMainCert(r.ModuleBase, String(r.properties.Certificate), certificates, ctx)
 	signed := android.PathForModuleOut(ctx, "signed", r.Name()+".apk")
 	var lineageFile android.Path