libavb_rs: add descriptor extraction

Adds support for extractubg the descriptors contained in a vbmeta
image. The descriptor data isn't usable yet, will come in a follow-up
CL.

Bug: b/290110273
Test: atest
Change-Id: I5e5dada9412b0784374da83d54e60b2a592c1ff6

4 files changed, 268 insertions, 0 deletions

diff --git a/rust/src/descriptor/mod.rs b/rust/src/descriptor/mod.rs
new file mode 100644
index 0000000..279f425
--- /dev/null
+++ b/rust/src/descriptor/mod.rs
@@ -0,0 +1,230 @@
+// Copyright 2023, 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.
+
+//! Descriptor extraction and handling.
+//!
+//! Descriptors are information encoded into vbmeta images which can be
+//! extracted from the resulting data after performing verification.
+
+extern crate alloc;
+
+use crate::VbmetaData;
+use alloc::vec::Vec;
+use avb_bindgen::{
+    avb_descriptor_foreach, avb_descriptor_validate_and_byteswap, AvbDescriptor, AvbDescriptorTag,
+};
+use core::{ffi::c_void, mem::size_of, slice};
+
+/// A single descriptor.
+// TODO(b/290110273): add support for full descriptor contents.
+#[derive(Debug)]
+pub enum Descriptor<'a> {
+    /// Wraps `AvbPropertyDescriptor`.
+    Property(&'a [u8]),
+    /// Wraps `AvbHashtreeDescriptor`.
+    Hashtree(&'a [u8]),
+    /// Wraps `AvbHashDescriptor`.
+    Hash(&'a [u8]),
+    /// Wraps `AvbKernelCmdlineDescriptor`.
+    KernelCommandline(&'a [u8]),
+    /// Wraps `AvbChainPartitionDescriptor`.
+    ChainPartition(&'a [u8]),
+    /// Unknown descriptor type.
+    Unknown(&'a [u8]),
+}
+
+/// Possible errors when extracting descriptors.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum DescriptorError {
+    /// libavb rejected the descriptor header.
+    InvalidHeader,
+    /// A value in the descriptor was invalid.
+    InvalidValue,
+}
+
+/// `Result` type for `DescriptorError` errors.
+pub type DescriptorResult<T> = Result<T, DescriptorError>;
+
+impl<'a> Descriptor<'a> {
+    /// Extracts the fully-typed descriptor from the generic `AvbDescriptor` header.
+    ///
+    /// # Arguments
+    /// * `raw_descriptor`: the raw `AvbDescriptor` pointing into the vbmeta image.
+    ///
+    /// # Returns
+    /// The fully-typed `Descriptor`, or `DescriptorError` if parsing the descriptor failed.
+    ///
+    /// # Safety
+    /// `raw_descriptor` must point to a valid `AvbDescriptor`, including the `num_bytes_following`
+    /// data contents, that lives at least as long as `'a`.
+    unsafe fn new(raw_descriptor: *const AvbDescriptor) -> DescriptorResult<Self> {
+        // Transform header to host-endian.
+        let mut descriptor = AvbDescriptor {
+            tag: 0,
+            num_bytes_following: 0,
+        };
+        // SAFETY: both args point to valid `AvbDescriptor` objects.
+        if !unsafe { avb_descriptor_validate_and_byteswap(raw_descriptor, &mut descriptor) } {
+            return Err(DescriptorError::InvalidHeader);
+        }
+
+        // Extract the descriptor header and contents bytes. The descriptor sub-type headers
+        // include the top-level header as the first member, so we need to grab the entire
+        // descriptor including the top-level header.
+        let num_bytes_following = descriptor
+            .num_bytes_following
+            .try_into()
+            .map_err(|_| DescriptorError::InvalidValue)?;
+        let total_size = size_of::<AvbDescriptor>()
+            .checked_add(num_bytes_following)
+            .ok_or(DescriptorError::InvalidValue)?;
+
+        // SAFETY: `raw_descriptor` points to the header plus `num_bytes_following` bytes.
+        let contents = unsafe { slice::from_raw_parts(raw_descriptor as *const u8, total_size) };
+
+        match descriptor.tag.try_into() {
+            Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_PROPERTY) => Ok(Descriptor::Property(contents)),
+            Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_HASHTREE) => Ok(Descriptor::Hashtree(contents)),
+            Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_HASH) => Ok(Descriptor::Hash(contents)),
+            Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_KERNEL_CMDLINE) => {
+                Ok(Descriptor::KernelCommandline(contents))
+            }
+            Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_CHAIN_PARTITION) => {
+                Ok(Descriptor::ChainPartition(contents))
+            }
+            _ => Ok(Descriptor::Unknown(contents)),
+        }
+    }
+}
+
+/// Returns a vector of descriptors extracted from the given vbmeta image.
+///
+/// # Arguments
+/// * `vbmeta`: the `VbmetaData` object to extract descriptors from.
+///
+/// # Returns
+/// The descriptors, or `DescriptorError` if any error occurred.
+///
+/// # Safety
+/// `vbmeta` must have been validated by `slot_verify()`.
+pub(crate) unsafe fn get_descriptors(vbmeta: &VbmetaData) -> DescriptorResult<Vec<Descriptor>> {
+    let mut result = Ok(Vec::new());
+
+    // Use `avb_descriptor_foreach()` to grab all the descriptor pointers in `vmbeta.data()`.
+    // This implementation processes all the descriptors immediately, so that any error is
+    // detected here and working with descriptors can be error-free.
+    //
+    // SAFETY:
+    // * the caller ensures that `vbmeta` has been validated by `slot_verify()`, which satisfies
+    //   the libavb `avb_vbmeta_image_verify()` requirement.
+    // * `avb_descriptor_foreach()` ensures the validity of each descriptor pointer passed to
+    //   the `fill_descriptors_vec()` callback.
+    // * our lifetimes guarantee that the raw descriptor data in `vbmeta` will remain unchanged for
+    //   the lifetime of the returned `Descriptor` objects.
+    // * the `user_data` param is a valid `DescriptorResult<Vec<Descriptor>>` with no other
+    //   concurrent access.
+    unsafe {
+        // We can ignore the return value of this function since we use the passed-in `result`
+        // to convey success/failure as well as more detailed error info.
+        avb_descriptor_foreach(
+            vbmeta.data().as_ptr(),
+            vbmeta.data().len(),
+            Some(fill_descriptors_vec),
+            &mut result as *mut _ as *mut c_void,
+        );
+    }
+
+    result
+}
+
+/// Adds the given descriptor to the `Vec` pointed to by `user_data`.
+///
+/// Serves as a C callback for use with `avb_descriptor_foreach()`.
+///
+/// # Returns
+/// True on success, false on failure (which will stop iteration early).
+///
+/// # Safety
+/// * `descriptor` must point to a valid `AvbDescriptor`, including the `num_bytes_following`
+///   data contents, which remains valid and unmodified for the lifetime of the `Descriptor` objects
+///   in `user_data`.
+/// * `user_data` must point to a valid `DescriptorResult<Vec<Descriptor>>` with no other concurrent
+///   access.
+unsafe extern "C" fn fill_descriptors_vec(
+    descriptor: *const AvbDescriptor,
+    user_data: *mut c_void,
+) -> bool {
+    // SAFETY: `user_data` gives exclusive access to a valid `DescriptorResult<Vec<Descriptor>>`.
+    let result = unsafe { (user_data as *mut DescriptorResult<Vec<Descriptor>>).as_mut() };
+    // We can always unwrap here because we never pass a NULL pointer as `user_data`.
+    let result = result.unwrap();
+
+    // SAFETY: caller ensures that `descriptor` points to a valid `AvbDescriptor` with header and
+    // body contents, which remains unmodified at least as long as the new `Descriptor`.
+    match unsafe { Descriptor::new(descriptor) } {
+        Ok(d) => {
+            // We can always unwrap here because this function will never be called with an error
+            // in `result`, since we stop iteration as soon as we encounter an error.
+            result.as_mut().unwrap().push(d);
+            true
+        }
+        Err(e) => {
+            // Set the error and stop iteration early.
+            *result = Err(e);
+            false
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn new_unknown_descriptor() {
+        // A fake descriptor which is valid but with an unknown tag.
+        let data: &[u8] = &[
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x42, // tag = 0x42u64 (BE)
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x08, // num_bytes_following = 8u64 (BE)
+            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, // fake contents
+        ];
+
+        // SAFETY: we've crafted a valid descriptor in `data`.
+        let descriptor = unsafe { Descriptor::new(data.as_ptr() as *const _) }.unwrap();
+
+        let contents = match descriptor {
+            Descriptor::Unknown(c) => c,
+            d => panic!("Expected Unknown descriptor, got {d:?}"),
+        };
+        assert_eq!(data, contents);
+    }
+
+    #[test]
+    fn new_invalid_descriptor_length_fails() {
+        // `avb_descriptor_validate_and_byteswap()` should detect and reject descriptors whose
+        // `num_bytes_following` is not 8-byte aligned.
+        let data: &[u8] = &[
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x42, // tag = 0x42u64 (BE)
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, // num_bytes_following = 7u64 (BE)
+            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, // fake contents
+        ];
+
+        assert_eq!(
+            // SAFETY: we've created an invalid descriptor in a way that should be detected and
+            // fail safely without triggering any undefined behavior.
+            unsafe { Descriptor::new(data.as_ptr() as *const _) }.unwrap_err(),
+            DescriptorError::InvalidHeader
+        );
+    }
+}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 2d52c38..18f0f26 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -26,10 +26,12 @@
 // panic_handler and eh_personality conditional on actually building a dylib.
 #![cfg_attr(not(any(test, android_dylib)), no_std)]
 
+mod descriptor;
 mod error;
 mod ops;
 mod verify;
 
+pub use descriptor::{Descriptor, DescriptorError, DescriptorResult};
 pub use error::{
     IoError, IoResult, SlotVerifyError, SlotVerifyNoDataResult, SlotVerifyResult,
     VbmetaVerifyError, VbmetaVerifyResult,
diff --git a/rust/src/verify.rs b/rust/src/verify.rs
index 0c6e9ad..e972226 100644
--- a/rust/src/verify.rs
+++ b/rust/src/verify.rs
@@ -17,13 +17,17 @@
 //! This module is responsible for all the conversions required to pass information between
 //! libavb and Rust for verifying images.
 
+extern crate alloc;
+
 use crate::{
+    descriptor::{get_descriptors, Descriptor, DescriptorResult},
     error::{
         slot_verify_enum_to_result, vbmeta_verify_enum_to_result, SlotVerifyError,
         SlotVerifyNoDataResult, SlotVerifyResult, VbmetaVerifyResult,
     },
     ops, Ops,
 };
+use alloc::vec::Vec;
 use avb_bindgen::{
     avb_slot_verify, avb_slot_verify_data_free, AvbPartitionData, AvbSlotVerifyData, AvbVBMetaData,
 };
@@ -91,6 +95,18 @@ impl VbmetaData {
     pub fn verify_result(&self) -> VbmetaVerifyResult<()> {
         vbmeta_verify_enum_to_result(self.0.verify_result)
     }
+
+    /// Extracts the descriptors from the vbmeta image.
+    ///
+    /// Note that this function allocates memory to hold the `Descriptor` objects.
+    ///
+    /// # Returns
+    /// A vector of descriptors, or `DescriptorError` on failure.
+    pub fn descriptors(&self) -> DescriptorResult<Vec<Descriptor>> {
+        // SAFETY: the only way to get a `VbmetaData` object is via the return value of
+        // `slot_verify()`, so we know we have been properly validated.
+        unsafe { get_descriptors(self) }
+    }
 }
 
 impl fmt::Display for VbmetaData {
diff --git a/rust/tests/verify_tests.rs b/rust/tests/verify_tests.rs
index 353edf3..fff7796 100644
--- a/rust/tests/verify_tests.rs
+++ b/rust/tests/verify_tests.rs
@@ -602,3 +602,23 @@ fn corrupted_image_verification_data_display() {
         r#"slot: "", vbmeta: ["vbmeta": Ok(())], images: ["test_part": Err(Verification(None))]"#
     );
 }
+
+#[test]
+fn one_image_gives_single_descriptor() {
+    let mut ops = test_ops_one_image_one_vbmeta();
+
+    let result = verify_one_image_one_vbmeta(&mut ops);
+
+    let data = result.unwrap();
+    assert_eq!(data.vbmeta_data()[0].descriptors().unwrap().len(), 1);
+}
+
+#[test]
+fn two_images_gives_two_descriptors() {
+    let mut ops = test_ops_two_images_one_vbmeta();
+
+    let result = verify_two_images_one_vbmeta(&mut ops);
+
+    let data = result.unwrap();
+    assert_eq!(data.vbmeta_data()[0].descriptors().unwrap().len(), 2);
+}