Add a parameter for extra apks to payload config

Users can pass extra apks and corresponding idsigs to the VM, by setting
extra_apks property to the payload config.

Bug: 205224817
Test: add extra_apks, vm run-app, see /dev/block/by-name/extra-apk-0 and
/dev/block/by-name/extra-idsig-0

Change-Id: I7908788a163d6ff7b90bb008fc326eb23d1611bb
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 6556b87..d04da0e 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -33,6 +33,7 @@
 import android.system.virtualizationservice.PartitionType;
 import android.system.virtualizationservice.VirtualMachineAppConfig;
 import android.system.virtualizationservice.VirtualMachineState;
+import android.util.JsonReader;
 
 import java.io.File;
 import java.io.FileInputStream;
@@ -40,13 +41,17 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.zip.ZipFile;
 
 /**
  * A handle to the virtual machine. The virtual machine is local to the app which created the
@@ -67,6 +72,9 @@
     /** Name of the idsig file for a VM */
     private static final String IDSIG_FILE = "idsig";
 
+    /** Name of the idsig files for extra APKs. */
+    private static final String EXTRA_IDSIG_FILE_PREFIX = "extra_idsig_";
+
     /** Name of the virtualization service. */
     private static final String SERVICE_NAME = "android.system.virtualizationservice";
 
@@ -100,6 +108,22 @@
     /** Path to the idsig file for this VM. */
     private final @NonNull File mIdsigFilePath;
 
+    private static class ExtraApkSpec {
+        public final File apk;
+        public final File idsig;
+
+        ExtraApkSpec(File apk, File idsig) {
+            this.apk = apk;
+            this.idsig = idsig;
+        }
+    }
+
+    /**
+     * List of extra apks. Apks are specified by the vm config, and corresponding idsigs are to be
+     * generated.
+     */
+    private final @NonNull List<ExtraApkSpec> mExtraApks;
+
     /** Size of the instance image. 10 MB. */
     private static final long INSTANCE_FILE_SIZE = 10 * 1024 * 1024;
 
@@ -128,16 +152,18 @@
     }
 
     private VirtualMachine(
-            @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config) {
+            @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
+            throws VirtualMachineException {
         mPackageName = context.getPackageName();
         mName = name;
         mConfig = config;
+        mConfigFilePath = getConfigFilePath(context, name);
 
         final File vmRoot = new File(context.getFilesDir(), VM_DIR);
         final File thisVmDir = new File(vmRoot, mName);
-        mConfigFilePath = new File(thisVmDir, CONFIG_FILE);
         mInstanceFilePath = new File(thisVmDir, INSTANCE_IMAGE_FILE);
         mIdsigFilePath = new File(thisVmDir, IDSIG_FILE);
+        mExtraApks = setupExtraApks(context, config, thisVmDir);
     }
 
     /**
@@ -198,11 +224,10 @@
     /** Loads a virtual machine that is already created before. */
     /* package */ static @NonNull VirtualMachine load(
             @NonNull Context context, @NonNull String name) throws VirtualMachineException {
-        VirtualMachine vm = new VirtualMachine(context, name, /* config */ null);
-
-        try (FileInputStream input = new FileInputStream(vm.mConfigFilePath)) {
-            VirtualMachineConfig config = VirtualMachineConfig.from(input);
-            vm.mConfig = config;
+        File configFilePath = getConfigFilePath(context, name);
+        VirtualMachineConfig config;
+        try (FileInputStream input = new FileInputStream(configFilePath)) {
+            config = VirtualMachineConfig.from(input);
         } catch (FileNotFoundException e) {
             // The VM doesn't exist.
             return null;
@@ -210,6 +235,8 @@
             throw new VirtualMachineException(e);
         }
 
+        VirtualMachine vm = new VirtualMachine(context, name, config);
+
         // If config file exists, but the instance image file doesn't, it means that the VM is
         // corrupted. That's different from the case that the VM doesn't exist. Throw an exception
         // instead of returning null.
@@ -292,6 +319,9 @@
 
         try {
             mIdsigFilePath.createNewFile();
+            for (ExtraApkSpec extraApk : mExtraApks) {
+                extraApk.idsig.createNewFile();
+            }
         } catch (IOException e) {
             // If the file already exists, exception is not thrown.
             throw new VirtualMachineException("failed to create idsig file", e);
@@ -320,9 +350,20 @@
             service.createOrUpdateIdsigFile(
                     appConfig.apk, ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_WRITE));
 
+            for (ExtraApkSpec extraApk : mExtraApks) {
+                service.createOrUpdateIdsigFile(
+                        ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY),
+                        ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_WRITE));
+            }
+
             // Re-open idsig file in read-only mode
             appConfig.idsig = ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_ONLY);
             appConfig.instanceImage = ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_WRITE);
