Move to zerocopy: partition info

Signed-off-by: Balint Dobszay <balint.dobszay@arm.com>
Change-Id: Iee2a126a02c311dfefcbadd4866580ad92891183
diff --git a/src/ffa_v1_1.rs b/src/ffa_v1_1.rs
index e16ec06..755003f 100644
--- a/src/ffa_v1_1.rs
+++ b/src/ffa_v1_1.rs
@@ -73,3 +73,63 @@
     /// Offset 24, length 8: Reserved (MBZ)
     pub(crate) reserved: u64,
 }
+
+/// Table 13.37: Partition information descriptor
+#[derive(Default, FromBytes, IntoBytes, KnownLayout, Immutable)]
+#[repr(C, packed)]
+pub(crate) struct partition_info_descriptor {
+    /// Offset 0, length 2: 16-bit ID of the partition, stream or auxiliary endpoint.
+    pub(crate) partition_id: u16,
+    /// Offset 2, length 2:
+    /// - Number of execution contexts implemented by this partition if Bit\[5:4\] = b’00 in the
+    ///   Partition properties field.
+    /// - ID of the proxy endpoint for a dependent peripheral device if Bit\[5:4\] = b’10 in the
+    ///   Partition properties field.
+    /// - Reserved and MBZ for all other encodings of the Partition properties field.
+    pub(crate) exec_ctx_count_or_proxy_id: u16,
+    /// Offset 4, length 4: Flags to determine partition properties.
+    /// - Bit\[3:0\] has the following encoding if Bit\[5:4\] = b’00. It is Reserved and MBZ otherwise.
+    ///   + Bit\[0\] has the following encoding:
+    ///     * b’0: Does not support receipt of direct requests
+    ///     * b’1: Supports receipt of direct requests. Count of execution contexts must be either 1
+    ///       or equal to the number of PEs in the system.
+    ///   + bit\[1\] has the following encoding:
+    ///     * b’0: Cannot send direct requests.
+    ///     * b’1: Can send direct requests.
+    ///   + bit\[2\] has the following encoding:
+    ///     * b’0: Cannot send and receive indirect messages.
+    ///     * b’1: Can send and receive indirect messages.
+    ///   + bit\[3\] has the following encoding:
+    ///     * b’0: Does not support receipt of notifications.
+    ///     * b’1: Supports receipt of notifications.
+    /// - bit\[5:4\] has the following encoding:
+    ///   + b’00: Partition ID is a PE endpoint ID.
+    ///   + b’01: Partition ID is a SEPID for an independent peripheral device.
+    ///   + b’10: Partition ID is a SEPID for an dependent peripheral device.
+    ///   + b’11: Partition ID is an auxiliary ID.
+    /// - bit\[6\] has the following encoding:
+    ///   + b’0: Partition must not be informed about each VM that is created by the Hypervisor.
+    ///   + b’1: Partition must be informed about each VM that is created by the Hypervisor.
+    ///   + bit\[6\] is used only if the following conditions are true. It is Reserved (MBZ) in all
+    ///     other scenarios.
+    ///     * This ABI is invoked at the Non-secure physical FF-A instance.
+    ///     * The partition is an SP that supports receipt of direct requests i.e. Bit\[0\] = b’1.
+    /// - bit\[7\] has the following encoding:
+    ///   + b’0: Partition must not be informed about each VM that is destroyed by the Hypervisor.
+    ///   + b’1: Partition must be informed about each VM that is destroyed by the Hypervisor.
+    ///   + bit\[7\] is used only if the following conditions are true. It is Reserved (MBZ) in all
+    ///     other scenarios.
+    ///     * This ABI is invoked at the Non-secure physical FF-A instance.
+    ///     * The partition is an SP that supports receipt of direct requests i.e. Bit\[0\] = b’1.
+    /// - bit\[8\] has the following encoding:
+    ///   + b’0: Partition runs in the AArch32 execution state.
+    ///   + b’1: Partition runs in the AArch64 execution state.
+    /// - bit\[31:9\]: Reserved (MBZ).
+    pub(crate) partition_props: u32,
+    /// Offset 8, length 16:
+    /// - UUID of the partition, stream or auxiliary endpoint if the Nil UUID was specified in w1-w4
+    ///   as an input parameter.
+    /// - This field is reserved and MBZ if a non-Nil UUID was was specified in w1-w4 as an input
+    ///   parameter.
+    pub(crate) uuid: [u8; 16],
+}
diff --git a/src/partition_info.rs b/src/partition_info.rs
index 5301cb7..a8025b3 100644
--- a/src/partition_info.rs
+++ b/src/partition_info.rs
@@ -1,121 +1,323 @@
 // SPDX-FileCopyrightText: Copyright 2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
