libavb_rs: add hash tree descriptor support

Adds the HashtreeDescriptor struct which provides a safe API to access
hash tree descriptors.

Bug: b/290110273
Test: atest libavb_rs_test libavb_rs_uuid_test libavb_rs_unittest libavb_rs_uuid_unittest
Change-Id: Ia1634feca54c566f6d5aa203ef3696b6318b413e

5 files changed, 266 insertions, 40 deletions

diff --git a/rust/Android.bp b/rust/Android.bp
index cafc675..c7c0128 100644
--- a/rust/Android.bp
+++ b/rust/Android.bp
@@ -32,12 +32,14 @@ rust_defaults {
     bindgen_flags: [
         "--constified-enum-module=AvbDescriptorTag",
         "--bitfield-enum=AvbHashDescriptorFlags",
+        "--bitfield-enum=AvbHashtreeDescriptorFlags",
         "--bitfield-enum=AvbSlotVerifyFlags",
         "--default-enum-style rust",
         "--with-derive-default",
         "--with-derive-custom=Avb.*Descriptor=FromZeroes,FromBytes",
         "--allowlist-type=AvbDescriptorTag",
         "--allowlist-type=AvbHashDescriptorFlags",
+        "--allowlist-type=AvbHashtreeDescriptorFlags",
         "--allowlist-function=.*",
         "--allowlist-var=AVB.*",
         "--use-core",
@@ -264,6 +266,7 @@ rust_defaults {
         ":avbrs_test_vbmeta",
         ":avbrs_test_vbmeta_2_parts",
         ":avbrs_test_vbmeta_persistent_digest",
+        ":avbrs_test_vbmeta_with_hashtree",
         ":avbrs_test_vbmeta_with_property",
     ],
     rustlibs: ["libhex"],
@@ -328,6 +331,18 @@ genrule {
     cmd: "$(location avbtool) add_hash_footer --image $(location :avbrs_test_image) --partition_name test_part_persistent_digest --dynamic_partition_size --do_not_append_vbmeta_image --use_persistent_digest --output_vbmeta_image $(out)",
 }
 
+// Unsigned vbmeta blob containing a hastree descriptor for partition name
+// "test_part_hashtree".
+genrule {
+    name: "avbrs_test_image_descriptor_hashtree",
+    tools: ["avbtool"],
+    srcs: [":avbrs_test_image"],
+    out: ["avbrs_test_image_descriptor_hashtree.img"],
+    // Generating FEC values requires the `fec` tool to be on $PATH, which does
+    // not seems to be possible here. For now pass `--do_not_generate_fec`.
+    cmd: "$(location avbtool) add_hashtree_footer --image $(location :avbrs_test_image) --partition_name test_part_hashtree --partition_size 0 --salt B000 --do_not_append_vbmeta_image --output_vbmeta_image $(out) --do_not_generate_fec",
+}
+
 // Standalone vbmeta image signing the test image descriptor.
 genrule {
     name: "avbrs_test_vbmeta",
@@ -377,6 +392,18 @@ genrule {
     cmd: "$(location avbtool) make_vbmeta_image --prop test_prop_key:test_prop_value --key $(location :avb_testkey_rsa4096) --algorithm SHA512_RSA4096 --include_descriptors_from_image $(location :avbrs_test_image_descriptor) --output $(out)",
 }
 
+// Standalone vbmeta image with the test image hashtree descriptor.
+genrule {
+    name: "avbrs_test_vbmeta_with_hashtree",
+    tools: ["avbtool"],
+    srcs: [
+        ":avbrs_test_image_descriptor_hashtree",
+        ":avb_testkey_rsa4096",
+    ],
+    out: ["test_vbmeta_with_hashtree.img"],
+    cmd: "$(location avbtool) make_vbmeta_image --key $(location :avb_testkey_rsa4096) --algorithm SHA512_RSA4096 --include_descriptors_from_image $(location :avbrs_test_image_descriptor_hashtree) --output $(out)",
+}
+
 // Combined test image + signed vbmeta footer for "test_part".
 avb_add_hash_footer {
     name: "avbrs_test_image_with_vbmeta_footer",
diff --git a/rust/src/descriptor/hashtree.rs b/rust/src/descriptor/hashtree.rs
new file mode 100644
index 0000000..5b44eaf
--- /dev/null
+++ b/rust/src/descriptor/hashtree.rs
@@ -0,0 +1,173 @@
+// Copyright 2024, 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.
+
+//! Hashtree descriptors.
+
+use super::{
+    util::{parse_descriptor, split_slice, ValidateAndByteswap, ValidationFunc},
+    DescriptorError, DescriptorResult,
+};
+use avb_bindgen::{avb_hashtree_descriptor_validate_and_byteswap, AvbHashtreeDescriptor};
+use core::{ffi::CStr, str::from_utf8};
+
+/// `AvbHashtreeDescriptorFlags`; see libavb docs for details.
+pub use avb_bindgen::AvbHashtreeDescriptorFlags as HashtreeDescriptorFlags;
+
+/// Wraps a Hashtree descriptor stored in a vbmeta image.
+#[derive(Debug, PartialEq, Eq)]
+pub struct HashtreeDescriptor<'a> {
+    /// DM-Verity version.
+    pub dm_verity_version: u32,
+
+    /// Hashed image size.
+    pub image_size: u64,
+
+    /// Offset to the root block of the hash tree.
+    pub tree_offset: u64,
+
+    /// Hash tree size.
+    pub tree_size: u64,
+
+    /// Data block size in bytes.
+    pub data_block_size: u32,
+
+    /// Hash block size in bytes.
+    pub hash_block_size: u32,
+
+    /// Number of forward error correction roots.
+    pub fec_num_roots: u32,
+
+    /// Offset to the forward error correction data.
+    pub fec_offset: u64,
+
+    /// Forward error correction data size.
+    pub fec_size: u64,
+
+    /// Hash algorithm name.
+    pub hash_algorithm: &'a str,
+
+    /// Flags.
+    pub flags: HashtreeDescriptorFlags,
+
+    /// Partition name.
+    pub partition_name: &'a str,
+
+    /// Salt used to hash the image.
+    pub salt: &'a [u8],
+
+    /// Image root hash digest.
+    pub root_digest: &'a [u8],
+}
+
+// SAFETY: `VALIDATE_AND_BYTESWAP_FUNC` is the correct libavb validator for this descriptor type.
+unsafe impl ValidateAndByteswap for AvbHashtreeDescriptor {
+    const VALIDATE_AND_BYTESWAP_FUNC: ValidationFunc<Self> =
+        avb_hashtree_descriptor_validate_and_byteswap;
+}
+
+impl<'a> HashtreeDescriptor<'a> {
+    /// Extract a `HashtreeDescriptor` from the given descriptor contents.
+    ///
+    /// # Arguments
+    /// * `contents`: descriptor contents, including the header, in raw big-endian format.
+    ///
+    /// # Returns
+    /// The new descriptor, or `DescriptorError` if the given `contents` aren't a valid
+    /// `AvbHashtreeDescriptor`.
+    pub(super) fn new(contents: &'a [u8]) -> DescriptorResult<Self> {
+        // Descriptor contains: header + name + salt + digest.
+        let descriptor = parse_descriptor::<AvbHashtreeDescriptor>(contents)?;
+        let (partition_name, remainder) =
+            split_slice(descriptor.body, descriptor.header.partition_name_len)?;
+        let (salt, remainder) = split_slice(remainder, descriptor.header.salt_len)?;
+        let (root_digest, _) = split_slice(remainder, descriptor.header.root_digest_len)?;
+
+        // Extract the hash algorithm from the original raw header since the temporary
+        // byte-swapped header doesn't live past this function.
+        // The hash algorithm is a nul-terminated UTF-8 string which is identical in the raw
+        // and byteswapped headers.
+        let hash_algorithm = CStr::from_bytes_until_nul(&descriptor.raw_header.hash_algorithm)
+            .map_err(|_| DescriptorError::InvalidValue)?
+            .to_str()
+            .map_err(|_| DescriptorError::InvalidUtf8)?;
+
+        Ok(Self {
+            dm_verity_version: descriptor.header.dm_verity_version,
+            image_size: descriptor.header.image_size,
+            tree_offset: descriptor.header.tree_offset,
+            tree_size: descriptor.header.tree_size,
+            data_block_size: descriptor.header.data_block_size,
+            hash_block_size: descriptor.header.hash_block_size,
+            fec_num_roots: descriptor.header.fec_num_roots,
+            fec_offset: descriptor.header.fec_offset,
+            fec_size: descriptor.header.fec_size,
+            hash_algorithm,
+            partition_name: from_utf8(partition_name).map_err(|_| DescriptorError::InvalidUtf8)?,
+            salt,
+            root_digest,
+            flags: HashtreeDescriptorFlags(descriptor.header.flags),
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::mem::size_of;
+
+    /// A valid hashtree descriptor in raw big-endian format.
+    const TEST_HASHTREE_DESCRIPTOR: &[u8] = &[
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0xE0, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00,
+        0x00, 0x10, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x73, 0x68, 0x61,
+        0x31, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x12, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x74, 0x65, 0x73, 0x74, 0x5F, 0x70, 0x61, 0x72, 0x74, 0x5F, 0x68, 0x61, 0x73, 0x68, 0x74,
+        0x72, 0x65, 0x65, 0x99, 0xCE, 0xC4, 0x29, 0x60, 0x61, 0xCF, 0xBD, 0xE7, 0xD2, 0x17, 0xE2,
+        0x88, 0x99, 0x05, 0x39, 0xAB, 0x70, 0x6D, 0xD0, 0x4C, 0x77, 0x76, 0xF8, 0xFD, 0xD2, 0x2B,
+        0xF4, 0xC4, 0x7F, 0x31, 0x1B, 0x7B, 0x7B, 0xA5, 0xEF, 0x42, 0x8D, 0x7B, 0xE8, 0x00, 0x00,
+    ];
+
+    #[test]
+    fn new_hashtree_descriptor_success() {
+        let descriptor = HashtreeDescriptor::new(TEST_HASHTREE_DESCRIPTOR);
+        assert!(descriptor.is_ok());
+    }
+
+    #[test]
+    fn new_hashtree_descriptor_too_short_header_fails() {
+        let bad_header_size = size_of::<AvbHashtreeDescriptor>() - 1;
+        assert_eq!(
+            HashtreeDescriptor::new(&TEST_HASHTREE_DESCRIPTOR[..bad_header_size]).unwrap_err(),
+            DescriptorError::InvalidHeader
+        );
+    }
+
+    #[test]
+    fn new_hashtree_descriptor_too_short_contents_fails() {
+        // The last 2 bytes are padding, so we need to drop 3 bytes to trigger an error.
+        let bad_contents_size = TEST_HASHTREE_DESCRIPTOR.len() - 3;
+        assert_eq!(
+            HashtreeDescriptor::new(&TEST_HASHTREE_DESCRIPTOR[..bad_contents_size]).unwrap_err(),
+            DescriptorError::InvalidSize
+        );
+    }
+}
diff --git a/rust/src/descriptor/mod.rs b/rust/src/descriptor/mod.rs
index ba741dd..82b434e 100644
--- a/rust/src/descriptor/mod.rs
+++ b/rust/src/descriptor/mod.rs
@@ -20,6 +20,7 @@
 extern crate alloc;
 
 mod hash;
+mod hashtree;
 mod property;
 mod util;
 
@@ -31,16 +32,17 @@ use avb_bindgen::{
 use core::{ffi::c_void, mem::size_of, slice};
 
 pub use hash::{HashDescriptor, HashDescriptorFlags};
+pub use hashtree::{HashtreeDescriptor, HashtreeDescriptorFlags};
 pub use property::PropertyDescriptor;
 
 /// A single descriptor.
 // TODO(b/290110273): add support for full descriptor contents.
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum Descriptor<'a> {
     /// Wraps `AvbPropertyDescriptor`.
     Property(PropertyDescriptor<'a>),
     /// Wraps `AvbHashtreeDescriptor`.
-    Hashtree(&'a [u8]),
+    Hashtree(HashtreeDescriptor<'a>),
     /// Wraps `AvbHashDescriptor`.
     Hash(HashDescriptor<'a>),
     /// Wraps `AvbKernelCmdlineDescriptor`.
@@ -110,7 +112,9 @@ impl<'a> Descriptor<'a> {
             Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_PROPERTY) => {
                 Ok(Descriptor::Property(PropertyDescriptor::new(contents)?))
             }
-            Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_HASHTREE) => Ok(Descriptor::Hashtree(contents)),
+            Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_HASHTREE) => {
+                Ok(Descriptor::Hashtree(HashtreeDescriptor::new(contents)?))
+            }
             Ok(AvbDescriptorTag::AVB_DESCRIPTOR_TAG_HASH) => {
                 Ok(Descriptor::Hash(HashDescriptor::new(contents)?))
             }
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index a5419f0..312f71e 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -33,7 +33,7 @@ mod verify;
 
 pub use descriptor::{
     Descriptor, DescriptorError, DescriptorResult, HashDescriptor, HashDescriptorFlags,
-    PropertyDescriptor,
+    HashtreeDescriptor, HashtreeDescriptorFlags, PropertyDescriptor,
 };
 pub use error::{
     IoError, IoResult, SlotVerifyError, SlotVerifyNoDataResult, SlotVerifyResult,
diff --git a/rust/tests/verify_tests.rs b/rust/tests/verify_tests.rs
index 61b0576..7c5bb0a 100644
--- a/rust/tests/verify_tests.rs
+++ b/rust/tests/verify_tests.rs
@@ -16,8 +16,9 @@
 
 use crate::test_ops::TestOps;
 use avb::{
-    slot_verify, Descriptor, HashDescriptor, HashDescriptorFlags, HashtreeErrorMode, IoError,
-    PropertyDescriptor, SlotVerifyData, SlotVerifyError, SlotVerifyFlags, SlotVerifyResult,
+    slot_verify, Descriptor, HashDescriptor, HashDescriptorFlags, HashtreeDescriptor,
+    HashtreeDescriptorFlags, HashtreeErrorMode, IoError, PropertyDescriptor, SlotVerifyData,
+    SlotVerifyError, SlotVerifyFlags, SlotVerifyResult,
 };
 use hex::decode;
 use std::{ffi::CString, fs};
@@ -28,14 +29,12 @@ use uuid::uuid;
 const TEST_IMAGE_PATH: &str = "test_image.img";
 const TEST_IMAGE_SIZE: usize = 16 * 1024;
 const TEST_IMAGE_SALT_HEX: &str = "1000";
-// Expected digest determined by examining the vbmeta image with `avbtool info_image`.
-const TEST_IMAGE_DIGEST_HEX: &str =
-    "89e6fd3142917b8c34ac7d30897a907a71bd3bf5d9b39d00bf938b41dcf3b84f";
-const TEST_IMAGE_HASH_ALGO: &str = "sha256"; // Default value, we don't explicitly set this.
+const TEST_HASHTREE_SALT_HEX: &str = "B000";
 const TEST_VBMETA_PATH: &str = "test_vbmeta.img";
 const TEST_VBMETA_2_PARTITIONS_PATH: &str = "test_vbmeta_2_parts.img";
 const TEST_VBMETA_PERSISTENT_DIGEST_PATH: &str = "test_vbmeta_persistent_digest.img";
 const TEST_VBMETA_WITH_PROPERTY_PATH: &str = "test_vbmeta_with_property.img";
+const TEST_VBMETA_WITH_HASHTREE_PATH: &str = "test_vbmeta_with_hashtree.img";
 const TEST_IMAGE_WITH_VBMETA_FOOTER_PATH: &str = "avbrs_test_image_with_vbmeta_footer.img";
 const TEST_IMAGE_WITH_VBMETA_FOOTER_FOR_BOOT_PATH: &str =
     "avbrs_test_image_with_vbmeta_footer_for_boot.img";
@@ -44,10 +43,19 @@ const TEST_PARTITION_NAME: &str = "test_part";
 const TEST_PARTITION_SLOT_C_NAME: &str = "test_part_c";
 const TEST_PARTITION_2_NAME: &str = "test_part_2";
 const TEST_PARTITION_PERSISTENT_DIGEST_NAME: &str = "test_part_persistent_digest";
+const TEST_PARTITION_HASH_TREE_NAME: &str = "test_part_hashtree";
 const TEST_VBMETA_ROLLBACK_LOCATION: usize = 0; // Default value, we don't explicitly set this.
 const TEST_PROPERTY_KEY: &str = "test_prop_key";
 const TEST_PROPERTY_VALUE: &[u8] = b"test_prop_value";
 
+// Expected values determined by examining the vbmeta image with `avbtool info_image`.
+// Images can be found in <out>/soong/.intermediates/external/avb/rust/.
+const TEST_IMAGE_DIGEST_HEX: &str =
+    "89e6fd3142917b8c34ac7d30897a907a71bd3bf5d9b39d00bf938b41dcf3b84f";
+const TEST_IMAGE_HASH_ALGO: &str = "sha256";
+const TEST_HASHTREE_DIGEST_HEX: &str = "5373fc4ee3dd898325eeeffb5a1dbb041900c5f1";
+const TEST_HASHTREE_ALGORITHM: &str = "sha1";
+
 /// Initializes a `TestOps` object such that verification will succeed on `TEST_PARTITION_NAME`.
 fn test_ops_one_image_one_vbmeta() -> TestOps {
     let mut ops = TestOps::default();
@@ -640,22 +648,16 @@ fn verify_hash_descriptor() {
 
     let data = result.unwrap();
     let descriptors = &data.vbmeta_data()[0].descriptors().unwrap();
-    let descriptor = match &descriptors[0] {
-        Descriptor::Hash(d) => d,
-        d => panic!("Expected hash descriptor, got {:?}", d),
-    };
 
-    assert_eq!(
-        descriptor,
-        &HashDescriptor {
-            image_size: TEST_IMAGE_SIZE as u64,
-            hash_algorithm: TEST_IMAGE_HASH_ALGO,
-            flags: HashDescriptorFlags(0),
-            partition_name: TEST_PARTITION_NAME,
-            salt: &decode(TEST_IMAGE_SALT_HEX).unwrap(),
-            digest: &decode(TEST_IMAGE_DIGEST_HEX).unwrap()
-        }
-    )
+    let expected = HashDescriptor {
+        image_size: TEST_IMAGE_SIZE as u64,
+        hash_algorithm: TEST_IMAGE_HASH_ALGO,
+        flags: HashDescriptorFlags(0),
+        partition_name: TEST_PARTITION_NAME,
+        salt: &decode(TEST_IMAGE_SALT_HEX).unwrap(),
+        digest: &decode(TEST_IMAGE_DIGEST_HEX).unwrap(),
+    };
+    assert!(descriptors.contains(&Descriptor::Hash(expected)));
 }
 
 #[test]
@@ -668,20 +670,40 @@ fn verify_property_descriptor() {
 
     let data = result.unwrap();
     let descriptors = &data.vbmeta_data()[0].descriptors().unwrap();
-    let descriptor = descriptors
-        .iter()
-        .filter_map(|d| match d {
-            Descriptor::Property(p) => Some(p),
-            _ => None,
-        })
-        .next()
-        .unwrap();
 
-    assert_eq!(
-        descriptor,
-        &PropertyDescriptor {
-            key: TEST_PROPERTY_KEY,
-            value: TEST_PROPERTY_VALUE
-        }
-    )
+    let expected = PropertyDescriptor {
+        key: TEST_PROPERTY_KEY,
+        value: TEST_PROPERTY_VALUE,
+    };
+    assert!(descriptors.contains(&Descriptor::Property(expected)));
+}
+
+#[test]
+fn verify_hashtree_descriptor() {
+    let mut ops = test_ops_one_image_one_vbmeta();
+    // Replace vbmeta with the version containing a hashtree descriptor.
+    ops.add_partition("vbmeta", fs::read(TEST_VBMETA_WITH_HASHTREE_PATH).unwrap());
+
+    let result = verify_one_image_one_vbmeta(&mut ops);
+
+    let data = result.unwrap();
+    let descriptors = &data.vbmeta_data()[0].descriptors().unwrap();
+
+    let expected = HashtreeDescriptor {
+        dm_verity_version: 1,
+        image_size: TEST_IMAGE_SIZE as u64,
+        tree_offset: TEST_IMAGE_SIZE as u64,
+        tree_size: 4096,
+        data_block_size: 4096,
+        hash_block_size: 4096,
+        fec_num_roots: 0,
+        fec_offset: 0,
+        fec_size: 0,
+        hash_algorithm: TEST_HASHTREE_ALGORITHM,
+        flags: HashtreeDescriptorFlags(0),
+        partition_name: TEST_PARTITION_HASH_TREE_NAME,
+        salt: &decode(TEST_HASHTREE_SALT_HEX).unwrap(),
+        root_digest: &decode(TEST_HASHTREE_DIGEST_HEX).unwrap(),
+    };
+    assert!(descriptors.contains(&Descriptor::Hashtree(expected)));
 }