[microdroid] Refator payload apk verification

Refactoring the payload apk verification by removing some
unnecessary vectors created during the verification.

Bug: 271275206
Test: atest MicrodroidTests
Change-Id: Ia0fcca8d78015b1bdf3bd14e52deaed9ba05d869
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
index 6900ea5..b16a1e1 100644
--- a/microdroid_manager/src/instance.rs
+++ b/microdroid_manager/src/instance.rs
@@ -279,12 +279,24 @@
     pub apex_data: Vec<ApexData>,
 }
 
+impl MicrodroidData {
+    pub fn extra_apk_root_hash_eq(&self, i: usize, root_hash: &[u8]) -> bool {
+        self.extra_apks_data.get(i).map_or(false, |apk| apk.root_hash_eq(root_hash))
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct ApkData {
     pub root_hash: Box<RootHash>,
     pub pubkey: Box<[u8]>,
 }
 
+impl ApkData {
+    pub fn root_hash_eq(&self, root_hash: &[u8]) -> bool {
+        self.root_hash.as_ref() == root_hash
+    }
+}
+
 pub type RootHash = [u8];
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 8732be1..c78b20f 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -424,13 +424,12 @@
         .as_ref()
         .ok_or_else(|| MicrodroidError::InvalidConfig("No task in VM config".to_string()))?;
 
-    if config.extra_apks.len() != verified_data.extra_apks_data.len() {
-        return Err(anyhow!(
-            "config expects {} extra apks, but found {}",
-            config.extra_apks.len(),
-            verified_data.extra_apks_data.len()
-        ));
-    }
+    ensure!(
+        config.extra_apks.len() == verified_data.extra_apks_data.len(),
+        "config expects {} extra apks, but found {}",
+        config.extra_apks.len(),
+        verified_data.extra_apks_data.len()
+    );
     mount_extra_apks(&config, &mut zipfuse)?;
 
     // Wait until apex config is done. (e.g. linker configuration for apexes)
@@ -567,9 +566,9 @@
     let start_time = SystemTime::now();
 
     // Verify main APK
-    let root_hash = saved_data.map(|d| &d.apk_data.root_hash);
     let root_hash_from_idsig = get_apk_root_hash_from_idsig(MAIN_APK_IDSIG_PATH)?;
-    let root_hash_trustful = root_hash == Some(&root_hash_from_idsig);
+    let root_hash_trustful =
+        saved_data.map(|d| d.apk_data.root_hash_eq(root_hash_from_idsig.as_ref())).unwrap_or(false);
 
     // If root_hash can be trusted, pass it to apkdmverity so that it uses the passed root_hash
     // instead of the value read from the idsig file.
@@ -597,42 +596,36 @@
         sorted(glob(EXTRA_APK_PATH_PATTERN)?.collect::<Result<Vec<_>, _>>()?).collect::<Vec<_>>();
     let extra_idsigs =
         sorted(glob(EXTRA_IDSIG_PATH_PATTERN)?.collect::<Result<Vec<_>, _>>()?).collect::<Vec<_>>();
-    if extra_apks.len() != extra_idsigs.len() {
-        return Err(anyhow!(
-            "Extra apks/idsigs mismatch: {} apks but {} idsigs",
-            extra_apks.len(),
-            extra_idsigs.len()
-        ));
-    }
-    let extra_apks_count = extra_apks.len();
+    ensure!(
+        extra_apks.len() == extra_idsigs.len(),
+        "Extra apks/idsigs mismatch: {} apks but {} idsigs",
+        extra_apks.len(),
+        extra_idsigs.len()
+    );
 
-    let (extra_apk_names, extra_root_hashes_from_idsig): (Vec<_>, Vec<_>) = extra_idsigs
+    let extra_root_hashes_from_idsig: Vec<_> = extra_idsigs
         .iter()
-        .enumerate()
-        .map(|(i, extra_idsig)| {
-            (
-                format!("extra-apk-{}", i),
-                get_apk_root_hash_from_idsig(extra_idsig)
-                    .expect("Can't find root hash from extra idsig"),
-            )
-        })
-        .unzip();
-
-    let saved_extra_root_hashes: Vec<_> = saved_data
-        .map(|d| d.extra_apks_data.iter().map(|apk_data| &apk_data.root_hash).collect())
-        .unwrap_or_else(Vec::new);
-    let extra_root_hashes_trustful: Vec<_> = extra_root_hashes_from_idsig
-        .iter()
-        .enumerate()
-        .map(|(i, root_hash_from_idsig)| {
-            saved_extra_root_hashes.get(i).copied() == Some(root_hash_from_idsig)
+        .map(|idsig| {
+            get_apk_root_hash_from_idsig(idsig).expect("Can't find root hash from extra idsig")
         })
         .collect();
 
-    for i in 0..extra_apks_count {
+    let extra_root_hashes_trustful: Vec<_> = if let Some(data) = saved_data {
+        extra_root_hashes_from_idsig
+            .iter()
+            .enumerate()
+            .map(|(i, root_hash)| data.extra_apk_root_hash_eq(i, root_hash))
+            .collect()
+    } else {
+        vec![false; extra_root_hashes_from_idsig.len()]
+    };
+    let extra_apk_names: Vec<_> =
+        (0..extra_apks.len()).map(|i| format!("extra-apk-{}", i)).collect();
+
+    for (i, extra_apk) in extra_apks.iter().enumerate() {
         apkdmverity_arguments.push({
             ApkDmverityArgument {
-                apk: extra_apks[i].to_str().unwrap(),
+                apk: extra_apk.to_str().unwrap(),
                 idsig: extra_idsigs[i].to_str().unwrap(),
                 name: &extra_apk_names[i],
                 saved_root_hash: if extra_root_hashes_trustful[i] {