commit 951de3ab4aefd53723fa3280b343c03f3803b5c7
author Alice Wang <aliceywang@google.com> Wed Jan 11 14:13:29 2023 +0000
committer Alice Wang <aliceywang@google.com> Wed Jan 11 16:06:38 2023 +0000
tree f9aba546cf0fa6a2bd31e53e08d297afc377dbce
parent f0fe7051d6101b96839ea50cc8cc31830d9f4858

[avb] Add PartitionName to facilitate partition name conversion

Bug: 256148034
Test: m pvmfw_img && atest libpvmfw_avb.test
Change-Id: I5e5be160d7f1eb820128fd93fa98936e8717d0ec

pvmfw/avb/src/verify.rs

diff --git a/pvmfw/avb/src/verify.rs b/pvmfw/avb/src/verify.rs
index 8eca130..0bab960 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -354,6 +354,41 @@
     }
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+enum PartitionName {
+    Kernel,
+    InitrdNormal,
+    InitrdDebug,
+}
+
+impl PartitionName {
+    const KERNEL_PARTITION_NAME: &[u8] = b"bootloader\0";
+    const INITRD_NORMAL_PARTITION_NAME: &[u8] = b"initrd_normal\0";
+    const INITRD_DEBUG_PARTITION_NAME: &[u8] = b"initrd_debug\0";
+
+    fn as_cstr(&self) -> &CStr {
+        let partition_name = 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()
+    }
+}
+
+impl TryFrom<&CStr> for PartitionName {
+    type Error = AvbIOError;
+
+    fn try_from(partition_name: &CStr) -> Result<Self, Self::Error> {
+        match partition_name.to_bytes_with_nul() {
+            Self::KERNEL_PARTITION_NAME => Ok(Self::Kernel),
+            Self::INITRD_NORMAL_PARTITION_NAME => Ok(Self::InitrdNormal),
+            Self::INITRD_DEBUG_PARTITION_NAME => Ok(Self::InitrdDebug),
+            _ => Err(AvbIOError::NoSuchPartition),
+        }
+    }
+}
+
 struct Payload<'a> {
     kernel: &'a [u8],
     initrd: Option<&'a [u8]>,
@@ -372,22 +407,17 @@
 }
 
 impl<'a> Payload<'a> {
-    const KERNEL_PARTITION_NAME: &[u8] = b"bootloader\0";
-    const INITRD_NORMAL_PARTITION_NAME: &[u8] = b"initrd_normal\0";
-    const INITRD_DEBUG_PARTITION_NAME: &[u8] = b"initrd_debug\0";
-
     const MAX_NUM_OF_HASH_DESCRIPTORS: usize = 3;
 
     fn get_partition(&self, partition_name: *const c_char) -> Result<&[u8], AvbIOError> {
         is_not_null(partition_name)?;
         // SAFETY: It is safe as the raw pointer `partition_name` is a nonnull pointer.
         let partition_name = unsafe { CStr::from_ptr(partition_name) };
-        match partition_name.to_bytes_with_nul() {
-            Self::KERNEL_PARTITION_NAME => Ok(self.kernel),
-            Self::INITRD_NORMAL_PARTITION_NAME | Self::INITRD_DEBUG_PARTITION_NAME => {
+        match partition_name.try_into()? {
+            PartitionName::Kernel => Ok(self.kernel),
+            PartitionName::InitrdNormal | PartitionName::InitrdDebug => {
                 self.initrd.ok_or(AvbIOError::NoSuchPartition)
             }
-            _ => Err(AvbIOError::NoSuchPartition),
         }
     }
 
@@ -445,8 +475,7 @@
     trusted_public_key: &[u8],
 ) -> Result<(), AvbImageVerifyError> {
     let mut payload = Payload { kernel, initrd, trusted_public_key };
-    let kernel = CStr::from_bytes_with_nul(Payload::KERNEL_PARTITION_NAME).unwrap();
-    let requested_partitions = [kernel];
+    let requested_partitions = [PartitionName::Kernel.as_cstr()];
     payload.verify_partitions(&requested_partitions)
 }