Merge "[LSC] Add LOCAL_LICENSE_KINDS to packages/modules/Virtualization"
diff --git a/apkdmverity/src/main.rs b/apkdmverity/src/main.rs
index cabeb35..b240c85 100644
--- a/apkdmverity/src/main.rs
+++ b/apkdmverity/src/main.rs
@@ -240,8 +240,7 @@
         }
 
         run_test(modified_apk.as_slice(), idsig.as_ref(), "incorrect_apk", |ctx| {
-            let ret = fs::read(&ctx.result.mapper_device).map_err(|e| e.kind());
-            assert_eq!(ret, Err(std::io::ErrorKind::Other));
+            fs::read(&ctx.result.mapper_device).expect_err("Should fail");
         });
     }
 
@@ -261,8 +260,7 @@
         }
 
         run_test(apk.as_ref(), modified_idsig.as_slice(), "incorrect_merkle_tree", |ctx| {
-            let ret = fs::read(&ctx.result.mapper_device).map_err(|e| e.kind());
-            assert_eq!(ret, Err(std::io::ErrorKind::Other));
+            fs::read(&ctx.result.mapper_device).expect_err("Should fail");
         });
     }
 
@@ -284,9 +282,7 @@
             // Read around the modified location causes an error
             let f = File::open(&ctx.result.mapper_device).unwrap();
             let mut buf = vec![0; 10]; // just read 10 bytes
-            let ret = f.read_at(&mut buf, MODIFIED_OFFSET).map_err(|e| e.kind());
-            assert!(ret.is_err());
-            assert_eq!(ret, Err(std::io::ErrorKind::Other));
+            f.read_at(&mut buf, MODIFIED_OFFSET).expect_err("Should fail");
         });
     }
 
@@ -359,11 +355,17 @@
         let apk = include_bytes!("../testdata/test.apk");
         let idsig = include_bytes!("../testdata/test.apk.idsig");
         let roothash = V4Signature::from(Cursor::new(&idsig)).unwrap().hashing_info.raw_root_hash;
-        run_test_with_hash(apk.as_ref(), idsig.as_ref(), "correct", Some(&roothash), |ctx| {
-            let verity = fs::read(&ctx.result.mapper_device).unwrap();
-            let original = fs::read(&ctx.result.data_device).unwrap();
-            assert_eq!(verity.len(), original.len()); // fail fast
-            assert_eq!(verity.as_slice(), original.as_slice());
-        });
+        run_test_with_hash(
+            apk.as_ref(),
+            idsig.as_ref(),
+            "correct_custom_roothash",
+            Some(&roothash),
+            |ctx| {
+                let verity = fs::read(&ctx.result.mapper_device).unwrap();
+                let original = fs::read(&ctx.result.data_device).unwrap();
+                assert_eq!(verity.len(), original.len()); // fail fast
+                assert_eq!(verity.as_slice(), original.as_slice());
+            },
+        );
     }
 }
diff --git a/compos/tests/java/android/compos/test/ComposKeyTestCase.java b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
index 6a3f755..64fd969 100644
--- a/compos/tests/java/android/compos/test/ComposKeyTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
@@ -159,7 +159,7 @@
                         apkName,
                         packageName,
                         "assets/vm_config.json",
-                        /* debug */ true);
+                        /* debug */ false);
         adbConnectToMicrodroid(getDevice(), mCid);
     }
 
