[avb] Verify vbmeta for payload with no initrd

Verifies that vbmeta doesn't have initrd hash descriptors when
payload has initrd == None.

Bug: 256148034
Test: m pvmfw_img && atest libpvmfw_avb.integration_test
Change-Id: I20a0d4798f535fa78fc5d56415653a860d1d2ae9
diff --git a/pvmfw/avb/Android.bp b/pvmfw/avb/Android.bp
index 1c5e91a..0527dfb 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -36,6 +36,8 @@
         ":microdroid_initrd_normal",
         ":microdroid_initrd_debuggable",
         ":test_image_with_one_hashdesc",
+        ":test_image_with_non_initrd_hashdesc",
+        ":test_image_with_prop_desc",
         ":unsigned_test_image",
     ],
     prefer_rlib: true,
@@ -64,6 +66,38 @@
     cmd: "$(location avbtool) generate_test_image --image_size 16384 --output $(out)",
 }
 
+avb_gen_vbmeta_image {
+    name: "test_non_initrd_hashdesc",
+    src: ":unsigned_test_image",
+    partition_name: "non_initrd11",
+    salt: "2222",
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_non_initrd_hashdesc",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "1111",
+    include_descriptors_from_images: [
+        ":test_non_initrd_hashdesc",
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_prop_desc",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "1111",
+    props: [
+        {
+            name: "mock_prop",
+            value: "3333",
+        },
+    ],
+}
+
 avb_add_hash_footer {
     name: "test_image_with_one_hashdesc",
     src: ":unsigned_test_image",
diff --git a/pvmfw/avb/src/verify.rs b/pvmfw/avb/src/verify.rs
index 9ba67bf..664da27 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -15,15 +15,21 @@
 //! This module handles the pvmfw payload verification.
 
 use crate::error::{slot_verify_result_to_verify_payload_result, AvbSlotVerifyError};
-use avb_bindgen::{avb_slot_verify, AvbHashtreeErrorMode, AvbIOResult, AvbOps, AvbSlotVerifyFlags};
+use avb_bindgen::{
+    avb_descriptor_foreach, avb_hash_descriptor_validate_and_byteswap, avb_slot_verify,
+    avb_slot_verify_data_free, AvbDescriptor, AvbHashDescriptor, AvbHashtreeErrorMode, AvbIOResult,
+    AvbOps, AvbSlotVerifyData, AvbSlotVerifyFlags, AvbVBMetaData,
+};
 use core::{
     ffi::{c_char, c_void, CStr},
+    mem::{size_of, MaybeUninit},
     ptr::{self, NonNull},
     slice,
 };
 
 const NULL_BYTE: &[u8] = b"\0";
 
+#[derive(Debug)]
 enum AvbIOError {
     /// AVB_IO_RESULT_ERROR_OOM,
     #[allow(dead_code)]
@@ -234,6 +240,89 @@
     write(out_is_trusted, public_key == trusted_public_key)
 }
 
+extern "C" fn search_initrd_hash_descriptor(
+    descriptor: *const AvbDescriptor,
+    user_data: *mut c_void,
+) -> bool {
+    try_search_initrd_hash_descriptor(descriptor, user_data).is_ok()
+}
+
+fn try_search_initrd_hash_descriptor(
+    descriptor: *const AvbDescriptor,
+    user_data: *mut c_void,
+) -> Result<(), AvbIOError> {
+    let hash_desc = AvbHashDescriptorRef::try_from(descriptor)?;
+    if matches!(
+        hash_desc.partition_name()?.try_into(),
+        Ok(PartitionName::InitrdDebug) | Ok(PartitionName::InitrdNormal),
+    ) {
+        write(user_data as *mut bool, true)?;
+    }
+    Ok(())
+}
+
+/// `hash_desc` only contains the metadata like fields length and flags of the descriptor.
+/// The data itself is contained in `ptr`.
+struct AvbHashDescriptorRef {
+    hash_desc: AvbHashDescriptor,
+    ptr: *const AvbDescriptor,
+}
+
+impl TryFrom<*const AvbDescriptor> for AvbHashDescriptorRef {
+    type Error = AvbIOError;
+
+    fn try_from(descriptor: *const AvbDescriptor) -> Result<Self, Self::Error> {
+        is_not_null(descriptor)?;
+        // SAFETY: It is safe as the raw pointer `descriptor` is a nonnull pointer and
+        // we have validated that it is of hash descriptor type.
+        let hash_desc = unsafe {
+            let mut desc = MaybeUninit::uninit();
+            if !avb_hash_descriptor_validate_and_byteswap(
+                descriptor as *const AvbHashDescriptor,
+                desc.as_mut_ptr(),
+            ) {
+                return Err(AvbIOError::Io);
+            }
+            desc.assume_init()
+        };
+        Ok(Self { hash_desc, ptr: descriptor })
+    }
+}
+
+impl AvbHashDescriptorRef {
+    fn check_is_in_range(&self, index: usize) -> Result<(), AvbIOError> {
+        let parent_desc = self.hash_desc.parent_descriptor;
+        let total_len = usize_checked_add(
+            size_of::<AvbDescriptor>(),
+            to_usize(parent_desc.num_bytes_following)?,
+        )?;
+        if index <= total_len {
+            Ok(())
+        } else {
+            Err(AvbIOError::Io)
+        }
+    }
+
+    /// Returns the non null-terminated partition name.
+    fn partition_name(&self) -> Result<&[u8], AvbIOError> {
+        let partition_name_offset = size_of::<AvbHashDescriptor>();
+        let partition_name_len = to_usize(self.hash_desc.partition_name_len)?;
+        self.check_is_in_range(usize_checked_add(partition_name_offset, partition_name_len)?)?;
+        let desc = self.ptr as *const u8;
+        // SAFETY: The descriptor has been validated as nonnull and the partition name is
+        // contained within the image.
+        unsafe { Ok(slice::from_raw_parts(desc.add(partition_name_offset), partition_name_len)) }
+    }
+}
+
+fn to_usize<T: TryInto<usize>>(num: T) -> Result<usize, AvbIOError> {
+    num.try_into().map_err(|_| AvbIOError::InvalidValueSize)
+}
+
+fn usize_checked_add(x: usize, y: usize) -> Result<usize, AvbIOError> {
+    x.checked_add(y).ok_or(AvbIOError::InvalidValueSize)
+}
+
 fn write<T>(ptr: *mut T, value: T) -> Result<(), AvbIOError> {
     let ptr = to_nonnull(ptr)?;
     // SAFETY: It is safe as the raw pointer `ptr` is a nonnull pointer.
@@ -274,12 +363,20 @@
     const INITRD_DEBUG_PARTITION_NAME: &[u8] = b"initrd_debug\0";
 
     fn as_cstr(&self) -> &CStr {
-        let partition_name = match self {
+        CStr::from_bytes_with_nul(self.as_bytes()).unwrap()
+    }
+
+    fn as_non_null_terminated_bytes(&self) -> &[u8] {
+        let partition_name = self.as_bytes();
+        &partition_name[..partition_name.len() - 1]
+    }
+
+    fn as_bytes(&self) -> &[u8] {
+        match self {
             Self::Kernel => Self::KERNEL_PARTITION_NAME,
             Self::InitrdNormal => Self::INITRD_NORMAL_PARTITION_NAME,
             Self::InitrdDebug => Self::INITRD_DEBUG_PARTITION_NAME,
-        };
-        CStr::from_bytes_with_nul(partition_name).unwrap()
+        }
     }
 }
 
@@ -296,6 +393,59 @@
     }
 }
 
