Add test API for extra APKs

Bug: 303201498
Test: atest MicrodroidTests
Change-Id: Iaae9274d704c6cbf47be31902820872c996a701e
diff --git a/javalib/api/test-current.txt b/javalib/api/test-current.txt
index 958005f..5aff93f 100644
--- a/javalib/api/test-current.txt
+++ b/javalib/api/test-current.txt
@@ -7,12 +7,14 @@
   }
 
   public final class VirtualMachineConfig {
+    method @FlaggedApi("RELEASE_AVF_ENABLE_MULTI_TENANT_MICRODROID_VM") @NonNull public java.util.List<java.lang.String> getExtraApks();
     method @FlaggedApi("RELEASE_AVF_ENABLE_VENDOR_MODULES") @Nullable public String getOs();
     method @Nullable public String getPayloadConfigPath();
     method public boolean isVmConsoleInputSupported();
   }
 
   public static final class VirtualMachineConfig.Builder {
+    method @FlaggedApi("RELEASE_AVF_ENABLE_MULTI_TENANT_MICRODROID_VM") @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder addExtraApk(@NonNull String);
     method @FlaggedApi("RELEASE_AVF_ENABLE_VENDOR_MODULES") @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setOs(@NonNull String);
     method @NonNull @RequiresPermission(android.system.virtualmachine.VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION) public android.system.virtualmachine.VirtualMachineConfig.Builder setPayloadConfigPath(@NonNull String);
     method @FlaggedApi("RELEASE_AVF_ENABLE_VENDOR_MODULES") @NonNull @RequiresPermission(android.system.virtualmachine.VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION) public android.system.virtualmachine.VirtualMachineConfig.Builder setVendorDiskImage(@NonNull java.io.File);
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 16f9631..5025e88 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -54,6 +54,8 @@
 import android.annotation.WorkerThread;
 import android.content.ComponentCallbacks2;
 import android.content.Context;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
 import android.content.res.Configuration;
 import android.os.Binder;
 import android.os.IBinder;
@@ -76,7 +78,6 @@
 
 import java.io.File;
 import java.io.FileInputStream;
-import java.io.FileOutputStream;
 import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
 import java.io.IOException;
@@ -807,10 +808,30 @@
                     createVmInputPipes();
                 }
 
+                VirtualMachineConfig vmConfig = getConfig();
                 VirtualMachineAppConfig appConfig =
-                        getConfig().toVsConfig(mContext.getPackageManager());
+                        vmConfig.toVsConfig(mContext.getPackageManager());
                 appConfig.name = mName;
 
+                if (!vmConfig.getExtraApks().isEmpty()) {
+                    // Extra APKs were specified directly, rather than via config file.
+                    // We've already populated the file names for the extra APKs and IDSigs
+                    // (via setupExtraApks). But we also need to open the APK files and add
+                    // fds for them to the payload config.
+                    // This isn't needed when the extra APKs are specified in a config file - then
+                    // Virtualization Manager opens them itself.
+                    List<ParcelFileDescriptor> extraApkFiles = new ArrayList<>(mExtraApks.size());
+                    for (ExtraApkSpec extraApk : mExtraApks) {
+                        try {
+                            extraApkFiles.add(
+                                    ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY));
+                        } catch (FileNotFoundException e) {
+                            throw new VirtualMachineException("Failed to open extra APK", e);
+                        }
+                    }
+                    appConfig.payload.getPayloadConfig().extraApks = extraApkFiles;
+                }
+
                 try {
                     createIdSigs(service, appConfig);
                 } catch (FileNotFoundException e) {
@@ -1239,6 +1260,46 @@
         return result.toString();
     }
 
