Make DTBO for VM containing vendor pubkey when vendor image exists

Bug: 285855436
Test: adb shell /apex/com.android.virt/bin/vm run-microdroid --vendor /data/local/tmp/test_microdroid_vendor_image.img
Test: atest virtualizationmanager_device_test

Change-Id: Ib1b8301d572d2546350fb1e5e9b11060b430b756
diff --git a/libs/vbmeta/Android.bp b/libs/vbmeta/Android.bp
index ae83703..4fb6ae4 100644
--- a/libs/vbmeta/Android.bp
+++ b/libs/vbmeta/Android.bp
@@ -18,6 +18,9 @@
 rust_library {
     name: "libvbmeta_rust",
     defaults: ["libvbmeta_rust.defaults"],
+    apex_available: [
+        "com.android.virt",
+    ],
 }
 
 rust_test_host {
diff --git a/virtualizationmanager/Android.bp b/virtualizationmanager/Android.bp
index 12d8724..33897b2 100644
--- a/virtualizationmanager/Android.bp
+++ b/virtualizationmanager/Android.bp
@@ -54,6 +54,7 @@
         "libshared_child",
         "libstatslog_virtualization_rust",
         "libtombstoned_client_rust",
+        "libvbmeta_rust",
         "libvm_control",
         "libvmconfig",
         "libzip",
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 4024b04..70e9be9 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -62,6 +62,7 @@
 };
 use disk::QcowFile;
 use lazy_static::lazy_static;
+use libfdt::Fdt;
 use log::{debug, error, info, warn};
 use microdroid_payload_config::{OsConfig, Task, TaskType, VmPayloadConfig};
 use nix::unistd::pipe;
@@ -70,7 +71,7 @@
 use semver::VersionReq;
 use std::collections::HashSet;
 use std::convert::TryInto;
-use std::ffi::CStr;
+use std::ffi::{CStr, CString};
 use std::fs::{canonicalize, read_dir, remove_file, File, OpenOptions};
 use std::io::{BufRead, BufReader, Error, ErrorKind, Write};
 use std::num::{NonZeroU16, NonZeroU32};
@@ -78,6 +79,7 @@
 use std::os::unix::raw::pid_t;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, Weak};
+use vbmeta::VbMetaImage;
 use vmconfig::VmConfig;
 use vsock::VsockStream;
 use zip::ZipArchive;
@@ -101,6 +103,9 @@
 
 const UNFORMATTED_STORAGE_MAGIC: &str = "UNFORMATTED-STORAGE";
 
+/// Roughly estimated sufficient size for storing vendor public key into DTBO.
+const EMPTY_VENDOR_DT_OVERLAY_BUF_SIZE: usize = 10000;
+
 /// crosvm requires all partitions to be a multiple of 4KiB.
 const PARTITION_GRANULARITY_BYTES: u64 = 4096;
 
@@ -361,6 +366,22 @@
             check_gdb_allowed(config)?;
         }
 
+        let vendor_public_key = extract_vendor_public_key(config)
+            .context("Failed to extract vendor public key")
+            .or_service_specific_exception(-1)?;
+        let dtbo_vendor = if let Some(vendor_public_key) = vendor_public_key {
+            let dtbo_for_vendor_image = temporary_directory.join("dtbo_vendor");
+            create_dtbo_for_vendor_image(&vendor_public_key, &dtbo_for_vendor_image)
+                .context("Failed to write vendor_public_key")
+                .or_service_specific_exception(-1)?;
+            let file = File::open(dtbo_for_vendor_image)
+                .context("Failed to open dtbo_vendor")
+                .or_service_specific_exception(-1)?;
+            Some(file)
+        } else {
+            None
+        };
+
         let debug_level = match config {
             VirtualMachineConfig::AppConfig(config) => config.debugLevel,
             _ => DebugLevel::NONE,
@@ -506,6 +527,7 @@
             detect_hangup: is_app_config,
             gdb_port,
             vfio_devices,
+            dtbo_vendor,
         };
         let instance = Arc::new(
             VmInstance::new(
@@ -524,6 +546,68 @@
     }
 }
 