+impl TryFrom<&[u8]> for PartitionName {
+    type Error = AvbIOError;
+
+    fn try_from(non_null_terminated_name: &[u8]) -> Result<Self, Self::Error> {
+        match non_null_terminated_name {
+            x if x == Self::Kernel.as_non_null_terminated_bytes() => Ok(Self::Kernel),
+            x if x == Self::InitrdNormal.as_non_null_terminated_bytes() => Ok(Self::InitrdNormal),
+            x if x == Self::InitrdDebug.as_non_null_terminated_bytes() => Ok(Self::InitrdDebug),
+            _ => Err(AvbIOError::NoSuchPartition),
+        }
+    }
+}
+
+struct AvbSlotVerifyDataWrap(*mut AvbSlotVerifyData);
+
+impl TryFrom<*mut AvbSlotVerifyData> for AvbSlotVerifyDataWrap {
+    type Error = AvbSlotVerifyError;
+
+    fn try_from(data: *mut AvbSlotVerifyData) -> Result<Self, Self::Error> {
+        is_not_null(data).map_err(|_| AvbSlotVerifyError::Io)?;
+        Ok(Self(data))
+    }
+}
+
+impl Drop for AvbSlotVerifyDataWrap {
+    fn drop(&mut self) {
+        // SAFETY: This is safe because `self.0` is checked nonnull when the
+        // instance is created. We can free this pointer when the instance is
+        // no longer needed.
+        unsafe {
+            avb_slot_verify_data_free(self.0);
+        }
+    }
+}
+
+impl AsRef<AvbSlotVerifyData> for AvbSlotVerifyDataWrap {
+    fn as_ref(&self) -> &AvbSlotVerifyData {
+        // This is safe because `self.0` is checked nonnull when the instance is created.
+        as_ref(self.0).unwrap()
+    }
+}
+
+impl AvbSlotVerifyDataWrap {
+    fn vbmeta_images(&self) -> Result<&[AvbVBMetaData], AvbSlotVerifyError> {
+        let data = self.as_ref();
+        is_not_null(data.vbmeta_images).map_err(|_| AvbSlotVerifyError::Io)?;
+        // SAFETY: It is safe as the raw pointer `data.vbmeta_images` is a nonnull pointer.
+        let vbmeta_images =
+            unsafe { slice::from_raw_parts(data.vbmeta_images, data.num_vbmeta_images) };
+        Ok(vbmeta_images)
+    }
+}
+
 struct Payload<'a> {
     kernel: &'a [u8],
     initrd: Option<&'a [u8]>,
