Merge "Remove obsolete TODOs."
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..2f15cb5 100644
--- a/compos/composd/src/instance_manager.rs
+++ b/compos/composd/src/instance_manager.rs
@@ -43,18 +43,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.cpu_set = system_properties::read(DEX2OAT_CPU_SET_PROP_NAME).ok();
+        let mut vm_parameters = new_vm_parameters()?;
+        vm_parameters.config_path = Some(PREFER_STAGED_VM_CONFIG_PATH.to_owned());
         self.start_instance(PENDING_INSTANCE_DIR, vm_parameters)
     }
 
     pub fn start_test_instance(&self) -> Result<Arc<CompOsInstance>> {
-        let vm_parameters = VmParameters { debug_mode: true, ..Default::default() };
+        let mut vm_parameters = new_vm_parameters()?;
+        vm_parameters.debug_mode = true;
         self.start_instance(TEST_INSTANCE_DIR, vm_parameters)
     }
 
@@ -86,6 +82,19 @@
     }
 }
 
+fn new_vm_parameters() -> Result<VmParameters> {
+    let 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)
+        }
+    };
+    let cpu_set = system_properties::read(DEX2OAT_CPU_SET_PROP_NAME).ok();
+    Ok(VmParameters { cpus, cpu_set, ..Default::default() })
+}
+
 // Ensures we only run one instance at a time.
 // Valid states:
 // Starting: is_starting is true, running_instance is None.
diff --git a/compos/native/compos_native.cpp b/compos/native/compos_native.cpp
index f529421..81280fd 100644
--- a/compos/native/compos_native.cpp
+++ b/compos/native/compos_native.cpp
@@ -42,7 +42,7 @@
     bssl::UniquePtr<RSA> key_pair(RSA_new());
 
     // This function specifies that the public exponent is always 65537, which is good because
-    // that's  what odsign is expecting.
+    // that's what odsign is expecting.
     if (!RSA_generate_key_fips(key_pair.get(), KEY_BITS, /*callback=*/nullptr)) {
         return make_key_error("Failed to generate key pair");
     }
diff --git a/compos/src/compilation.rs b/compos/src/compilation.rs
index 850a0a8..e8f55f8 100644
--- a/compos/src/compilation.rs
+++ b/compos/src/compilation.rs
@@ -14,14 +14,14 @@
  * limitations under the License.
  */
 
-use anyhow::{bail, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use log::{debug, info, warn};
 use minijail::{self, Minijail};
 use regex::Regex;
+use std::collections::HashMap;
 use std::env;
 use std::ffi::OsString;
 use std::fs::read_dir;
-use std::os::unix::io::AsRawFd;
 use std::path::{self, Path, PathBuf};
 use std::process::Command;
 
@@ -34,15 +34,11 @@
     },
     IAuthFsService::IAuthFsService,
 };
-use authfs_aidl_interface::binder::{ParcelFileDescriptor, Strong};
+use authfs_aidl_interface::binder::Strong;
 use compos_common::odrefresh::ExitCode;
 
 const FD_SERVER_PORT: i32 = 3264; // TODO: support dynamic port
 
-/// The number that represents the file descriptor number expecting by the task. The number may be
-/// meaningless in the current process.
-pub type PseudoRawFd = i32;
-
 pub struct OdrefreshContext<'a> {
     system_dir_fd: i32,
     output_dir_fd: i32,
@@ -114,19 +110,22 @@
     let authfs = authfs_service.mount(&authfs_config)?;
     let mountpoint = PathBuf::from(authfs.getMountPoint()?);
 
+    // Make a copy of our environment as the basis of the one we will give odrefresh
+    let mut odrefresh_vars = EnvMap::from_current_env();
+
     let mut android_root = mountpoint.clone();
     android_root.push(context.system_dir_fd.to_string());
     android_root.push("system");
-    env::set_var("ANDROID_ROOT", &android_root);
+    odrefresh_vars.set("ANDROID_ROOT", path_to_str(&android_root)?);
     debug!("ANDROID_ROOT={:?}", &android_root);
 
     let art_apex_data = mountpoint.join(context.output_dir_fd.to_string());
-    env::set_var("ART_APEX_DATA", &art_apex_data);
+    odrefresh_vars.set("ART_APEX_DATA", path_to_str(&art_apex_data)?);
     debug!("ART_APEX_DATA={:?}", &art_apex_data);
 
     let staging_dir = mountpoint.join(context.staging_dir_fd.to_string());
 