diff --git a/libs/libavb_rs/Android.bp b/libs/libavb_rs/Android.bp
new file mode 100644
index 0000000..1035498
--- /dev/null
+++ b/libs/libavb_rs/Android.bp
@@ -0,0 +1,31 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_bindgen {
+    name: "libavb_bindgen",
+    wrapper_src: "bindgen/avb.h",
+    crate_name: "avb_bindgen",
+    source_stem: "bindings",
+    bindgen_flags: [
+        "--size_t-is-usize",
+        "--allowlist-function=.*",
+    ],
+    static_libs: [
+        "libavb",
+    ],
+    shared_libs: [
+        "libcrypto",
+    ],
+    cflags: ["-DBORINGSSL_NO_CXX"],
+}
+
+rust_test {
+    name: "libavb_bindgen_test",
+    srcs: [":libavb_bindgen"],
+    crate_name: "avb_bindgen_test",
+    test_suites: ["general-tests"],
+    auto_gen_config: true,
+    clippy_lints: "none",
+    lints: "none",
+}
diff --git a/libs/libavb_rs/bindgen/avb.h b/libs/libavb_rs/bindgen/avb.h
new file mode 100644
index 0000000..b3d5385
--- /dev/null
+++ b/libs/libavb_rs/bindgen/avb.h
@@ -0,0 +1,19 @@
+/*
+ * Copyright (C) 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.
+ */
+
+#pragma once
+
+#include <libavb/libavb.h>
diff --git a/microdroid/payload/metadata.proto b/microdroid/payload/metadata.proto
index 4c32dde..5ae2158 100644
--- a/microdroid/payload/metadata.proto
+++ b/microdroid/payload/metadata.proto
@@ -31,14 +31,13 @@
 
 message ApexPayload {
   // Required.
-  // The apex name.
   string name = 1;
-
   string partition_name = 2;
 
   // Optional.
-  // When specified, the public key used to sign the apex should match with it.
+  // When specified, apex payload should be verified with the public key and root digest.
   bytes public_key = 3;
+  bytes root_digest = 4;
 }
 
 message ApkPayload {
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index 9957689..721f9fa 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -13,6 +13,7 @@
         "android.system.virtualmachineservice-rust",
         "libanyhow",
         "libapkverify",
+        "libavb_bindgen",
         "libbinder_rpc_unstable_bindgen",
         "libbinder_rs",
         "libbyteorder",
@@ -63,4 +64,5 @@
             enabled: false,
         },
     },
+    data: ["tests/data/*"],
 }
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
index e431f51..47230e3 100644
--- a/microdroid_manager/src/instance.rs
+++ b/microdroid_manager/src/instance.rs
@@ -328,5 +328,6 @@
 #[derive(Debug, Serialize, Deserialize, PartialEq)]
 pub struct ApexData {
     pub name: String,
-    pub pubkey: Vec<u8>,
+    pub public_key: Vec<u8>,
+    pub root_digest: Vec<u8>,
 }
diff --git a/microdroid_manager/src/ioutil.rs b/microdroid_manager/src/ioutil.rs
index ab82e05..8ab2413 100644
--- a/microdroid_manager/src/ioutil.rs
+++ b/microdroid_manager/src/ioutil.rs
@@ -15,6 +15,8 @@
 //! IO utilities
 
 use anyhow::{anyhow, Result};
+use log::debug;
+use std::fmt::Debug;
 use std::fs::File;
 use std::io;
 use std::path::Path;
@@ -24,7 +26,8 @@
 const SLEEP_DURATION: Duration = Duration::from_millis(5);
 
 /// waits for a file with a timeout and returns it