@@ -328,7 +478,10 @@
         }
     }
 
-    fn verify_partitions(&mut self, partition_names: &[&CStr]) -> Result<(), AvbSlotVerifyError> {
+    fn verify_partitions(
+        &mut self,
+        partition_names: &[&CStr],
+    ) -> Result<AvbSlotVerifyDataWrap, AvbSlotVerifyError> {
         if partition_names.len() > Self::MAX_NUM_OF_HASH_DESCRIPTORS {
             return Err(AvbSlotVerifyError::InvalidArgument);
         }
@@ -356,7 +509,7 @@
             validate_public_key_for_partition: Some(validate_public_key_for_partition),
         };
         let ab_suffix = CStr::from_bytes_with_nul(NULL_BYTE).unwrap();
-        let out_data = ptr::null_mut();
+        let mut out_data = MaybeUninit::uninit();
         // SAFETY: It is safe to call `avb_slot_verify()` as the pointer arguments (`ops`,
         // `requested_partitions` and `ab_suffix`) passed to the method are all valid and
         // initialized. The last argument `out_data` is allowed to be null so that nothing
@@ -368,10 +521,37 @@
                 ab_suffix.as_ptr(),
                 AvbSlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION,
                 AvbHashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_RESTART_AND_INVALIDATE,
-                out_data,
+                out_data.as_mut_ptr(),
             )
         };
-        slot_verify_result_to_verify_payload_result(result)
+        slot_verify_result_to_verify_payload_result(result)?;
+        // SAFETY: This is safe because `out_data` has been properly initialized after
+        // calling `avb_slot_verify` and it returns OK.
+        let out_data = unsafe { out_data.assume_init() };
+        out_data.try_into()
+    }
+}
+
+fn verify_vbmeta_has_no_initrd_descriptor(
+    vbmeta_image: &AvbVBMetaData,
+) -> Result<(), AvbSlotVerifyError> {
+    is_not_null(vbmeta_image.vbmeta_data).map_err(|_| AvbSlotVerifyError::Io)?;
+    let mut has_unexpected_descriptor = false;
+    // SAFETY: It is safe as the raw pointer `vbmeta_image.vbmeta_data` is a nonnull pointer.
+    if !unsafe {
+        avb_descriptor_foreach(
+            vbmeta_image.vbmeta_data,
+            vbmeta_image.vbmeta_size,
+            Some(search_initrd_hash_descriptor),
+            &mut has_unexpected_descriptor as *mut _ as *mut c_void,
+        )
+    } {
+        return Err(AvbSlotVerifyError::InvalidMetadata);
+    }
+    if has_unexpected_descriptor {
+        Err(AvbSlotVerifyError::InvalidMetadata)
+    } else {
+        Ok(())
     }
 }
 
