Merge "Pass odrefresh its own environment"
diff --git a/compos/Android.bp b/compos/Android.bp
index 658f8bf..c54348a 100644
--- a/compos/Android.bp
+++ b/compos/Android.bp
@@ -24,7 +24,6 @@
         "liblog_rust",
         "libminijail_rust",
         "libnix",
-        "libnum_cpus",
         "libodsign_proto_rust",
         "libprotobuf",
         "libregex",
diff --git a/compos/composd/Android.bp b/compos/composd/Android.bp
index 3b545e5..55a3107 100644
--- a/compos/composd/Android.bp
+++ b/compos/composd/Android.bp
@@ -18,6 +18,7 @@
         "libcompos_common",
         "libcomposd_native_rust",
         "libminijail_rust",
+        "libnum_cpus",
         "libnix",
         "liblibc",
         "liblog_rust",
diff --git a/compos/composd/src/instance_manager.rs b/compos/composd/src/instance_manager.rs
index 8950f20..e8c1d9a 100644
--- a/compos/composd/src/instance_manager.rs
+++ b/compos/composd/src/instance_manager.rs
@@ -45,10 +45,14 @@
     pub fn start_pending_instance(&self) -> Result<Arc<CompOsInstance>> {
         let config_path = Some(PREFER_STAGED_VM_CONFIG_PATH.to_owned());
         let mut vm_parameters = VmParameters { config_path, ..Default::default() };
-        vm_parameters.cpus = NonZeroU32::from_str(
-            &system_properties::read(DEX2OAT_THREADS_PROP_NAME).unwrap_or_default(),
-        )
-        .ok();
+        vm_parameters.cpus = match system_properties::read(DEX2OAT_THREADS_PROP_NAME) {
+            Ok(s) => Some(NonZeroU32::from_str(&s)?),
+            Err(_) => {
+                // dex2oat uses all CPUs by default. To match the behavior, give the VM all CPUs by
+                // default.
+                NonZeroU32::new(num_cpus::get() as u32)
+            }
+        };
         vm_parameters.cpu_set = system_properties::read(DEX2OAT_CPU_SET_PROP_NAME).ok();
         self.start_instance(PENDING_INSTANCE_DIR, vm_parameters)
     }
diff --git a/compos/tests/java/android/compos/test/ComposKeyTestCase.java b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
index d59d3d9..49235fe 100644
--- a/compos/tests/java/android/compos/test/ComposKeyTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
@@ -38,13 +38,25 @@
 @RunWith(DeviceJUnit4ClassRunner.class)
 public final class ComposKeyTestCase extends VirtualizationTestCaseBase {
 
-    /** Wait time for service to be ready on boot */
+    /**
+     * Wait time for service to be ready on boot
+     */
     private static final int READY_LATENCY_MS = 10 * 1000; // 10 seconds
 
-    /** Path to compos_key_cmd tool */
+    /**
+     * Path to compos_key_cmd tool
+     */
     private static final String COMPOS_KEY_CMD_BIN = "/apex/com.android.compos/bin/compos_key_cmd";
 
-    /** Config of the test VM. This is a path inside the APK. */
+    /**
+     * Path to the com.android.compos.payload APK
+     */
+    private static final String COMPOS_PAYLOAD_APK_PATH =
+            "/apex/com.android.compos/app/CompOSPayloadApp/CompOSPayloadApp.apk";
+
+    /**
+     * Config of the test VM. This is a path inside the APK.
+     */
     private static final String VM_TEST_CONFIG_PATH = "assets/vm_test_config.json";
 
     private String mCid;
@@ -134,7 +146,9 @@
                         getDevice(),
                         getBuild(),
                         /* apkName, no need to install */ null,
-                        packageName,
+                        COMPOS_PAYLOAD_APK_PATH,
+                        /* packageName - not needed, we know the path */ null,
+                        /* extraIdSigPaths */ null,
                         VM_TEST_CONFIG_PATH,
                         /* debug */ true,
                         /* use default memoryMib */ 0,
diff --git a/microdroid/payload/metadata.proto b/microdroid/payload/metadata.proto
index 9e60b38..2e92f55 100644
--- a/microdroid/payload/metadata.proto
+++ b/microdroid/payload/metadata.proto
@@ -43,6 +43,10 @@
   // The timestamp in seconds when the APEX was last updated. This should match the value in
   // apex-info-list.xml.
   uint64 last_update_seconds = 5;
+
+  // Required.
+  // Whether the APEX is a factory version or not.
+  bool is_factory = 6;
 }
 
 message ApkPayload {
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
index 1068792..5a77198 100644
--- a/microdroid_manager/src/instance.rs
+++ b/microdroid_manager/src/instance.rs
@@ -336,4 +336,5 @@
     pub public_key: Vec<u8>,
     pub root_digest: Vec<u8>,
     pub last_update_seconds: u64,
+    pub is_factory: bool,
 }
