Handle split APKs

If no APK path is set, but we know the binary payload name, search for
a split APK containing the payload. If there aren't any splits, or we
don't find it, fall back to the primary APK as previously.

Bug: 247116743
Test: atest MicrodroidTests
Change-Id: I5b28b61500a3fe3564e637c76bb52a4d197a1ff6
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 668d7dc..93e65db 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -31,11 +31,13 @@
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
+import android.os.Build;
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
 import android.sysprop.HypervisorProperties;
 import android.system.virtualizationservice.VirtualMachineAppConfig;
 import android.system.virtualizationservice.VirtualMachinePayloadConfig;
+import android.util.Log;
 
 import java.io.File;
 import java.io.FileInputStream;
@@ -47,6 +49,7 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.Objects;
+import java.util.zip.ZipFile;
 
 /**
  * Represents a configuration of a virtual machine. A configuration consists of hardware
@@ -57,6 +60,7 @@
  */
 @SystemApi
 public final class VirtualMachineConfig {
+    private static final String TAG = "VirtualMachineConfig";
     private static final String[] EMPTY_STRING_ARRAY = {};
 
     // These define the schema of the config file persisted on disk.
@@ -296,8 +300,7 @@
 
     /**
      * Returns the absolute path of the APK which should contain the binary payload that will
-     * execute within the VM. Returns null if no specific path has been set, so the primary APK will
-     * be used.
+     * execute within the VM. Returns null if no specific path has been set.
      *
      * @hide
      */
@@ -445,18 +448,7 @@
             throws VirtualMachineException {
         VirtualMachineAppConfig vsConfig = new VirtualMachineAppConfig();
 
-        String apkPath = mApkPath;
-        if (apkPath == null) {
-            try {
-                ApplicationInfo appInfo =
-                        packageManager.getApplicationInfo(
-                                mPackageName, PackageManager.ApplicationInfoFlags.of(0));
-                // This really is the path to the APK, not a directory.
-                apkPath = appInfo.sourceDir;
-            } catch (PackageManager.NameNotFoundException e) {
-                throw new VirtualMachineException("Package not found", e);
-            }
-        }
+        String apkPath = (mApkPath != null) ? mApkPath : findPayloadApk(packageManager);
 
         try {
             vsConfig.apk = ParcelFileDescriptor.open(new File(apkPath), MODE_READ_ONLY);
@@ -495,6 +487,45 @@
         return vsConfig;
     }
 
+    private String findPayloadApk(PackageManager packageManager) throws VirtualMachineException {
+        ApplicationInfo appInfo;
+        try {
+            appInfo =
+                    packageManager.getApplicationInfo(
+                            mPackageName, PackageManager.ApplicationInfoFlags.of(0));
+        } catch (PackageManager.NameNotFoundException e) {
+            throw new VirtualMachineException("Package not found", e);
+        }
+
+        String[] splitApkPaths = appInfo.splitSourceDirs;
+        String[] abis = Build.SUPPORTED_64_BIT_ABIS;
+
+        // If there are split APKs, and we know the payload binary name, see if we can find a
+        // split APK containing the binary.
+        if (mPayloadBinaryName != null && splitApkPaths != null && abis.length != 0) {
+            String[] libraryNames = new String[abis.length];
+            for (int i = 0; i < abis.length; i++) {
+                libraryNames[i] = "lib/" + abis[i] + "/" + mPayloadBinaryName;
+            }
+
+            for (String path : splitApkPaths) {
+                try (ZipFile zip = new ZipFile(path)) {
+                    for (String name : libraryNames) {
+                        if (zip.getEntry(name) != null) {
+                            Log.i(TAG, "Found payload in " + path);
+                            return path;
+                        }
+                    }
+                } catch (IOException e) {
+                    Log.w(TAG, "Failed to scan split APK: " + path, e);
+                }
+            }
+        }
+
+        // This really is the path to the APK, not a directory.
+        return appInfo.sourceDir;
+    }
+
     private int bytesToMebiBytes(long mMemoryBytes) {
         long oneMebi = 1024 * 1024;
         // We can't express requests for more than 2 exabytes, but then they're not going to succeed
@@ -596,7 +627,8 @@
 
         /**
          * Sets the absolute path of the APK containing the binary payload that will execute within
-         * the VM. If not set explicitly, defaults to the primary APK of the context.
+         * the VM. If not set explicitly, defaults to the split APK containing the payload, if there
+         * is one, and otherwise the primary APK of the context.
          *
          * @hide
          */