@@ -382,6 +562,16 @@
     trusted_public_key: &[u8],
 ) -> Result<(), AvbSlotVerifyError> {
     let mut payload = Payload { kernel, initrd, trusted_public_key };
-    let requested_partitions = [PartitionName::Kernel.as_cstr()];
-    payload.verify_partitions(&requested_partitions)
+    let kernel_verify_result = payload.verify_partitions(&[PartitionName::Kernel.as_cstr()])?;
+    let vbmeta_images = kernel_verify_result.vbmeta_images()?;
+    if vbmeta_images.len() != 1 {
+        // There can only be one VBMeta, from the 'boot' partition.
+        return Err(AvbSlotVerifyError::InvalidMetadata);
+    }
+    if payload.initrd.is_none() {
+        verify_vbmeta_has_no_initrd_descriptor(&vbmeta_images[0])?;
+    }
+    // TODO(b/256148034): Check the vbmeta doesn't have hash descriptors other than
+    // boot, initrd_normal, initrd_debug.
+    Ok(())
 }
diff --git a/pvmfw/avb/tests/api_test.rs b/pvmfw/avb/tests/api_test.rs
index 945c838..fc6bb1c 100644
--- a/pvmfw/avb/tests/api_test.rs
+++ b/pvmfw/avb/tests/api_test.rs
@@ -29,6 +29,8 @@
 const INITRD_NORMAL_IMG_PATH: &str = "microdroid_initrd_normal.img";
 const INITRD_DEBUG_IMG_PATH: &str = "microdroid_initrd_debuggable.img";
 const TEST_IMG_WITH_ONE_HASHDESC_PATH: &str = "test_image_with_one_hashdesc.img";
+const TEST_IMG_WITH_PROP_DESC_PATH: &str = "test_image_with_prop_desc.img";
+const TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH: &str = "test_image_with_non_initrd_hashdesc.img";
 const UNSIGNED_TEST_IMG_PATH: &str = "unsigned_test.img";
 
 const PUBLIC_KEY_RSA2048_PATH: &str = "data/testkey_rsa2048_pub.bin";
@@ -57,15 +59,39 @@
 
 #[test]
 fn payload_expecting_no_initrd_passes_verification_with_no_initrd() -> Result<()> {
-    let kernel = fs::read(TEST_IMG_WITH_ONE_HASHDESC_PATH)?;
-    let public_key = load_trusted_public_key()?;
-
-    assert_eq!(Ok(()), verify_payload(&kernel, None, &public_key));
-    Ok(())
+    assert_payload_verification_with_no_initrd_eq(
+        &fs::read(TEST_IMG_WITH_ONE_HASHDESC_PATH)?,
+        &load_trusted_public_key()?,
+        Ok(()),
+    )
 }
 
-// TODO(b/256148034): Test that kernel with two hashdesc and no initrd fails verification.
-// e.g. payload_expecting_initrd_fails_verification_with_no_initrd
+#[test]
+fn payload_with_non_initrd_descriptor_passes_verification_with_no_initrd() -> Result<()> {
+    assert_payload_verification_with_no_initrd_eq(
+        &fs::read(TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH)?,
+        &load_trusted_public_key()?,
+        Ok(()),
+    )
+}
+
+#[test]
+fn payload_with_prop_descriptor_fails_verification_with_no_initrd() -> Result<()> {
+    assert_payload_verification_with_no_initrd_eq(
+        &fs::read(TEST_IMG_WITH_PROP_DESC_PATH)?,
+        &load_trusted_public_key()?,
+        Err(AvbSlotVerifyError::InvalidMetadata),
+    )
+}
+
+#[test]
+fn payload_expecting_initrd_fails_verification_with_no_initrd() -> Result<()> {
+    assert_payload_verification_with_no_initrd_eq(
+        &load_latest_signed_kernel()?,
+        &load_trusted_public_key()?,
+        Err(AvbSlotVerifyError::InvalidMetadata),
+    )
+}
 
 #[test]
 fn payload_with_empty_public_key_fails_verification() -> Result<()> {
@@ -208,6 +234,15 @@
     Ok(vbmeta_header)
 }
 
+fn assert_payload_verification_with_no_initrd_eq(
+    kernel: &[u8],
+    trusted_public_key: &[u8],
+    expected_result: Result<(), AvbSlotVerifyError>,
+) -> Result<()> {
+    assert_eq!(expected_result, verify_payload(kernel, /*initrd=*/ None, trusted_public_key));
+    Ok(())
+}
+
 fn assert_payload_verification_fails(
     kernel: &[u8],
     initrd: &[u8],