+fn extract_vendor_public_key(config: &VirtualMachineConfig) -> Result<Option<Vec<u8>>> {
+    let VirtualMachineConfig::AppConfig(config) = config else {
+        return Ok(None);
+    };
+    let Some(custom_config) = &config.customConfig else {
+        return Ok(None);
+    };
+    let Some(file) = custom_config.vendorImage.as_ref() else {
+        return Ok(None);
+    };
+
+    let file = clone_file(file)?;
+    let size = file.metadata().context("Failed to get metadata from microdroid-vendor.img")?.len();
+    let vbmeta = VbMetaImage::verify_reader_region(&file, 0, size)
+        .context("Failed to get vbmeta from microdroid-vendor.img")?;
+    let vendor_public_key = vbmeta
+        .public_key()
+        .ok_or(anyhow!("No public key is extracted from microdroid-vendor.img"))?
+        .to_vec();
+
+    Ok(Some(vendor_public_key))
+}
+
+fn create_dtbo_for_vendor_image(vendor_public_key: &[u8], dtbo: &PathBuf) -> Result<()> {
+    if dtbo.exists() {
+        return Err(anyhow!("DTBO file already exists"));
+    }
+
+    let mut buf = vec![0; EMPTY_VENDOR_DT_OVERLAY_BUF_SIZE];
+    let fdt = Fdt::create_empty_tree(buf.as_mut_slice())
+        .map_err(|e| anyhow!("Failed to create FDT: {:?}", e))?;
+    let mut root = fdt.root_mut().map_err(|e| anyhow!("Failed to get root node: {:?}", e))?;
+
+    let fragment_node_name = CString::new("fragment@0")?;
+    let mut fragment_node = root
+        .add_subnode(fragment_node_name.as_c_str())
+        .map_err(|e| anyhow!("Failed to create fragment node: {:?}", e))?;
+    let target_path_prop_name = CString::new("target-path")?;
+    let target_path = CString::new("/")?;
+    fragment_node
+        .setprop(target_path_prop_name.as_c_str(), target_path.to_bytes_with_nul())
+        .map_err(|e| anyhow!("Failed to set target-path: {:?}", e))?;
+    let overlay_node_name = CString::new("__overlay__")?;
+    let mut overlay_node = fragment_node
+        .add_subnode(overlay_node_name.as_c_str())
+        .map_err(|e| anyhow!("Failed to create overlay node: {:?}", e))?;
+
+    let avf_node_name = CString::new("avf")?;
+    let mut avf_node = overlay_node
+        .add_subnode(avf_node_name.as_c_str())
+        .map_err(|e| anyhow!("Failed to create avf node: {:?}", e))?;
+    let vendor_public_key_name = CString::new("vendor_public_key")?;
+    avf_node
+        .setprop(vendor_public_key_name.as_c_str(), vendor_public_key)
+        .map_err(|e| anyhow!("Failed to set avf/vendor_public_key: {:?}", e))?;
+
+    fdt.pack().map_err(|e| anyhow!("Failed to pack fdt: {:?}", e))?;
+    let mut file = File::create(dtbo)?;
+    file.write_all(fdt.as_slice())?;
+    Ok(file.flush()?)
+}
+
 fn write_zero_filler(zero_filler_path: &Path) -> Result<()> {
     let file = OpenOptions::new()
         .create_new(true)
@@ -1375,4 +1459,58 @@
         append_kernel_param("bar=42", &mut vm_config);
         assert_eq!(vm_config.params, Some("foo=5 bar=42".to_owned()))
     }
+
+    #[test]
+    fn test_create_dtbo_for_vendor_image() -> Result<()> {
+        let vendor_public_key = String::from("foo");
+        let vendor_public_key = vendor_public_key.as_bytes();
+
+        let tmp_dir = tempfile::TempDir::new()?;
+        let dtbo_path = tmp_dir.path().to_path_buf().join("bar");
+
+        create_dtbo_for_vendor_image(vendor_public_key, &dtbo_path)?;
+
+        let data = std::fs::read(dtbo_path)?;
+        let fdt = Fdt::from_slice(&data).unwrap();
+
+        let fragment_node_path = CString::new("/fragment@0")?;
+        let fragment_node = fdt.node(fragment_node_path.as_c_str()).unwrap();
+        let Some(fragment_node) = fragment_node else {
+            bail!("fragment_node shouldn't be None.");
+        };
+        let target_path_prop_name = CString::new("target-path")?;
+        let target_path_from_dtbo =
+            fragment_node.getprop(target_path_prop_name.as_c_str()).unwrap();
+        let target_path_expected = CString::new("/")?;
+        assert_eq!(target_path_from_dtbo, Some(target_path_expected.to_bytes_with_nul()));
+
+        let avf_node_path = CString::new("/fragment@0/__overlay__/avf")?;
+        let avf_node = fdt.node(avf_node_path.as_c_str()).unwrap();
+        let Some(avf_node) = avf_node else {
+            bail!("avf_node shouldn't be None.");
+        };
+        let vendor_public_key_name = CString::new("vendor_public_key")?;
+        let key_from_dtbo = avf_node.getprop(vendor_public_key_name.as_c_str()).unwrap();
+        assert_eq!(key_from_dtbo, Some(vendor_public_key));
+
+        tmp_dir.close()?;
+        Ok(())
+    }
+
+    #[test]
+    fn test_create_dtbo_for_vendor_image_throws_error_if_already_exists() -> Result<()> {
+        let vendor_public_key = String::from("foo");
+        let vendor_public_key = vendor_public_key.as_bytes();
+
+        let tmp_dir = tempfile::TempDir::new()?;
+        let dtbo_path = tmp_dir.path().to_path_buf().join("bar");
+
+        create_dtbo_for_vendor_image(vendor_public_key, &dtbo_path)?;
+
+        let ret_second_trial = create_dtbo_for_vendor_image(vendor_public_key, &dtbo_path);
+        assert!(ret_second_trial.is_err(), "should fail");
+
+        tmp_dir.close()?;
+        Ok(())
+    }
 }
diff --git a/virtualizationmanager/src/crosvm.rs b/virtualizationmanager/src/crosvm.rs
index bb6066f..b842574 100644
--- a/virtualizationmanager/src/crosvm.rs
+++ b/virtualizationmanager/src/crosvm.rs
@@ -116,6 +116,7 @@
     pub detect_hangup: bool,
     pub gdb_port: Option<NonZeroU16>,
     pub vfio_devices: Vec<VfioDevice>,
+    pub dtbo_vendor: Option<File>,
 }
 
 /// A disk image to pass to crosvm for a VM.
@@ -886,6 +887,8 @@
         .arg("--socket")
         .arg(add_preserved_fd(&mut preserved_fds, &control_server_socket.as_raw_descriptor()));
 
+    // TODO(b/285855436): Pass dtbo_vendor after --device-tree-overlay crosvm option is supported.
+
     append_platform_devices(&mut command, &config)?;
 
     debug!("Preserving FDs {:?}", preserved_fds);