+            List<ParcelFileDescriptor> extraIdsigs = new ArrayList<>();
+            for (ExtraApkSpec extraApk : mExtraApks) {
+                extraIdsigs.add(ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_ONLY));
+            }
+            appConfig.extraIdsigs = extraIdsigs;
 
             android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
                     android.system.virtualizationservice.VirtualMachineConfig.appConfig(appConfig);
@@ -426,6 +467,9 @@
             throw new VirtualMachineException("Virtual machine is not stopped");
         }
         final File vmRootDir = mConfigFilePath.getParentFile();
+        for (ExtraApkSpec extraApks : mExtraApks) {
+            extraApks.idsig.delete();
+        }
         mConfigFilePath.delete();
         mInstanceFilePath.delete();
         mIdsigFilePath.delete();
@@ -507,4 +551,76 @@
         sb.append(")");
         return sb.toString();
     }
+
+    private static List<String> parseExtraApkListFromPayloadConfig(JsonReader reader)
+            throws VirtualMachineException {
+        /**
+         * JSON schema from packages/modules/Virtualization/microdroid/payload/config/src/lib.rs:
+         *
+         * <p>{ "extra_apks": [ { "path": "/system/app/foo.apk", }, ... ], ... }
+         */
+        try {
+            List<String> apks = new ArrayList<>();
+
+            reader.beginObject();
+            while (reader.hasNext()) {
+                if (reader.nextName().equals("extra_apks")) {
+                    reader.beginArray();
+                    while (reader.hasNext()) {
+                        reader.beginObject();
+                        String name = reader.nextName();
+                        if (name.equals("path")) {
+                            apks.add(reader.nextString());
+                        } else {
+                            reader.skipValue();
+                        }
+                        reader.endObject();
+                    }
+                    reader.endArray();
+                } else {
+                    reader.skipValue();
+                }
+            }
+            reader.endObject();
+            return apks;
+        } catch (IOException e) {
+            throw new VirtualMachineException(e);
+        }
+    }
+
+    /**
+     * 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 {
+        try {
+            ZipFile zipFile = new ZipFile(context.getPackageCodePath());
+            String payloadPath = config.getPayloadConfigPath();
+            InputStream inputStream =
+                    zipFile.getInputStream(zipFile.getEntry(config.getPayloadConfigPath()));
+            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)));
+            }
+
+            return extraApks;
+        } catch (IOException e) {
+            throw new VirtualMachineException("Couldn't parse extra apks from the vm config", e);
+        }
+    }
+
+    private static File getConfigFilePath(@NonNull Context context, @NonNull String name) {
+        final File vmRoot = new File(context.getFilesDir(), VM_DIR);
+        final File thisVmDir = new File(vmRoot, name);
+        return new File(thisVmDir, CONFIG_FILE);
+    }
 }
diff --git a/microdroid/payload/config/src/lib.rs b/microdroid/payload/config/src/lib.rs
index 2547f3d..67e8feb 100644
--- a/microdroid/payload/config/src/lib.rs
+++ b/microdroid/payload/config/src/lib.rs
@@ -31,6 +31,10 @@
     #[serde(default)]
     pub apexes: Vec<ApexConfig>,
 
+    /// Extra APKs to be passed to a VM
+    #[serde(default)]
+    pub extra_apks: Vec<ApkConfig>,
+
     /// Tells VirtualizationService to use staged APEXes if possible
     #[serde(default)]
     pub prefer_staged: bool,
@@ -91,3 +95,10 @@
     /// The name of APEX
     pub name: String,
 }
+
+/// APK config
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub struct ApkConfig {
+    /// The path of APK
+    pub path: String,
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
index 073c088..0cb187c 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
@@ -23,6 +23,9 @@
     /** idsig for an APK */
     ParcelFileDescriptor idsig;
 
+    /** Idsigs for the extra APKs. Must match with the extra_apks in the payload config. */
+    List<ParcelFileDescriptor> extraIdsigs;
+
     /** instance.img that has per-instance data */
     ParcelFileDescriptor instanceImage;
 
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 55eb19b..8f7a69a 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -20,7 +20,7 @@
     VirtualMachineRawConfig::VirtualMachineRawConfig,
 };
 use android_system_virtualizationservice::binder::ParcelFileDescriptor;
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use binder::wait_for_interface;
 use log::{error, info};
 use microdroid_metadata::{ApexPayload, ApkPayload, Metadata};
