fix: QueryDownstreamIdentifiersResponse and Device parsing

Signed-off-by: leongross <leon.gross@9elements.com>

Author: leongross

src/message/firmware_update/query_devid.rs

@@ -213,7 +213,7 @@ mod test {
     use super::*;
     use crate::protocol::firmware_update::{Descriptor, DescriptorType};
 
-    #[ignore]
+    //#[ignore]
     #[test]
     fn test_query_device_identifiers_resp_codec() {
         let instance_id = 0;

src/message/firmware_update/query_downstream.rs

@@ -129,6 +129,43 @@ pldm_completion_code! {
 
 #[derive(Debug, Clone, PartialEq)]
 #[repr(C)]
+/// The total structure for QueryDownstreamIdentifiersResponse looks as follows:
+/// ```text
+/// QueryDownstreamIdentifiersResponse
+///  completion_code (u8)
+///  next_data_transfer_handle (u32)
+///  transfer_flag (u8)
+///  --- portion (variable-length)
+///    downstream_devices_length_i (u32)
+///    number_of_downstream_devices_i (u16)
+///    downstream_devices_index_i (u16)
+///    ---
+///      downstream_device_index_ij (u16)
+///      downstream_descriptor_count_ij (u8)
+///      ---
+///        descriptor_type_ijk (u16)
+///        descriptor_length_ijk (u16)
+///        descriptor_data_ijk (variable-length L_ijk)
+///           ...
+///        descriptor_type_ij(k+1) (u16)
+///        descriptor_length_ij(k+1) (u16)
+///        descriptor_data_ij(k+1) (variable-length L_ij(k+1))
+///           ...
+///     ...
+///     downstream_devices_index_i(j+1) (u16)
+///     downstream_device_count_i(j+1) (u8)
+///     ---
+///       descriptor_type_(i(j+1)k) (u16)
+///       descriptor_length_(i(j+1)k) (u16)
+///       descriptor_data_(i(j+1)k) (variable-length L_(i(j+1)k))
+///         ...
+///      ...
+///    downstream_devices_length_(i+1) (u32)
+///    number_of_downstream_devices_(i+1) (u16)
+///    downstream_devices_index_(i+1) (u16)
+///    ...
+/// ...
+/// ```
 pub struct QueryDownstreamIdentifiersResponse<'a> {
     pub hdr: PldmMsgHeader<[u8; PLDM_MSG_HEADER_LEN]>,
 
@@ -227,10 +264,10 @@ impl QueryDownstreamIdentifiersResponse<'_> {
     }
 }
 
-/// Iterate over all available [DownstreamDevice] in the response portion.
 impl<'a> Iterator for QueryDownstreamIdentifiersResponse<'a> {
     type Item = DownstreamDevice<'a>;
 
+    /// Iterate over all available [DownstreamDevice] in the response portion.
     fn next(&mut self) -> Option<Self::Item> {
         let portion_hdr = self.try_get_portion_header().ok()?;
         if self.portion_iter_current >= portion_hdr.number_of_downstream_devices as usize {
@@ -249,14 +286,18 @@ impl<'a> Iterator for QueryDownstreamIdentifiersResponse<'a> {
             return None;
         }
 
-        let hdr = self
+        let hdr: DownstreamDevicesHeader = self
             .try_get_downstream_device_header(&self.portion[offset..])
             .ok()?;
         let device_header_size = size_of::<DownstreamDevicesHeader>();
-        let descriptors_size = hdr.downstream_descriptor_count as usize * size_of::<Descriptor>();
-        let total_device_size = device_header_size + descriptors_size;
 
-        let next_offset = offset + total_device_size;
+        let desc_size = Descriptor::try_get_descriptor_length_from_blob(
+            &self.portion[offset + device_header_size..],
+            hdr.downstream_descriptor_count as usize,
+        )
+        .ok()?;
+
+        let next_offset = offset + device_header_size + desc_size;
         if next_offset > self.portion.len() {
             self.portion_iter_current = 0;
             self.portion_offset_next = 0;
@@ -423,19 +464,27 @@ impl Iterator for DownstreamDevice<'_> {
             return None;
         }
 
-        let descriptor_size = size_of::<Descriptor>();
-        if self.downstream_descriptors.len() < descriptor_size {
-            self._iter_dev_count = 0;
-            self._iter_offset = 0;
+        // Read the descriptor length while skipping the type field
+        let descriptor_length = u16::read_from_bytes(
+            &self.downstream_descriptors
+                [self._iter_offset + size_of::<u16>()..self._iter_offset + size_of::<u16>() * 2],
+        )
+        .ok()? as usize;
+
+        // check bounds
+        if self._iter_offset + size_of::<u16>() * 2 + descriptor_length
+            > self.downstream_descriptors.len()
+        {
             return None;
         }
 
-        let descriptor = Descriptor::try_read_from_prefix(
-            &self.downstream_descriptors[self._iter_offset..self._iter_offset + descriptor_size],
+        let descriptor = Descriptor::decode(
+            &self.downstream_descriptors
+                [self._iter_offset..self._iter_offset + 2 * size_of::<u16>() + descriptor_length],
         )
-        .ok()?
-        .0;
-        self._iter_offset += descriptor_size;
+        .ok()?;
+
+        self._iter_offset += size_of::<u16>() * 2 + descriptor_length;
         self._iter_dev_count += 1;
 
         Some(descriptor)
@@ -1141,7 +1190,7 @@ impl RequestDownstreamDeviceUpdateResponse {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::codec::PldmCodec;
+    use crate::{codec::PldmCodec, protocol::firmware_update::DescriptorType};
 
     #[test]
     fn test_query_downstream_devices_request_codec() {
@@ -1389,15 +1438,18 @@ mod tests {
         let descriptors: [Descriptor; 3] = [descriptor.clone(), descriptor.clone(), descriptor];
 
         const DESC_LEN: usize = size_of::<Descriptor>();
-        let mut descriptor_bytes: [u8; DESC_LEN * 3] = [0u8; DESC_LEN * 3];
+        let mut descriptor_bytes_max: [u8; DESC_LEN * 3] = [0u8; DESC_LEN * 3];
         let mut offset = 0;
+
+        // encoding
         for desc in descriptors.iter() {
-            descriptor_bytes[offset..offset + DESC_LEN].copy_from_slice(&desc.as_bytes());
-            offset += DESC_LEN;
+            let encoded = desc.encode(&mut descriptor_bytes_max[offset..]).unwrap();
+            offset += encoded;
         }
 
         let downstream_device_index = DownstreamDeviceIndex::try_from(1).unwrap();
-        let downstream_device = DownstreamDevice::new(downstream_device_index, &descriptor_bytes);
+        let downstream_device =
+            DownstreamDevice::new(downstream_device_index, &descriptor_bytes_max);
 
         let mut buffer = [0u8; 256];
         let bytes_written = downstream_device.encode(&mut buffer).unwrap();
@@ -1407,6 +1459,7 @@ mod tests {
 
     #[test]
     fn test_iterator_query_downstream_identifiers_response() {
+        const DSC_DATA_LEN: usize = 16;
         let instance_id: InstanceId = 0x01;
         let ph: PortionHeader = PortionHeader {
             downstream_devices_length: 1,
@@ -1417,22 +1470,21 @@ mod tests {
             downstream_descriptor_count: 3,
         };
         let dsc_0: Descriptor = Descriptor {
-            descriptor_type: 0xff,
-            descriptor_length: 0x02,
+            descriptor_type: DescriptorType::VendorDefined as u16,
+            descriptor_length: DSC_DATA_LEN as u16,
             descriptor_data: [0u8; 64],
         };
 
         let mut dsc_1 = dsc_0.clone();
-        dsc_1.descriptor_type = 0xfe;
-        dsc_1.descriptor_data = [1u8; 64];
+        dsc_1.descriptor_data[0..16].clone_from_slice(&[1u8; 16]);
 
         let mut dsc_2 = dsc_0.clone();
-        dsc_2.descriptor_type = 0xfd;
-        dsc_2.descriptor_data = [2u8; 64];
+        dsc_2.descriptor_data[0..16].clone_from_slice(&[2u8; 16]);
 
         const LEN: usize = size_of::<PortionHeader>()
             + size_of::<DownstreamDevicesHeader>()
-            + 3 * size_of::<Descriptor>();
+            + 3 * (2 * size_of::<u16>() + DSC_DATA_LEN); // type + length + data
+
         let mut offset = 0;
         let mut portion: [u8; LEN] = [0u8; LEN];
 
@@ -1443,13 +1495,10 @@ mod tests {
             .copy_from_slice(&dsdh.as_bytes());
         offset += size_of::<DownstreamDevicesHeader>();
 
-        portion[offset..offset + size_of::<Descriptor>()].copy_from_slice(&dsc_0.as_bytes());
-        offset += size_of::<Descriptor>();
-
-        portion[offset..offset + size_of::<Descriptor>()].copy_from_slice(&dsc_1.as_bytes());
-        offset += size_of::<Descriptor>();
-
-        portion[offset..offset + size_of::<Descriptor>()].copy_from_slice(&dsc_2.as_bytes());
+        for desc in [dsc_0.clone(), dsc_1.clone(), dsc_2.clone()].iter() {
+            let size = &desc.encode(&mut portion[offset..]).unwrap();
+            offset += size;
+        }
 
         let mut qdir = QueryDownstreamIdentifiersResponse::new(
             instance_id,

src/protocol/firmware_update.rs

@@ -5,7 +5,7 @@ use crate::error::PldmError;
 use bitfield::bitfield;
 use core::convert::TryFrom;
 use core::fmt;
-use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
+use zerocopy::{FromBytes, Immutable, IntoBytes};
 
 pub const PLDM_FWUP_COMPONENT_RELEASE_DATA_LEN: usize = 8;
 pub const PLDM_FWUP_BASELINE_TRANSFER_SIZE: usize = 32;
@@ -285,7 +285,7 @@ pub fn get_descriptor_length(descriptor_type: DescriptorType) -> usize {
     }
 }
 
-#[derive(Debug, Copy, Clone, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)]
+#[derive(Debug, Copy, Clone, PartialEq, FromBytes)]
 #[repr(C)]
 pub struct Descriptor {
     pub descriptor_type: u16,
@@ -331,6 +331,112 @@ impl Descriptor {
     pub fn codec_size_in_bytes(&self) -> usize {
         core::mem::size_of::<u16>() * 2 + self.descriptor_length as usize
     }
+
+    /// Calculates the total length of a list of descriptors from a blob.
+    ///
+    /// This simplifies the use case where multiple descriptors are used in a message,
+    /// and the total length needs to be calculated.
+    /// It follows the linked list of descriptors by decoding each one in sequence.
+    ///
+    /// The blob must start with the first descriptor, but may contain additional data after the last descriptor.
+    ///
+    /// ```rust
+    /// use pldm_lib::protocol::firmware_update::{Descriptor, DescriptorType, get_descriptor_length};
+    /// use crate::pldm_lib::codec::PldmCodec;
+    ///
+    /// let dsc_0 = Descriptor::new(DescriptorType::PciVendorId, &[0x11, 0x22]).unwrap();
+    /// let dsc_1 = Descriptor::new(DescriptorType::PciVendorId, &[0x33, 0x44]).unwrap();
+    /// const BLOB_LEN: usize = (size_of::<u16>() * 2 + 2 * size_of::<u8>()) * 2 as usize;
+    /// let mut blob = [0u8; BLOB_LEN];
+    ///
+    /// let offset = dsc_0.encode(&mut blob[0..]).unwrap();
+    /// let _ = dsc_1.encode(&mut blob[offset..]).unwrap();
+    ///
+    /// let total_len = Descriptor::try_get_descriptor_length_from_blob(&blob, 2).unwrap();
+    /// assert_eq!(total_len, dsc_0.codec_size_in_bytes() + dsc_1.codec_size_in_bytes());
+    /// ```
+    pub fn try_get_descriptor_length_from_blob(
+        blob: &[u8],
+        desc_count: usize,
+    ) -> Result<usize, PldmCodecError> {
+        let mut offset = 0;
+        let mut total_len = 0;
+
+        for _ in 0..desc_count {
+            let descriptor = Descriptor::decode(&blob[offset..])?;
+            total_len += descriptor.codec_size_in_bytes();
+            offset += descriptor.codec_size_in_bytes();
+        }
+
+        Ok(total_len)
+    }
+}
+
+impl PldmCodec for Descriptor {
+    /// Custom encode implementation of PldmCodec for Descriptor to handle variable-length.
+    ///
+    /// The descriptor_data field is variable-length, so we need to ensure that only the valid portion
+    /// of the array is encoded, not all data of the fixed size array of length [DESCRIPTOR_DATA_MAX_LEN].
+    fn encode(&self, buffer: &mut [u8]) -> Result<usize, PldmCodecError> {
+        if buffer.len() < self.codec_size_in_bytes() {
+            return Err(PldmCodecError::BufferTooShort);
+        }
+        let mut offset = 0;
+
+        self.descriptor_type
+            .write_to(&mut buffer[offset..offset + core::mem::size_of::<u16>()])
+            .unwrap();
+        offset += core::mem::size_of::<u16>();
+
+        self.descriptor_length
+            .write_to(&mut buffer[offset..offset + core::mem::size_of::<u16>()])
+            .unwrap();
+        offset += core::mem::size_of::<u16>();
+
+        self.descriptor_data[..self.descriptor_length as usize]
+            .write_to(&mut buffer[offset..offset + self.descriptor_length as usize])
+            .unwrap();
+        offset += self.descriptor_length as usize;
+
+        Ok(offset)
+    }
+
+    /// Custom decode implementation of PldmCodec for Descriptor to handle variable-length.
+    ///
+    /// The descriptor_data field is variable-length, so we need to ensure that only the valid portion
+    /// of the array is decoded, not all data of the fixed size array of length [DESCRIPTOR_DATA_MAX_LEN].
+    fn decode(buffer: &[u8]) -> Result<Self, PldmCodecError> {
+        let mut offset = 0;
+
+        let descriptor_type = u16::read_from_bytes(
+            buffer
+                .get(offset..offset + core::mem::size_of::<u16>())
+                .ok_or(PldmCodecError::BufferTooShort)?,
+        )
+        .unwrap();
+        offset += core::mem::size_of::<u16>();
+
+        let descriptor_length = u16::read_from_bytes(
+            buffer
+                .get(offset..offset + core::mem::size_of::<u16>())
+                .ok_or(PldmCodecError::BufferTooShort)?,
+        )
+        .unwrap();
+        offset += core::mem::size_of::<u16>();
+
+        let mut descriptor_data = [0u8; DESCRIPTOR_DATA_MAX_LEN];
+        descriptor_data[..descriptor_length as usize].copy_from_slice(
+            buffer
+                .get(offset..offset + descriptor_length as usize)
+                .ok_or(PldmCodecError::BufferTooShort)?,
+        );
+
+        Ok(Descriptor {
+            descriptor_type,
+            descriptor_length,
+            descriptor_data,
+        })
+    }
 }
 
 bitfield! {
@@ -868,10 +974,28 @@ mod test {
             descriptor.descriptor_length,
             get_descriptor_length(DescriptorType::Uuid) as u16
         );
-        let mut buffer = [0u8; 512];
+        let mut buffer = [0u8; core::mem::size_of::<Descriptor>()];
         descriptor.encode(&mut buffer).unwrap();
+
         let decoded_descriptor = Descriptor::decode(&buffer).unwrap();
         assert_eq!(descriptor, decoded_descriptor);
+
+        let raw_descriptor = [
+            0xff, 0xff, // Type=VendorDefined
+            0x10, 0x00, // size=16
+            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E,
+            0x0F, 0x10,
+        ];
+        let decoded_raw_descriptor = Descriptor::decode(&raw_descriptor).unwrap();
+        assert_eq!(
+            decoded_raw_descriptor.descriptor_type,
+            DescriptorType::VendorDefined as u16,
+        );
+        assert_eq!(decoded_raw_descriptor.descriptor_length, 16,);
+        assert_eq!(
+            &decoded_raw_descriptor.descriptor_data[..16],
+            &raw_descriptor[4..20]
+        );
     }
 
     #[test]