diff --git a/microdroid_manager/src/payload.rs b/microdroid_manager/src/payload.rs
index 661af5f..48535f3 100644
--- a/microdroid_manager/src/payload.rs
+++ b/microdroid_manager/src/payload.rs
@@ -48,6 +48,7 @@
                 public_key: result.public_key,
                 root_digest: result.root_digest,
                 last_update_seconds: apex.last_update_seconds,
+                is_factory: apex.is_factory,
             })
         })
         .collect()
@@ -63,6 +64,7 @@
                 public_key: data.public_key.clone(),
                 root_digest: data.root_digest.clone(),
                 last_update_seconds: data.last_update_seconds,
+                is_factory: data.is_factory,
                 ..Default::default()
             })
             .collect(),
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index 678fe84..e15f1ae 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -204,6 +204,24 @@
             Optional<Integer> numCpus,
             Optional<String> cpuAffinity)
             throws DeviceNotAvailableException {
+        return startMicrodroid(androidDevice, buildInfo, apkName, null, packageName,
+                extraIdsigPaths, configPath, debug,
+                memoryMib, numCpus, cpuAffinity);
+    }
+
+    public static String startMicrodroid(
+            ITestDevice androidDevice,
+            IBuildInfo buildInfo,
+            String apkName,
+            String apkPath,
+            String packageName,
+            String[] extraIdsigPaths,
+            String configPath,
+            boolean debug,
+            int memoryMib,
+            Optional<Integer> numCpus,
+            Optional<String> cpuAffinity)
+            throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(androidDevice);
 
         // Install APK if necessary
@@ -215,9 +233,11 @@
         // Get the path to the installed apk. Note that
         // getDevice().getAppPackageInfo(...).getCodePath() doesn't work due to the incorrect
         // parsing of the "=" character. (b/190975227). So we use the `pm path` command directly.
-        String apkPath = android.run("pm", "path", packageName);
-        assertTrue(apkPath.startsWith("package:"));
-        apkPath = apkPath.substring("package:".length());
+        if (apkPath == null) {
+            apkPath = android.run("pm", "path", packageName);
+            assertTrue(apkPath.startsWith("package:"));
+            apkPath = apkPath.substring("package:".length());
+        }
 
         android.run("mkdir", "-p", TEST_ROOT);
 
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 65adc25..e4a763c 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -66,6 +66,9 @@
     // The field claims to be milliseconds but is actually seconds.
     #[serde(rename = "lastUpdateMillis")]
     last_update_seconds: u64,
+
+    #[serde(rename = "isFactory")]
+    is_factory: bool,
 }
 
 impl ApexInfoList {
@@ -137,6 +140,8 @@
                         let metadata = metadata(&apex_info.path)?;
                         apex_info.last_update_seconds =
                             metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs();
+                        // by definition, staged apex can't be a factory apex.
+                        apex_info.is_factory = false;
                     }
                 }
             }
@@ -158,10 +163,12 @@
             .iter()
             .enumerate()
             .map(|(i, apex_name)| {
+                let apex_info = apex_list.get(apex_name)?;
                 Ok(ApexPayload {
                     name: apex_name.clone(),
                     partition_name: format!("microdroid-apex-{}", i),
-                    last_update_seconds: apex_list.get(apex_name)?.last_update_seconds,
+                    last_update_seconds: apex_info.last_update_seconds,
+                    is_factory: apex_info.is_factory,
                     ..Default::default()
                 })
             })
@@ -411,12 +418,14 @@
                     path: PathBuf::from("path0"),
                     has_classpath_jar: false,
                     last_update_seconds: 12345678,
+                    is_factory: true,
                 },
                 ApexInfo {
                     name: "has_classpath".to_string(),
                     path: PathBuf::from("path1"),
                     has_classpath_jar: true,
                     last_update_seconds: 87654321,
+                    is_factory: false,
                 },
             ],
         };