@@ -194,22 +194,34 @@
 ///   ..
 ///   microdroid-apk: apk
 ///   microdroid-apk-idsig: idsig
+///   extra-apk-0:   additional apk 0
+///   extra-idsig-0: additional idsig 0
+///   extra-apk-1:   additional apk 1
+///   extra-idsig-1: additional idsig 1
+///   ..
 fn make_payload_disk(
+    app_config: &VirtualMachineAppConfig,
     apk_file: File,
     idsig_file: File,
-    config_path: &str,
     vm_payload_config: &VmPayloadConfig,
     temporary_directory: &Path,
-    debug_level: DebugLevel,
 ) -> Result<DiskImage> {
+    if vm_payload_config.extra_apks.len() != app_config.extraIdsigs.len() {
+        bail!(
+            "payload config has {} apks, but app config has {} idsigs",
+            vm_payload_config.extra_apks.len(),
+            app_config.extraIdsigs.len()
+        );
+    }
+
     let pm = PackageManager::new()?;
     let apex_list = pm.get_apex_list(vm_payload_config.prefer_staged)?;
 
     // collect APEX names from config
-    let apexes = collect_apex_names(&apex_list, &vm_payload_config.apexes, debug_level);
+    let apexes = collect_apex_names(&apex_list, &vm_payload_config.apexes, app_config.debugLevel);
     info!("Microdroid payload APEXes: {:?}", apexes);
 
-    let metadata_file = make_metadata_file(config_path, &apexes, temporary_directory)?;
+    let metadata_file = make_metadata_file(&app_config.configPath, &apexes, temporary_directory)?;
     // put metadata at the first partition
     let mut partitions = vec![Partition {
         label: "payload-metadata".to_owned(),
@@ -237,6 +249,23 @@
         writable: false,
     });
 
+    // we've already checked that extra_apks and extraIdsigs are in the same size.
+    let extra_apks = &vm_payload_config.extra_apks;
+    let extra_idsigs = &app_config.extraIdsigs;
+    for (i, (extra_apk, extra_idsig)) in extra_apks.iter().zip(extra_idsigs.iter()).enumerate() {
+        partitions.push(Partition {
+            label: format!("extra-apk-{}", i),
+            image: Some(ParcelFileDescriptor::new(File::open(PathBuf::from(&extra_apk.path))?)),
+            writable: false,
+        });
+
+        partitions.push(Partition {
+            label: format!("extra-idsig-{}", i),
+            image: Some(ParcelFileDescriptor::new(extra_idsig.as_ref().try_clone()?)),
+            writable: false,
+        });
+    }
+
     Ok(DiskImage { image: None, partitions, writable: false })
 }
 
