Merge "Use staged APEXes for compilation"
diff --git a/compos/apk/assets/vm_config_staged.json b/compos/apk/assets/vm_config_staged.json
new file mode 100644
index 0000000..9c81e4e
--- /dev/null
+++ b/compos/apk/assets/vm_config_staged.json
@@ -0,0 +1,31 @@
+{
+  "version": 1,
+  "os": {
+    "name": "microdroid"
+  },
+  "task": {
+    "type": "executable",
+    "command": "/apex/com.android.compos/bin/compsvc",
+    "args": [
+      "--log_to_stderr"
+    ]
+  },
+  "prefer_staged": true,
+  "apexes": [
+    {
+      "name": "com.android.art"
+    },
+    {
+      "name": "com.android.compos"
+    },
+    {
+      "name": "{DEX2OATBOOTCLASSPATH}"
+    },
+    {
+      "name": "{BOOTCLASSPATH}"
+    },
+    {
+      "name": "{SYSTEMSERVERCLASSPATH}"
+    }
+  ]
+}
\ No newline at end of file
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 9e419f5..6c16bb0 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -17,7 +17,7 @@
 //! Support for starting CompOS in a VM and connecting to the service
 
 use crate::timeouts::timeouts;
-use crate::{COMPOS_APEX_ROOT, COMPOS_DATA_ROOT, COMPOS_VSOCK_PORT};
+use crate::{COMPOS_APEX_ROOT, COMPOS_DATA_ROOT, COMPOS_VSOCK_PORT, DEFAULT_VM_CONFIG_PATH};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
     IVirtualMachine::IVirtualMachine,
     IVirtualMachineCallback::{BnVirtualMachineCallback, IVirtualMachineCallback},
@@ -56,6 +56,8 @@
 pub struct VmParameters {
     /// Whether the VM should be debuggable.
     pub debug_mode: bool,
+    /// If present, overrides the path to the VM config JSON file
+    pub config_path: Option<String>,
 }
 
 impl VmInstance {
@@ -95,11 +97,12 @@
             (None, DebugLevel::NONE)
         };
 
+        let config_path = parameters.config_path.as_deref().unwrap_or(DEFAULT_VM_CONFIG_PATH);
         let config = VirtualMachineConfig::AppConfig(VirtualMachineAppConfig {
             apk: Some(apk_fd),
             idsig: Some(idsig_fd),
             instanceImage: Some(instance_fd),
-            configPath: "assets/vm_config.json".to_owned(),
+            configPath: config_path.to_owned(),
             debugLevel: debug_level,
             ..Default::default()
         });
diff --git a/compos/common/lib.rs b/compos/common/lib.rs
index 4bfa81f..9b07030 100644
--- a/compos/common/lib.rs
+++ b/compos/common/lib.rs
@@ -52,3 +52,10 @@
 
 /// The file that holds the instance image for a CompOS instance.
 pub const INSTANCE_IMAGE_FILE: &str = "instance.img";
+
+/// The path within our config APK of our default VM configuration file, used at boot time.
+pub const DEFAULT_VM_CONFIG_PATH: &str = "assets/vm_config.json";
+
+/// The path within our config APK of the VM configuration file we use when compiling staged
+/// APEXes before reboot.
+pub const PREFER_STAGED_VM_CONFIG_PATH: &str = "assets/vm_config_staged.json";
diff --git a/compos/compos_key_cmd/compos_key_cmd.cpp b/compos/compos_key_cmd/compos_key_cmd.cpp
index 560eb09..76ff06f 100644
--- a/compos/compos_key_cmd/compos_key_cmd.cpp
+++ b/compos/compos_key_cmd/compos_key_cmd.cpp
@@ -74,8 +74,9 @@
 constexpr const char* kConfigApkIdsigPath =
         "/apex/com.android.compos/etc/CompOSPayloadApp.apk.idsig";
 
-// This is a path inside the APK
-constexpr const char* kConfigFilePath = "assets/vm_config.json";
+// These are paths inside the APK
+constexpr const char* kDefaultConfigFilePath = "assets/vm_config.json";
+constexpr const char* kPreferStagedConfigFilePath = "assets/vm_config_staged.json";
 
 static bool writeBytesToFile(const std::vector<uint8_t>& bytes, const std::string& path) {
     std::string str(bytes.begin(), bytes.end());
@@ -186,11 +187,12 @@
 class TargetVm {
 public:
     TargetVm(int cid, const std::string& logFile, const std::string& instanceImageFile,
-             bool debuggable)
+             bool debuggable, bool preferStaged)
           : mCid(cid),
             mLogFile(logFile),
             mInstanceImageFile(instanceImageFile),
-            mDebuggable(debuggable) {}
+            mDebuggable(debuggable),
+            mPreferStaged(preferStaged) {}
 
     // Returns 0 if we are to connect to a local service, otherwise the CID of
     // either an existing VM or a VM we have started, depending on the command
@@ -251,7 +253,7 @@
         appConfig.apk = std::move(apkFd);
         appConfig.idsig = std::move(idsigFd);
         appConfig.instanceImage = std::move(instanceFd);
-        appConfig.configPath = kConfigFilePath;
+        appConfig.configPath = mPreferStaged ? kPreferStagedConfigFilePath : kDefaultConfigFilePath;
         appConfig.debugLevel = mDebuggable ? VirtualMachineAppConfig::DebugLevel::FULL
                                            : VirtualMachineAppConfig::DebugLevel::NONE;
         appConfig.memoryMib = 0; // Use default
@@ -297,6 +299,7 @@
     const std::string mLogFile;
     const std::string mInstanceImageFile;
     const bool mDebuggable;
+    const bool mPreferStaged;
     std::shared_ptr<Callback> mCallback;
     std::shared_ptr<IVirtualMachine> mVm;
 };