+use crate::ffa_v1_1::partition_info_descriptor;
+use thiserror::Error;
 use uuid::Uuid;
+use zerocopy::{FromBytes, IntoBytes};
 
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("Invalid buffer size")]
+    InvalidBufferSize,
+    #[error("Malformed descriptor")]
+    MalformedDescriptor,
+}
+
+impl From<Error> for crate::FfaError {
+    fn from(_value: Error) -> Self {
+        Self::InvalidParameters
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum PartitionIdType {
-    PeEndpoint(u16),
+    PeEndpoint { execution_ctx_count: u16 },
     SepidIndep,
-    SepidDep(u16),
+    SepidDep { proxy_endpoint_id: u16 },
     Aux,
 }
 
-pub struct PartInfoDesc {
-    pub partition_id: u16,
-    pub uuid: Uuid,
-    pub id_type: PartitionIdType,
+impl PartitionIdType {
+    const SHIFT: usize = 4;
+    const MASK: u32 = 0b11;
+    const PE_ENDPOINT: u32 = 0b00;
+    const SEPID_INDEP: u32 = 0b01;
+    const SEPID_DEP: u32 = 0b10;
+    const AUX: u32 = 0b11;
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct PartitionProperties {
+    /// Supports receipt of direct requests
     pub support_direct_req_rec: bool,
+    /// Can send direct requests
     pub support_direct_req_send: bool,
+    /// Can send and receive indirect messages
     pub support_indirect_msg: bool,
+    /// Supports receipt of notifications
     pub support_notif_rec: bool,
+    /// Must be informed about each VM that is created by the Hypervisor
     pub subscribe_vm_created: bool,
+    /// Must be informed about each VM that is destroyed by the Hypervisor
     pub subscribe_vm_destroyed: bool,
+    /// Partition runs in the AArch64 execution state
     pub is_aarch64: bool,
 }
 
-impl PartInfoDesc {
-    pub const SIZE: usize = 24;
+impl PartitionProperties {
+    const SUPPORT_DIRECT_REQ_REC_SHIFT: usize = 0;
+    const SUPPORT_DIRECT_REQ_SEND_SHIFT: usize = 1;
+    const SUPPORT_INDIRECT_MSG_SHIFT: usize = 2;
+    const SUPPORT_NOTIF_REC_SHIFT: usize = 3;
+    const SUBSCRIBE_VM_CREATED_SHIFT: usize = 6;
+    const SUBSCRIBE_VM_DESTROYED_SHIFT: usize = 7;
+    const IS_AARCH64_SHIFT: usize = 8;
 }
 
-pub fn create_partition_info(buf: &mut [u8], descriptors: &[PartInfoDesc], fill_uuid: bool) {
-    let mut offset = 0;
+struct PartPropWrapper(PartitionIdType, PartitionProperties);
 
-    for desc in descriptors {
-        // Offset 0, length 2: 16-bit ID of the partition, stream or auxiliary endpoint.
-        buf[offset..offset + 2].copy_from_slice(&desc.partition_id.to_le_bytes());
+impl From<PartPropWrapper> for (u32, u16) {
+    fn from(value: PartPropWrapper) -> Self {
+        let exec_ctx_count_or_proxy_id = match value.0 {
+            PartitionIdType::PeEndpoint {
+                execution_ctx_count,
+            } => execution_ctx_count,
+            PartitionIdType::SepidIndep => 0,
+            PartitionIdType::SepidDep { proxy_endpoint_id } => proxy_endpoint_id,
+            PartitionIdType::Aux => 0,
+        };
 
-        // Offset 2, length 2: Execution context count or Proxy partition ID
-        match desc.id_type {
-            PartitionIdType::PeEndpoint(exec_ctx_cnt) => {
-                buf[offset + 2..offset + 4].copy_from_slice(&exec_ctx_cnt.to_le_bytes())
-            }
-            PartitionIdType::SepidDep(id) => {
-                buf[offset + 2..offset + 4].copy_from_slice(&id.to_le_bytes())
-            }
-            _ => buf[offset + 2..offset + 4].fill(0),
-        }
+        let mut props = match value.0 {
+            PartitionIdType::PeEndpoint { .. } => {
+                let mut p = PartitionIdType::PE_ENDPOINT << PartitionIdType::SHIFT;
 
-        // Offset 4, length 4: Flags to determine partition properties.
-        let mut props = 0u32;
-        match desc.id_type {
-            PartitionIdType::PeEndpoint(_) => {
-                if desc.support_direct_req_rec {
-                    props |= 0b1;
-                    if desc.subscribe_vm_created {
-                        // TODO: check NS phys instance
-                        props |= 0b1 << 6
+                if value.1.support_direct_req_rec {
+                    p |= 1 << PartitionProperties::SUPPORT_DIRECT_REQ_REC_SHIFT;
+                    if value.1.subscribe_vm_created {
+                        // TODO: how to handle if ABI is invoked at NS phys instance?
+                        p |= 1 << PartitionProperties::SUBSCRIBE_VM_CREATED_SHIFT
                     }
-                    if desc.subscribe_vm_destroyed {
-                        // TODO: check NS phys instance
-                        props |= 0b1 << 7
+                    if value.1.subscribe_vm_destroyed {
+                        // TODO: how to handle if ABI is invoked at NS phys instance?
+                        p |= 1 << PartitionProperties::SUBSCRIBE_VM_DESTROYED_SHIFT
                     }
                 }
-                if desc.support_direct_req_send {
-                    props |= 0b1 << 1
+                if value.1.support_direct_req_send {
+                    p |= 1 << PartitionProperties::SUPPORT_DIRECT_REQ_SEND_SHIFT
                 }
-                if desc.support_indirect_msg {
-                    props |= 0b1 << 2
+                if value.1.support_indirect_msg {
+                    p |= 1 << PartitionProperties::SUPPORT_INDIRECT_MSG_SHIFT
                 }
-                if desc.support_notif_rec {
-                    props |= 0b1 << 3
+                if value.1.support_notif_rec {
+                    p |= 1 << PartitionProperties::SUPPORT_NOTIF_REC_SHIFT
+                }
+
+                p
+            }
+            PartitionIdType::SepidIndep => PartitionIdType::SEPID_INDEP << PartitionIdType::SHIFT,
+            PartitionIdType::SepidDep { .. } => {
+                PartitionIdType::SEPID_DEP << PartitionIdType::SHIFT
+            }
+            PartitionIdType::Aux => PartitionIdType::AUX << PartitionIdType::SHIFT,
+        };
+
+        if value.1.is_aarch64 {
+            props |= 1 << PartitionProperties::IS_AARCH64_SHIFT
+        }
+
+        (props, exec_ctx_count_or_proxy_id)
+    }
+}
+
+impl From<(u32, u16)> for PartPropWrapper {
+    fn from(value: (u32, u16)) -> Self {
+        let part_id_type = match (value.0 >> PartitionIdType::SHIFT) & PartitionIdType::MASK {
+            PartitionIdType::PE_ENDPOINT => PartitionIdType::PeEndpoint {
+                execution_ctx_count: value.1,
+            },
+            PartitionIdType::SEPID_INDEP => PartitionIdType::SepidIndep,
+            PartitionIdType::SEPID_DEP => PartitionIdType::SepidDep {
+                proxy_endpoint_id: value.1,
+            },
+            PartitionIdType::AUX => PartitionIdType::Aux,
+            _ => panic!(), // The match is exhaustive for a 2-bit value
+        };
+
+        let mut part_props = PartitionProperties::default();
+
+        if (value.0 >> PartitionIdType::SHIFT) & PartitionIdType::MASK
+            == PartitionIdType::PE_ENDPOINT
+        {
+            if (value.0 >> PartitionProperties::SUPPORT_DIRECT_REQ_REC_SHIFT) & 0b1 == 1 {
+                part_props.support_direct_req_rec = true;
+
+                if (value.0 >> PartitionProperties::SUBSCRIBE_VM_CREATED_SHIFT) & 0b1 == 1 {
+                    part_props.subscribe_vm_created = true;
+                }
+
+                if (value.0 >> PartitionProperties::SUBSCRIBE_VM_DESTROYED_SHIFT) & 0b1 == 1 {
+                    part_props.subscribe_vm_destroyed = true;
                 }
             }
-            PartitionIdType::SepidIndep => props |= 0b01 << 4,
-            PartitionIdType::SepidDep(_) => props |= 0b10 << 4,
-            PartitionIdType::Aux => props |= 0b11 << 4,
-        }
-        if desc.is_aarch64 {
-            props |= 0b1 << 8
-        }
-        buf[offset + 4..offset + 8].copy_from_slice(&props.to_le_bytes());
 
-        // Offset 8, length 16: Partition UUID if the Nil UUID was specified. Reserved (MBZ) otherwise
-        if fill_uuid {
-            buf[offset + 8..offset + 24].copy_from_slice(desc.uuid.as_bytes());
-        } else {
-            buf[offset + 8..offset + 24].fill(0);
+            if (value.0 >> PartitionProperties::SUPPORT_DIRECT_REQ_SEND_SHIFT) & 0b1 == 1 {
+                part_props.support_direct_req_send = true;
+            }
+
+            if (value.0 >> PartitionProperties::SUPPORT_INDIRECT_MSG_SHIFT) & 0b1 == 1 {
+                part_props.support_indirect_msg = true;
+            }
+
+            if (value.0 >> PartitionProperties::SUPPORT_NOTIF_REC_SHIFT) & 0b1 == 1 {
+                part_props.support_notif_rec = true;
+            }
         }
 
-        offset += PartInfoDesc::SIZE;
+        if (value.0 >> PartitionProperties::IS_AARCH64_SHIFT) & 0b1 == 1 {
+            part_props.is_aarch64 = true;
+        }
+
+        PartPropWrapper(part_id_type, part_props)
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct PartitionInfo {
+    pub uuid: Uuid,
+    pub partition_id: u16,
+    pub partition_id_type: PartitionIdType,
+    pub props: PartitionProperties,
+}
+
+impl PartitionInfo {
+    pub const DESC_SIZE: usize = size_of::<partition_info_descriptor>();
+
+    pub fn pack(descriptors: &[PartitionInfo], buf: &mut [u8], fill_uuid: bool) {
+        let mut offset = 0;
+
+        for desc in descriptors {
+            let mut desc_raw = partition_info_descriptor {
+                partition_id: desc.partition_id,
+                ..Default::default()
+            };
+
+            (
+                desc_raw.partition_props,
+                desc_raw.exec_ctx_count_or_proxy_id,
+            ) = PartPropWrapper(desc.partition_id_type, desc.props).into();
+
+            if fill_uuid {
+                desc_raw.uuid.copy_from_slice(desc.uuid.as_bytes());
+            }
+
+            desc_raw.write_to_prefix(&mut buf[offset..]).unwrap();
+            offset += Self::DESC_SIZE;
+        }
+    }
+}
+
+pub struct PartitionInfoIterator<'a> {
+    buf: &'a [u8],
+    offset: usize,
+    count: usize,
+}
+
+impl<'a> PartitionInfoIterator<'a> {
+    pub fn new(buf: &'a [u8], count: usize) -> Result<Self, Error> {
+        let Some(total_size) = count.checked_mul(PartitionInfo::DESC_SIZE) else {
+            return Err(Error::InvalidBufferSize);
+        };
+
+        if buf.len() < total_size {
+            return Err(Error::InvalidBufferSize);
+        }
+
+        Ok(Self {
+            buf,
+            offset: 0,
+            count,
+        })
+    }
+}
+
+impl Iterator for PartitionInfoIterator<'_> {
+    type Item = Result<PartitionInfo, Error>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.count > 0 {
+            let offset = self.offset;
+            self.offset += PartitionInfo::DESC_SIZE;
+            self.count -= 1;
+
+            let Ok(desc_raw) = partition_info_descriptor::ref_from_bytes(
+                &self.buf[offset..offset + PartitionInfo::DESC_SIZE],
+            ) else {
+                return Some(Err(Error::MalformedDescriptor));
+            };
+
+            let partition_id = desc_raw.partition_id;
+
+            let wrapper = PartPropWrapper::from((
+                desc_raw.partition_props,
+                desc_raw.exec_ctx_count_or_proxy_id,
+            ));
+
+            let uuid = Uuid::from_bytes(desc_raw.uuid);
+
+            let desc = PartitionInfo {
+                uuid,
+                partition_id,
+                partition_id_type: wrapper.0,
+                props: wrapper.1,
+            };
+
+            return Some(Ok(desc));
+        }
+
+        None
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-
     use uuid::uuid;
 
+    // TODO: add tests with a known correct partition info blob
+
     #[test]
     fn part_info() {
-        let desc = PartInfoDesc {
-            partition_id: 0x8001,
+        let desc1 = PartitionInfo {
             uuid: uuid!("12345678-1234-1234-1234-123456789abc"),
-            id_type: PartitionIdType::PeEndpoint(1),
-            support_direct_req_rec: true,
-            support_direct_req_send: true,
-            support_indirect_msg: false,
-            support_notif_rec: false,
-            subscribe_vm_created: true,
-            subscribe_vm_destroyed: true,
-            is_aarch64: true,
+            partition_id: 0x8001,
+            partition_id_type: PartitionIdType::PeEndpoint {
+                execution_ctx_count: 1,
+            },
+            props: PartitionProperties {
+                support_direct_req_rec: true,
+                support_direct_req_send: true,
+                support_indirect_msg: false,
+                support_notif_rec: false,
+                subscribe_vm_created: true,
+                subscribe_vm_destroyed: true,
+                is_aarch64: true,
+            },
+        };
+
+        let desc2 = PartitionInfo {
+            uuid: uuid!("abcdef00-abcd-dcba-1234-abcdef012345"),
+            partition_id: 0x8002,
+            partition_id_type: PartitionIdType::SepidIndep,
+            props: PartitionProperties {
+                support_direct_req_rec: false,
+                support_direct_req_send: false,
+                support_indirect_msg: false,
+                support_notif_rec: false,
+                subscribe_vm_created: false,
+                subscribe_vm_destroyed: false,
+                is_aarch64: true,
+            },
         };
 
         let mut buf = [0u8; 0xff];
-        create_partition_info(&mut buf, &[desc], true);
+        PartitionInfo::pack(&[desc1, desc2], &mut buf, true);
 
-        println!("{:#x?}", &buf[0..0x0f]);
+        let mut descriptors = PartitionInfoIterator::new(&buf, 2).unwrap();
+        let desc1_check = descriptors.next().unwrap().unwrap();
+        let desc2_check = descriptors.next().unwrap().unwrap();
 
-        assert_eq!(0x8001_u16, u16::from_le_bytes([buf[0], buf[1]]));
+        assert_eq!(desc1, desc1_check);
+        assert_eq!(desc2, desc2_check);
     }
 }