Build composite image in VirtualizationService.

Bug: 184131523
Test: atest VirtualizationTestCases
Test: ran microdroid manually
Change-Id: I24eb776bb3049a4cdc5f000607447a73bd0adeec
diff --git a/virtualizationservice/src/composite.rs b/virtualizationservice/src/composite.rs
new file mode 100644
index 0000000..eb738a7
--- /dev/null
+++ b/virtualizationservice/src/composite.rs
@@ -0,0 +1,112 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Functions for running `mk_cdisk`.
+
+mod config;
+
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::Partition::Partition as AidlPartition;
+use anyhow::{bail, Context, Error};
+use command_fds::{CommandFdExt, FdMapping};
+use config::{Config, Partition};
+use log::info;
+use std::fs::File;
+use std::os::unix::io::AsRawFd;
+use std::panic;
+use std::path::Path;
+use std::process::{Command, Stdio};
+use std::str;
+use std::thread;
+
+const MK_CDISK_PATH: &str = "/apex/com.android.virt/bin/mk_cdisk";
+
+/// Calls `mk_cdisk` to construct a composite disk image for the given list of partitions, and opens
+/// it ready to use. Returns the composite disk image file, and a list of FD mappings which must be
+/// applied to any process which wants to use it. This is necessary because the composite image
+/// contains paths of the form `/proc/self/fd/N` for the partition images.
+pub fn make_composite_image(
+    partitions: &[AidlPartition],
+    output_filename: &Path,
+) -> Result<(File, Vec<File>), Error> {
+    let (config_json, files) = make_config_json(partitions)?;
+    let fd_mappings: Vec<_> = files
+        .iter()
+        .map(|file| FdMapping { parent_fd: file.as_raw_fd(), child_fd: file.as_raw_fd() })
+        .collect();
+
+    let mut command = Command::new(MK_CDISK_PATH);
+    command
+        .arg("-") // Read config JSON from stdin.
+        .arg(&output_filename)
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped());
+    command.fd_mappings(fd_mappings)?;
+    let mut child = command.spawn().context("Failed to spawn mk_cdisk")?;
+    let stdin = child.stdin.take().unwrap();
+
+    // Write config to stdin of mk_cdisk on a separate thread to avoid deadlock, as it may not read
+    // all of stdin before it blocks on writing to stdout.
+    let writer_thread = thread::spawn(move || config_json.write_json(&stdin));
+    info!("Running {:?}", command);
+    let output = child.wait_with_output()?;
+    match writer_thread.join() {
+        Ok(result) => result?,
+        Err(panic_payload) => panic::resume_unwind(panic_payload),
+    }
+
+    if !output.status.success() {
+        info!("mk_cdisk stdout: {}", str::from_utf8(&output.stdout)?);
+        info!("mk_cdisk stderr: {}", str::from_utf8(&output.stderr)?);
+        bail!("mk_cdisk exited with error {}", output.status);
+    }
+
+    let composite_image = File::open(&output_filename)
+        .with_context(|| format!("Failed to open composite image {:?}", output_filename))?;
+
+    Ok((composite_image, files))
+}
+
+/// Given the AIDL config containing a list of partitions, with a [`ParcelFileDescriptor`] for each
+/// partition, return the list of file descriptors which must be passed to the mk_cdisk child
+/// process and the JSON configuration for it.
+fn make_config_json(partitions: &[AidlPartition]) -> Result<(Config, Vec<File>), Error> {
+    // File descriptors to pass to child process.
+    let mut files = vec![];
+
+    let partitions = partitions
+        .iter()
+        .map(|partition| {
+            // TODO(b/187187765): This shouldn't be an Option.
+            let file = partition
+                .image
+                .as_ref()
+                .context("Invalid partition image file descriptor")?
+                .as_ref()
+                .try_clone()
+                .context("Failed to clone partition image file descriptor")?;
+            let fd = file.as_raw_fd();
+            files.push(file);
+
+            Ok(Partition {
+                writable: partition.writable,
+                label: partition.label.to_owned(),
+                path: format!("/proc/self/fd/{}", fd).into(),
+            })
+        })
+        .collect::<Result<_, Error>>()?;
+    let config_json = Config { partitions };
+
+    Ok((config_json, files))
+}