@@ -543,17 +546,25 @@
     std::string imageFile;
     std::string logFile;
     bool debuggable = false;
+    bool preferStaged = false;
 
     for (;;) {
+        // Options with no associated value
         if (argc >= 2) {
             if (argv[1] == "--debug"sv) {
                 debuggable = true;
                 argc -= 1;
                 argv += 1;
                 continue;
+            } else if (argv[1] == "--staged"sv) {
+                preferStaged = true;
+                argc -= 1;
+                argv += 1;
+                continue;
             }
         }
         if (argc < 3) break;
+        // Options requiring a value
         if (argv[1] == "--cid"sv) {
             cid = atoi(argv[2]);
             if (cid == 0) {
@@ -571,7 +582,7 @@
         argv += 2;
     }
 
-    TargetVm vm(cid, logFile, imageFile, debuggable);
+    TargetVm vm(cid, logFile, imageFile, debuggable, preferStaged);
 
     if (argc == 4 && argv[1] == "generate"sv) {
         auto result = generate(vm, argv[2], argv[3]);
@@ -627,9 +638,11 @@
                   << "    <filename>.signature\n"
                   << "  make-instance <image file> Create an empty instance image file for a VM.\n"
                   << "\n"
-                  << "OPTIONS: --log <log file> --debug (--cid <cid> | --start <image file>)\n"
+                  << "OPTIONS: --log <log file> --debug --staged\n"
+                  << "    (--cid <cid> | --start <image file>)\n"
                   << "  Specify --log to write VM log to a file rather than stdout.\n"
                   << "  Specify --debug with --start to make the VM fully debuggable.\n"
+                  << "  Specify --staged with --start to prefer staged APEXes in the VM.\n"
                   << "  Specify --cid to connect to a VM rather than the host.\n"
                   << "  Specify --start to start a VM from the given instance image file and\n "
                   << "    connect to that.\n";
diff --git a/compos/composd/src/instance_manager.rs b/compos/composd/src/instance_manager.rs
index 767e9f7..24ae576 100644
--- a/compos/composd/src/instance_manager.rs
+++ b/compos/composd/src/instance_manager.rs
@@ -23,7 +23,7 @@
 use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
 use compos_aidl_interface::binder::Strong;
 use compos_common::compos_client::VmParameters;
-use compos_common::{PENDING_INSTANCE_DIR, TEST_INSTANCE_DIR};
+use compos_common::{PENDING_INSTANCE_DIR, PREFER_STAGED_VM_CONFIG_PATH, TEST_INSTANCE_DIR};
 use std::sync::{Arc, Mutex, Weak};
 use virtualizationservice::IVirtualizationService::IVirtualizationService;
 
@@ -44,11 +44,13 @@
     }
 
     pub fn start_pending_instance(&self) -> Result<Arc<CompOsInstance>> {
-        self.start_instance(PENDING_INSTANCE_DIR, VmParameters::default())
+        let config_path = Some(PREFER_STAGED_VM_CONFIG_PATH.to_owned());
+        let vm_parameters = VmParameters { config_path, ..Default::default() };
+        self.start_instance(PENDING_INSTANCE_DIR, vm_parameters)
     }
 
     pub fn start_test_instance(&self) -> Result<Arc<CompOsInstance>> {
-        let vm_parameters = VmParameters { debug_mode: true };
+        let vm_parameters = VmParameters { debug_mode: true, ..Default::default() };
         self.start_instance(TEST_INSTANCE_DIR, vm_parameters)
     }
 
diff --git a/compos/verify_key/verify_key.rs b/compos/verify_key/verify_key.rs
index 945acb4..e0ed5e5 100644
--- a/compos/verify_key/verify_key.rs
+++ b/compos/verify_key/verify_key.rs
@@ -105,8 +105,11 @@
     let instance_image = File::open(instance_image).context("Failed to open instance image")?;
 
     let virtualization_service = VmInstance::connect_to_virtualization_service()?;
-    let vm_instance =
-        VmInstance::start(&*virtualization_service, instance_image, &VmParameters { debug_mode })?;
+    let vm_instance = VmInstance::start(
+        &*virtualization_service,
+        instance_image,
+        &VmParameters { debug_mode, ..Default::default() },
+    )?;
     let service = vm_instance.get_service()?;
 
     let result = service.verifySigningKey(&blob, &public_key).context("Verifying signing key")?;