+    /**
+     * Reads the payload config inside the application, parses extra APK information, and then
+     * creates corresponding idsig file paths.
+     */
+    private static List<ExtraApkSpec> setupExtraApks(
+            @NonNull Context context, @NonNull VirtualMachineConfig config, @NonNull File vmDir)
+            throws VirtualMachineException {
+        String configPath = config.getPayloadConfigPath();
+        List<String> extraApks = config.getExtraApks();
+        if (configPath != null) {
+            return setupExtraApksFromConfigFile(context, vmDir, configPath);
+        } else if (!extraApks.isEmpty()) {
+            return setupExtraApksFromList(context, vmDir, extraApks);
+        } else {
+            return Collections.emptyList();
+        }
+    }
+
+    private static List<ExtraApkSpec> setupExtraApksFromConfigFile(
+            Context context, File vmDir, String configPath) throws VirtualMachineException {
+        try (ZipFile zipFile = new ZipFile(context.getPackageCodePath())) {
+            InputStream inputStream = zipFile.getInputStream(zipFile.getEntry(configPath));
+            List<String> apkList =
+                    parseExtraApkListFromPayloadConfig(
+                            new JsonReader(new InputStreamReader(inputStream)));
+
+            List<ExtraApkSpec> extraApks = new ArrayList<>(apkList.size());
+            for (int i = 0; i < apkList.size(); ++i) {
+                extraApks.add(
+                        new ExtraApkSpec(
+                                new File(apkList.get(i)),
+                                new File(vmDir, EXTRA_IDSIG_FILE_PREFIX + i)));
+            }
+
+            return extraApks;
+        } catch (IOException e) {
+            throw new VirtualMachineException("Couldn't parse extra apks from the vm config", e);
+        }
+    }
+
     private static List<String> parseExtraApkListFromPayloadConfig(JsonReader reader)
             throws VirtualMachineException {
         /*
@@ -1275,36 +1336,28 @@
         }
     }
 
-    /**
-     * Reads the payload config inside the application, parses extra APK information, and then
-     * creates corresponding idsig file paths.
-     */
-    private static List<ExtraApkSpec> setupExtraApks(
-            @NonNull Context context, @NonNull VirtualMachineConfig config, @NonNull File vmDir)
-            throws VirtualMachineException {
-        String configPath = config.getPayloadConfigPath();
-        if (configPath == null) {
-            return Collections.emptyList();
-        }
-        try (ZipFile zipFile = new ZipFile(context.getPackageCodePath())) {
-            InputStream inputStream =
-                    zipFile.getInputStream(zipFile.getEntry(configPath));
-            List<String> apkList =
-                    parseExtraApkListFromPayloadConfig(
-                            new JsonReader(new InputStreamReader(inputStream)));
-
-            List<ExtraApkSpec> extraApks = new ArrayList<>();
-            for (int i = 0; i < apkList.size(); ++i) {
-                extraApks.add(
-                        new ExtraApkSpec(
-                                new File(apkList.get(i)),
-                                new File(vmDir, EXTRA_IDSIG_FILE_PREFIX + i)));
+    private static List<ExtraApkSpec> setupExtraApksFromList(
+            Context context, File vmDir, List<String> extraApkInfo) throws VirtualMachineException {
+        int count = extraApkInfo.size();
+        List<ExtraApkSpec> extraApks = new ArrayList<>(count);
+        for (int i = 0; i < count; i++) {
+            String packageName = extraApkInfo.get(i);
+            ApplicationInfo appInfo;
+            try {
+                appInfo =
+                        context.getPackageManager()
+                                .getApplicationInfo(
+                                        packageName, PackageManager.ApplicationInfoFlags.of(0));
+            } catch (PackageManager.NameNotFoundException e) {
+                throw new VirtualMachineException("Extra APK package not found", e);
             }
 
-            return Collections.unmodifiableList(extraApks);
-        } catch (IOException e) {
-            throw new VirtualMachineException("Couldn't parse extra apks from the vm config", e);
+            extraApks.add(
+                    new ExtraApkSpec(
+                            new File(appInfo.sourceDir),
+                            new File(vmDir, EXTRA_IDSIG_FILE_PREFIX + i)));
         }
+        return extraApks;
     }
 
     private void importInstanceFrom(@NonNull ParcelFileDescriptor instanceFd)
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index e8ef195..9688789 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -49,7 +49,10 @@
 import java.io.OutputStream;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
 import java.util.Objects;
 import java.util.zip.ZipFile;
 
@@ -67,7 +70,7 @@
     private static String[] EMPTY_STRING_ARRAY = {};
 
     // These define the schema of the config file persisted on disk.
-    private static final int VERSION = 7;
+    private static final int VERSION = 8;
     private static final String KEY_VERSION = "version";
     private static final String KEY_PACKAGENAME = "packageName";
     private static final String KEY_APKPATH = "apkPath";
@@ -82,6 +85,7 @@
     private static final String KEY_VM_CONSOLE_INPUT_SUPPORTED = "vmConsoleInputSupported";
     private static final String KEY_VENDOR_DISK_IMAGE_PATH = "vendorDiskImagePath";
     private static final String KEY_OS = "os";
+    private static final String KEY_EXTRA_APKS = "extraApks";
 
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
@@ -140,6 +144,8 @@
     /** Absolute path to the APK file containing the VM payload. */
     @Nullable private final String mApkPath;
 
+    private final List<String> mExtraApks;
+
     @DebugLevel private final int mDebugLevel;
 
     /**
@@ -181,6 +187,7 @@
     private VirtualMachineConfig(
             @Nullable String packageName,
             @Nullable String apkPath,
+            List<String> extraApks,
             @Nullable String payloadConfigPath,
             @Nullable String payloadBinaryName,
             @DebugLevel int debugLevel,
@@ -195,6 +202,11 @@
         // This is only called from Builder.build(); the builder handles parameter validation.
         mPackageName = packageName;
         mApkPath = apkPath;
+        mExtraApks =
+                extraApks.isEmpty()
+                        ? Collections.emptyList()
+                        : Collections.unmodifiableList(
+                                Arrays.asList(extraApks.toArray(new String[0])));
         mPayloadConfigPath = payloadConfigPath;
         mPayloadBinaryName = payloadBinaryName;
         mDebugLevel = debugLevel;
@@ -292,6 +304,13 @@
             builder.setOs(os);
         }
 
+        String[] extraApks = b.getStringArray(KEY_EXTRA_APKS);
+        if (extraApks != null) {
+            for (String extraApk : extraApks) {
+                builder.addExtraApk(extraApk);
+            }
+        }
+
         return builder.build();
     }
 
@@ -331,6 +350,10 @@
             b.putString(KEY_VENDOR_DISK_IMAGE_PATH, mVendorDiskImage.getAbsolutePath());
         }
         b.putString(KEY_OS, mOs);
+        if (!mExtraApks.isEmpty()) {
+            String[] extraApks = mExtraApks.toArray(new String[0]);
+            b.putStringArray(KEY_EXTRA_APKS, extraApks);
+        }
         b.writeToStream(output);
     }
 
@@ -347,6 +370,19 @@
     }
 
     /**
+     * Returns the package names of any extra APKs that have been requested for the VM. They are
+     * returned in the order in which they were added via {@link Builder#addExtraApk}.
+     *
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi("RELEASE_AVF_ENABLE_MULTI_TENANT_MICRODROID_VM")
+    @NonNull
+    public List<String> getExtraApks() {
+        return mExtraApks;
+    }
+
+    /**
      * Returns the path within the APK to the payload config file that defines software aspects of
      * the VM.
      *
@@ -495,7 +531,8 @@
                 && Objects.equals(this.mPayloadConfigPath, other.mPayloadConfigPath)
                 && Objects.equals(this.mPayloadBinaryName, other.mPayloadBinaryName)
                 && Objects.equals(this.mPackageName, other.mPackageName)
-                && Objects.equals(this.mOs, other.mOs);
+                && Objects.equals(this.mOs, other.mOs)
+                && Objects.equals(this.mExtraApks, other.mExtraApks);
     }
 
     /**
@@ -623,6 +660,7 @@
 
         @Nullable private final String mPackageName;
         @Nullable private String mApkPath;
+        private final List<String> mExtraApks = new ArrayList<>();
         @Nullable private String mPayloadConfigPath;
         @Nullable private String mPayloadBinaryName;
         @DebugLevel private int mDebugLevel = DEBUG_LEVEL_NONE;
@@ -683,6 +721,10 @@
                     throw new IllegalStateException(
                             "setPayloadConfigPath and setOs may not both be called");
                 }
+                if (!mExtraApks.isEmpty()) {
+                    throw new IllegalStateException(
+                            "setPayloadConfigPath and addExtraApk may not both be called");
+                }
             } else {
                 if (mPayloadConfigPath != null) {
                     throw new IllegalStateException(
@@ -710,6 +752,7 @@
             return new VirtualMachineConfig(
                     packageName,
                     apkPath,
+                    mExtraApks,
                     mPayloadConfigPath,
                     mPayloadBinaryName,
                     mDebugLevel,
@@ -742,6 +785,21 @@
         }
 
         /**
+         * Specify the package name of an extra APK to be included in the VM. Each extra APK is
+         * mounted, in unzipped form, inside the VM, allowing access to the code and/or data within
+         * it. The VM entry point must be in the main APK.
+         *
+         * @hide
+         */
+        @TestApi
+        @FlaggedApi("RELEASE_AVF_ENABLE_MULTI_TENANT_MICRODROID_VM")
+        @NonNull
+        public Builder addExtraApk(@NonNull String packageName) {
+            mExtraApks.add(requireNonNull(packageName, "extra APK package name must not be null"));
+            return this;
+        }
+
+        /**
          * Sets the path within the APK to the payload config file that defines software aspects of
          * the VM. The file is a JSON file; see
          * packages/modules/Virtualization/microdroid/payload/config/src/lib.rs for the format.