@@ -298,12 +327,11 @@
     vm_config: &mut VirtualMachineRawConfig,
 ) -> Result<()> {
     vm_config.disks.push(make_payload_disk(
+        config,
         apk_file,
         idsig_file,
-        &config.configPath,
         vm_payload_config,
         temporary_directory,
-        config.debugLevel,
     )?);
 
     vm_config.disks[1].partitions.push(Partition {
diff --git a/vm/Android.bp b/vm/Android.bp
index 734f2d3..2d22562 100644
--- a/vm/Android.bp
+++ b/vm/Android.bp
@@ -7,16 +7,19 @@
     crate_name: "vm",
     srcs: ["src/main.rs"],
     edition: "2018",
+    prefer_rlib: true,
     rustlibs: [
         "android.system.virtualizationservice-rust",
         "libanyhow",
         "libenv_logger",
         "liblibc",
         "liblog_rust",
+        "libmicrodroid_payload_config",
         "libserde_json",
         "libserde",
         "libstructopt",
         "libvmconfig",
+        "libzip",
     ],
     apex_available: [
         "com.android.virt",
diff --git a/vm/src/main.rs b/vm/src/main.rs
index 87bcda7..d53305b 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -33,6 +33,9 @@
 const VIRTUALIZATION_SERVICE_BINDER_SERVICE_IDENTIFIER: &str =
     "android.system.virtualizationservice";
 
+#[derive(Debug)]
+struct Idsigs(Vec<PathBuf>);
+
 #[derive(StructOpt)]
 #[structopt(no_version, global_settings = &[AppSettings::DisableVersion])]
 enum Opt {
@@ -73,6 +76,10 @@
         /// in the VM config file.
         #[structopt(short, long)]
         mem: Option<u32>,
+
+        /// Paths to extra idsig files.
+        #[structopt(long)]
+        extra_idsigs: Vec<PathBuf>,
     },
     /// Run a virtual machine
     Run {
@@ -138,20 +145,30 @@
         .context("Failed to find VirtualizationService")?;
 
     match opt {
-        Opt::RunApp { apk, idsig, instance, config_path, daemonize, console, log, debug, mem } => {
-            command_run_app(
-                service,
-                &apk,
-                &idsig,
-                &instance,
-                &config_path,
-                daemonize,
-                console.as_deref(),
-                log.as_deref(),
-                debug,
-                mem,
-            )
-        }
+        Opt::RunApp {
+            apk,
+            idsig,
+            instance,
+            config_path,
+            daemonize,
+            console,
+            log,
+            debug,
+            mem,
+            extra_idsigs,
+        } => command_run_app(
+            service,
+            &apk,
+            &idsig,
+            &instance,
+            &config_path,
+            daemonize,
+            console.as_deref(),
+            log.as_deref(),
+            debug,
+            mem,
+            &extra_idsigs,
+        ),
         Opt::Run { config, daemonize, console } => {
             command_run(service, &config, daemonize, console.as_deref(), /* mem */ None)
         }
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 15775cb..1cd51a1 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -28,12 +28,14 @@
     BinderFeatures, DeathRecipient, IBinder, ParcelFileDescriptor, Strong,
 };
 use android_system_virtualizationservice::binder::{Interface, Result as BinderResult};
-use anyhow::{Context, Error};
+use anyhow::{bail, Context, Error};
+use microdroid_payload_config::VmPayloadConfig;
 use std::fs::File;
 use std::io::{self, BufRead, BufReader};
 use std::os::unix::io::{AsRawFd, FromRawFd};
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use vmconfig::{open_parcel_file, VmConfig};
+use zip::ZipArchive;
 
 /// Run a VM from the given APK, idsig, and config.
 #[allow(clippy::too_many_arguments)]
@@ -48,7 +50,23 @@
     log_path: Option<&Path>,
     debug_level: DebugLevel,
     mem: Option<u32>,
+    extra_idsigs: &[PathBuf],
 ) -> Result<(), Error> {
+    let extra_apks = parse_extra_apk_list(apk, config_path)?;
+    if extra_apks.len() != extra_idsigs.len() {
+        bail!(
+            "Found {} extra apks, but there are {} extra idsigs",
+            extra_apks.len(),
+            extra_idsigs.len()
+        )
+    }
+
+    for i in 0..extra_apks.len() {
+        let extra_apk_fd = ParcelFileDescriptor::new(File::open(&extra_apks[i])?);
+        let extra_idsig_fd = ParcelFileDescriptor::new(File::create(&extra_idsigs[i])?);
+        service.createOrUpdateIdsigFile(&extra_apk_fd, &extra_idsig_fd)?;
+    }
+
     let apk_file = File::open(apk).context("Failed to open APK file")?;
     let idsig_file = File::create(idsig).context("Failed to create idsig file")?;
 
@@ -69,9 +87,13 @@
         )?;
     }
 
+    let extra_idsig_files: Result<Vec<File>, _> = extra_idsigs.iter().map(File::open).collect();
+    let extra_idsig_fds = extra_idsig_files?.into_iter().map(ParcelFileDescriptor::new).collect();
+
     let config = VirtualMachineConfig::AppConfig(VirtualMachineAppConfig {
         apk: apk_fd.into(),
         idsig: idsig_fd.into(),
+        extraIdsigs: extra_idsig_fds,
         instanceImage: open_parcel_file(instance, true /* writable */)?.into(),
         configPath: config_path.to_owned(),
         debugLevel: debug_level,
@@ -204,6 +226,13 @@
     Ok(death_recipient)
 }
 
+fn parse_extra_apk_list(apk: &Path, config_path: &str) -> Result<Vec<String>, Error> {
+    let mut archive = ZipArchive::new(File::open(apk)?)?;
+    let config_file = archive.by_name(config_path)?;
+    let config: VmPayloadConfig = serde_json::from_reader(config_file)?;
+    Ok(config.extra_apks.into_iter().map(|x| x.path).collect())
+}
+
 #[derive(Debug)]
 struct VirtualMachineCallback {
     dead: AtomicFlag,