-pub fn wait_for_file<P: AsRef<Path>>(path: P, timeout: Duration) -> Result<File> {
+pub fn wait_for_file<P: AsRef<Path> + Debug>(path: P, timeout: Duration) -> Result<File> {
+    debug!("waiting for {:?}...", path);
     let begin = Instant::now();
     loop {
         match File::open(&path) {
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index ac7adc9..9c1792d 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -88,13 +88,20 @@
     Ok(ret)
 }
 
-fn main() -> Result<()> {
+fn main() {
+    if let Err(e) = try_main() {
+        error!("failed with {}", e);
+        std::process::exit(1);
+    }
+}
+
+fn try_main() -> Result<()> {
     kernlog::init()?;
     info!("started.");
 
-    let metadata = load_metadata()?;
+    let metadata = load_metadata().context("Failed to load payload metadata")?;
 
-    let mut instance = InstanceDisk::new()?;
+    let mut instance = InstanceDisk::new().context("Failed to load instance.img")?;
     let saved_data = instance.read_microdroid_data().context("Failed to read identity data")?;
 
     // Verify the payload before using it.
@@ -162,14 +169,15 @@
     // Start apkdmverity and wait for the dm-verify block
     system_properties::write("ctl.start", "apkdmverity")?;
 
-    // While waiting for apkdmverity to mount APK, gathers APEX pubkeys from payload.
+    // While waiting for apkdmverity to mount APK, gathers public keys and root digests from
+    // APEX payload.
     let apex_data_from_payload = get_apex_data_from_payload(metadata)?;
     if let Some(saved_data) = saved_data.map(|d| &d.apex_data) {
-        // For APEX payload, we don't support updating their pubkeys
+        // We don't support APEX updates. (assuming that update will change root digest)
         ensure!(saved_data == &apex_data_from_payload, "APEX payloads has changed.");
         let apex_metadata = to_metadata(&apex_data_from_payload);
-        // Pass metadata(with pubkeys) to apexd so that it uses the passed metadata
-        // instead of the default one (/dev/block/by-name/payload-metadata)
+        // Pass metadata(with public keys and root digests) to apexd so that it uses the passed
+        // metadata instead of the default one (/dev/block/by-name/payload-metadata)
         OpenOptions::new()
             .create_new(true)
             .write(true)
diff --git a/microdroid_manager/src/payload.rs b/microdroid_manager/src/payload.rs
index bf9d9f9..8ec6f74 100644
--- a/microdroid_manager/src/payload.rs
+++ b/microdroid_manager/src/payload.rs
@@ -17,14 +17,11 @@
 use crate::instance::ApexData;
 use crate::ioutil::wait_for_file;
 use anyhow::Result;
+use apex::verify;
 use log::info;
 use microdroid_metadata::{read_metadata, ApexPayload, Metadata};
-use std::fs::File;
-use std::io::Read;
 use std::time::Duration;
-use zip::ZipArchive;
 
-const APEX_PUBKEY_ENTRY: &str = "apex_pubkey";
 const PAYLOAD_METADATA_PATH: &str = "/dev/block/by-name/payload-metadata";
 const WAIT_TIMEOUT: Duration = Duration::from_secs(10);
 
@@ -35,29 +32,20 @@
     read_metadata(file)
 }
 
-/// Loads (name, pubkey) from payload APEXes
+/// Loads (name, public_key, root_digest) from payload APEXes
 pub fn get_apex_data_from_payload(metadata: &Metadata) -> Result<Vec<ApexData>> {
     metadata
         .apexes
         .iter()
         .map(|apex| {
             let name = apex.name.clone();
-            let partition = format!("/dev/block/by-name/{}", apex.partition_name);
-            let pubkey = get_pubkey_from_apex(&partition)?;
-            Ok(ApexData { name, pubkey })
+            let apex_path = format!("/dev/block/by-name/{}", apex.partition_name);
+            let result = verify(&apex_path)?;
+            Ok(ApexData { name, public_key: result.public_key, root_digest: result.root_digest })
         })
         .collect()
 }
 
-fn get_pubkey_from_apex(path: &str) -> Result<Vec<u8>> {
-    let f = File::open(path)?;
-    let mut z = ZipArchive::new(f)?;
-    let mut pubkey_file = z.by_name(APEX_PUBKEY_ENTRY)?;
-    let mut pubkey = Vec::new();
-    pubkey_file.read_to_end(&mut pubkey)?;
-    Ok(pubkey)
-}
-
 /// Convert vector of ApexData into Metadata
 pub fn to_metadata(apex_data: &[ApexData]) -> Metadata {
     Metadata {
@@ -65,10 +53,13 @@
             .iter()
             .map(|data| ApexPayload {
                 name: data.name.clone(),
-                public_key: data.pubkey.clone(),
+                public_key: data.public_key.clone(),
+                root_digest: data.root_digest.clone(),
                 ..Default::default()
             })
             .collect(),
         ..Default::default()
     }
 }
+
+mod apex;
diff --git a/microdroid_manager/src/payload/apex.rs b/microdroid_manager/src/payload/apex.rs
new file mode 100644
index 0000000..24c4f05
--- /dev/null
+++ b/microdroid_manager/src/payload/apex.rs
@@ -0,0 +1,225 @@
+// 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.
+
+//! Routines for handling APEX payload
+
+use anyhow::{anyhow, ensure, Result};
+use avb_bindgen::*;
+use std::ffi::{c_void, CStr};
+use std::fs::File;
+use std::io::{Read, Seek, SeekFrom};
+use std::mem::{size_of, zeroed};
+use std::ops::Deref;
+use std::ptr::null_mut;
+use std::slice::{from_raw_parts, from_raw_parts_mut};
+use zip::ZipArchive;
+
+const APEX_PUBKEY_ENTRY: &str = "apex_pubkey";
+const APEX_PAYLOAD_ENTRY: &str = "apex_payload.img";
+
+/// Verification result holds public key and root digest of apex_payload.img
+pub struct ApexVerificationResult {
+    pub public_key: Vec<u8>,
+    pub root_digest: Vec<u8>,
+}
+
+/// Verify APEX payload by AVB verification and return public key and root digest
+pub fn verify(path: &str) -> Result<ApexVerificationResult> {
+    let apex_file = File::open(path)?;
+    let (public_key, image_offset, image_size) = get_public_key_and_image_info(&apex_file)?;
+    let root_digest = verify_vbmeta(apex_file, image_offset, image_size, &public_key)?;
+    Ok(ApexVerificationResult { public_key, root_digest })
+}
+
+fn get_public_key_and_image_info(apex_file: &File) -> Result<(Vec<u8>, u64, u64)> {
+    let mut z = ZipArchive::new(apex_file)?;
+
+    let mut public_key = Vec::new();
+    z.by_name(APEX_PUBKEY_ENTRY)?.read_to_end(&mut public_key)?;
+
+    let (image_offset, image_size) =
+        z.by_name(APEX_PAYLOAD_ENTRY).map(|f| (f.data_start(), f.size()))?;
+
+    Ok((public_key, image_offset, image_size))
+}
+
+// Manual addition of a missing enum
+#[allow(non_camel_case_types, dead_code)]
+#[repr(u8)]
+enum AvbDescriptorTag {
+    AVB_DESCRIPTOR_TAG_PROPERTY = 0,
+    AVB_DESCRIPTOR_TAG_HASHTREE,
+    AVB_DESCRIPTOR_TAG_HASH,
+    AVB_DESCRIPTOR_TAG_KERNEL_CMDLINE,
+    AVB_DESCRIPTOR_TAG_CHAIN_PARTITION,
+}
+
+const FOOTER_SIZE: usize = size_of::<AvbFooter>();
+const HASHTREE_DESCRIPTOR_SIZE: usize = size_of::<AvbHashtreeDescriptor>();
+
+/// Verify VBmeta image and return root digest
+fn verify_vbmeta<R: Read + Seek>(
+    image: R,
+    offset: u64,
+    size: u64,
+    public_key: &[u8],
+) -> Result<Vec<u8>> {
+    let vbmeta = VbMeta::from(image, offset, size)?;
+    vbmeta.verify(public_key)?;
+    for &descriptor in vbmeta.descriptors()?.iter() {
+        if let Ok(hashtree_descriptor) = HashtreeDescriptor::from(descriptor) {
+            return hashtree_descriptor.root_digest();
+        }
+    }
+    Err(anyhow!("HashtreeDescriptor is not found."))
+}
+
+struct VbMeta {
+    data: Vec<u8>,
+}
+
+impl VbMeta {
+    // Read a VbMeta data from a given image
+    fn from<R: Read + Seek>(mut image: R, offset: u64, size: u64) -> Result<VbMeta> {
+        // Get AvbFooter first
+        image.seek(SeekFrom::Start(offset + size - FOOTER_SIZE as u64))?;
+        // SAFETY: AvbDescriptor is a "repr(C,packed)" struct from bindgen
+        let mut footer: AvbFooter = unsafe { zeroed() };
+        // SAFETY: safe to read because of seek(-FOOTER_SIZE) above
+        unsafe {
+            let footer_slice = from_raw_parts_mut(&mut footer as *mut _ as *mut u8, FOOTER_SIZE);
+            image.read_exact(footer_slice)?;
+            ensure!(avb_footer_validate_and_byteswap(&footer, &mut footer));
+        }
+        // Get VbMeta block
+        image.seek(SeekFrom::Start(offset + footer.vbmeta_offset))?;
+        let vbmeta_size = footer.vbmeta_size as usize;
+        let mut data = vec![0u8; vbmeta_size];
+        image.read_exact(&mut data)?;
+        Ok(VbMeta { data })
+    }
+    // Verify VbMeta image. Its enclosed public key should match with a given public key.
+    fn verify(&self, outer_public_key: &[u8]) -> Result<()> {
+        // SAFETY: self.data points to a valid VBMeta data and avb_vbmeta_image_verify should work fine
+        // with it
+        let public_key = unsafe {
+            let mut pk_ptr: *const u8 = null_mut();
+            let mut pk_len: usize = 0;
+            let res = avb_vbmeta_image_verify(
+                self.data.as_ptr(),
+                self.data.len(),
+                &mut pk_ptr,
+                &mut pk_len,
+            );
+            ensure!(
+                res == AvbVBMetaVerifyResult_AVB_VBMETA_VERIFY_RESULT_OK,
+                CStr::from_ptr(avb_vbmeta_verify_result_to_string(res))
+                    .to_string_lossy()
+                    .into_owned()
+            );
+            from_raw_parts(pk_ptr, pk_len)
+        };
+
+        ensure!(public_key == outer_public_key, "Public key mismatch with a given one.");
+        Ok(())
+    }
+    // Return a slice of AvbDescriptor pointers
+    fn descriptors(&self) -> Result<Descriptors> {
+        let mut num: usize = 0;
+        // SAFETY: ptr will be freed by Descriptor.
+        Ok(unsafe {
+            let ptr = avb_descriptor_get_all(self.data.as_ptr(), self.data.len(), &mut num);
+            ensure!(!ptr.is_null(), "VbMeta has no descriptors.");
+            let all = from_raw_parts(ptr, num);
+            Descriptors { ptr, all }
+        })
+    }
+}
+
+struct HashtreeDescriptor {
+    ptr: *const u8,
+    inner: AvbHashtreeDescriptor,
+}
+
+impl HashtreeDescriptor {
+    fn from(descriptor: *const AvbDescriptor) -> Result<HashtreeDescriptor> {
+        // SAFETY: AvbDescriptor is a "repr(C,packed)" struct from bindgen
+        let mut desc: AvbDescriptor = unsafe { zeroed() };
+        // SAFETY: both points to valid AvbDescriptor pointers
+        unsafe {
+            ensure!(avb_descriptor_validate_and_byteswap(descriptor, &mut desc));
+        }
+        ensure!({ desc.tag } == AvbDescriptorTag::AVB_DESCRIPTOR_TAG_HASHTREE as u64);
+        // SAFETY: AvbHashtreeDescriptor is a "repr(C, packed)" struct from bindgen
+        let mut hashtree_descriptor: AvbHashtreeDescriptor = unsafe { zeroed() };
+        // SAFETY: With tag == AVB_DESCRIPTOR_TAG_HASHTREE, descriptor should point to
+        // a AvbHashtreeDescriptor.
+        unsafe {
+            ensure!(avb_hashtree_descriptor_validate_and_byteswap(
+                descriptor as *const AvbHashtreeDescriptor,
+                &mut hashtree_descriptor,
+            ));
+        }
+        Ok(Self { ptr: descriptor as *const u8, inner: hashtree_descriptor })
+    }
+    fn root_digest(&self) -> Result<Vec<u8>> {
+        // SAFETY: digest_ptr should point to a valid buffer of root_digest_len
+        let root_digest = unsafe {
+            let digest_ptr = self.ptr.offset(
+                HASHTREE_DESCRIPTOR_SIZE as isize
+                    + self.inner.partition_name_len as isize
+                    + self.inner.salt_len as isize,
+            );
+            from_raw_parts(digest_ptr, self.inner.root_digest_len as usize)
+        };
+        Ok(root_digest.to_owned())
+    }
+}
+
+// Wraps pointer to a heap-allocated array of AvbDescriptor pointers
+struct Descriptors<'a> {
+    ptr: *mut *const AvbDescriptor,
+    all: &'a [*const AvbDescriptor],
+}
+
+// Wrapped pointer should be freed with avb_free.
+impl Drop for Descriptors<'_> {
+    fn drop(&mut self) {
+        // SAFETY: ptr is allocated by avb_descriptor_get_all
+        unsafe { avb_free(self.ptr as *mut c_void) }
+    }
+}
+
+impl<'a> Deref for Descriptors<'a> {
+    type Target = &'a [*const AvbDescriptor];
+    fn deref(&self) -> &Self::Target {
+        &self.all
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    fn to_hex_string(buf: &[u8]) -> String {
+        buf.iter().map(|b| format!("{:02x}", b)).collect()
+    }
+    #[test]
+    fn test_open_apex() {
+        let res = verify("tests/data/test.apex").unwrap();
+        assert_eq!(
+            to_hex_string(&res.root_digest),
+            "fe11ab17da0a3a738b54bdc3a13f6139cbdf91ec32f001f8d4bbbf8938e04e39"
+        );
+    }
+}
diff --git a/microdroid_manager/tests/data/README.md b/microdroid_manager/tests/data/README.md
new file mode 100644
index 0000000..82ebec6
--- /dev/null
+++ b/microdroid_manager/tests/data/README.md
@@ -0,0 +1,3 @@
+# Test data
+
+- test.apex: copied from system/apexshim/prebuilts/x86/com.android.apex.cts.shim.v1.apex
\ No newline at end of file
diff --git a/microdroid_manager/tests/data/test.apex b/microdroid_manager/tests/data/test.apex
new file mode 100644
index 0000000..fd79365
--- /dev/null
+++ b/microdroid_manager/tests/data/test.apex
Binary files differ
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index 97cd426..4c8f5eb 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -61,6 +61,9 @@
 
         // disconnect from microdroid
         tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
+
+        // remove any leftover files under test root
+        android.tryRun("rm", "-rf", TEST_ROOT + "*");
     }
 
     public static void cleanUpVirtualizationTestSetup(ITestDevice androidDevice)