-    set_classpaths(&android_root)?;
+    set_classpaths(&mut odrefresh_vars, &android_root)?;
 
     let mut args = vec![
         "odrefresh".to_string(),
@@ -147,7 +146,7 @@
     args.push("--compile".to_string());
 
     debug!("Running odrefresh with args: {:?}", &args);
-    let jail = spawn_jailed_task(odrefresh_path, &args, Vec::new() /* fd_mapping */)
+    let jail = spawn_jailed_task(odrefresh_path, &args, &odrefresh_vars.into_env())
         .context("Spawn odrefresh")?;
     let exit_code = match jail.wait() {
         Ok(_) => Result::<u8>::Ok(0),
@@ -173,9 +172,13 @@
     Ok(exit_code)
 }
 
-fn set_classpaths(android_root: &Path) -> Result<()> {
+fn path_to_str(path: &Path) -> Result<&str> {
+    path.to_str().ok_or_else(|| anyhow!("Bad path {:?}", path))
+}
+
+fn set_classpaths(odrefresh_vars: &mut EnvMap, android_root: &Path) -> Result<()> {
     let export_lines = run_derive_classpath(android_root)?;
-    load_classpath_vars(&export_lines)
+    load_classpath_vars(odrefresh_vars, &export_lines)
 }
 
 fn run_derive_classpath(android_root: &Path) -> Result<String> {
@@ -203,15 +206,14 @@
     String::from_utf8(result.stdout).context("Converting derive_classpath output")
 }
 
-fn load_classpath_vars(export_lines: &str) -> Result<()> {
+fn load_classpath_vars(odrefresh_vars: &mut EnvMap, export_lines: &str) -> Result<()> {
     // Each line should be in the format "export <var name> <value>"
     let pattern = Regex::new(r"^export ([^ ]+) ([^ ]+)$").context("Failed to construct Regex")?;
     for line in export_lines.lines() {
         if let Some(captures) = pattern.captures(line) {
             let name = &captures[1];
             let value = &captures[2];
-            // TODO(b/213416778) Don't modify our env, construct a fresh one for odrefresh
-            env::set_var(name, value);
+            odrefresh_vars.set(name, value);
         } else {
             warn!("Malformed line from derive_classpath: {}", line);
         }
@@ -238,14 +240,28 @@
     Ok(())
 }
 
-fn spawn_jailed_task(
-    executable: &Path,
-    args: &[String],
-    fd_mapping: Vec<(ParcelFileDescriptor, PseudoRawFd)>,
-) -> Result<Minijail> {
+fn spawn_jailed_task(executable: &Path, args: &[String], env_vars: &[String]) -> Result<Minijail> {
     // TODO(b/185175567): Run in a more restricted sandbox.
     let jail = Minijail::new()?;
-    let preserve_fds: Vec<_> = fd_mapping.iter().map(|(f, id)| (f.as_raw_fd(), *id)).collect();
-    let _pid = jail.run_remap(executable, preserve_fds.as_slice(), args)?;
+    let keep_fds = [];
+    let command = minijail::Command::new_for_path(executable, &keep_fds, args, Some(env_vars))?;
+    let _pid = jail.run_command(command)?;
     Ok(jail)
 }
+
+struct EnvMap(HashMap<String, String>);
+
+impl EnvMap {
+    fn from_current_env() -> Self {
+        Self(env::vars().collect())
+    }
+
+    fn set(&mut self, key: &str, value: &str) {
+        self.0.insert(key.to_owned(), value.to_owned());
+    }
+
+    fn into_env(self) -> Vec<String> {
+        // execve() expects an array of "k=v" strings, rather than a list of (k, v) pairs.
+        self.0.into_iter().map(|(k, v)| k + "=" + &v).collect()
+    }
+}
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/microdroid.json b/microdroid/microdroid.json
index c6c743d..0c294e9 100644
--- a/microdroid/microdroid.json
+++ b/microdroid/microdroid.json
@@ -37,6 +37,6 @@
       "writable": true
     }
   ],
-  "memory_mib": 2048,
+  "memory_mib": 256,
   "protected": false
 }
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 028e3ad..c661d44 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 {
@@ -136,6 +139,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;
                     }
                 }
             }
@@ -157,10 +162,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()
                 })
             })
@@ -410,12 +417,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,
                 },
             ],
         };