Implement a dedicated reader for vectored buffers

embediver · embediver · commit 793949a2dc68 · 2025-10-20T12:50:31.000+02:00
Signed-off-by: Marvin Gudel &lt;marvin.gudel@9elements.com&gt;
diff --git a/mctp-estack/src/fragment.rs b/mctp-estack/src/fragment.rs
@@ -10,7 +10,7 @@ use crate::fmt::{debug, error, info, trace, warn};
 
 use mctp::{Eid, Error, MsgIC, MsgType, Result, Tag};
 
-use crate::{AppCookie, MctpHeader};
+use crate::{util::VectorReader, AppCookie, MctpHeader};
 
 /// Fragments a MCTP message.
 ///
@@ -24,8 +24,8 @@ pub struct Fragmenter {
 
     cookie: Option<AppCookie>,
 
-    // A count of how many bytes have already been sent.
-    payload_used: usize,
+    // A reader to read from the payload vector
+    reader: VectorReader,
 }
 
 impl Fragmenter {
@@ -59,7 +59,7 @@ impl Fragmenter {
         };
 
         Ok(Self {
-            payload_used: 0,
+            reader: VectorReader::new(),
             header,
             typ,
             mtu,
@@ -123,9 +123,7 @@ impl Fragmenter {
             return SendOutput::success(self);
         }
 
-        let total_payload_len =
-            payload.iter().fold(0, |acc, part| acc + part.len());
-        if total_payload_len < self.payload_used {
+        if self.reader.is_exhausted(payload).is_err() {
             // Caller is passing varying payload buffers
             debug!("varying payload");
             return SendOutput::failure(Error::BadArgument, self);
@@ -149,14 +147,13 @@ impl Fragmenter {
             rest = &mut rest[1..];
         }
 
-        let remaining_payload_len = total_payload_len - self.payload_used;
-        let l = remaining_payload_len.min(rest.len());
-        let (d, rest) = rest.split_at_mut(l);
-        crate::util::copy_vectored(payload, self.payload_used, d);
-        self.payload_used += l;
+        let Ok(n) = self.reader.read(payload, &mut rest) else {
+            return SendOutput::failure(Error::BadArgument, self);
+        };
+        let rest = &rest[n..];
 
         // Add the header
-        if self.payload_used == total_payload_len {
+        if self.reader.is_exhausted(payload).unwrap() {
             self.header.eom = true;
         }
         // OK unwrap: seq and tag are valid.
diff --git a/mctp-estack/src/util.rs b/mctp-estack/src/util.rs
@@ -19,96 +19,136 @@ macro_rules! get_build_var {
     }};
 }
 
-/// Copy data from a vectored src to dest
+/// A reader to read a vector of byte slices
 ///
-/// Copies `dest.len()` bytes from payload to dest,
-/// starting after `offset` bytes.
-///
-/// ## Panics
-///
-/// This function will panic when not enough bytes are available to fill dest.
-/// Total size of `payload` has to be `atleast dest.len()` + `offset`.
-pub fn copy_vectored(src: &[&[u8]], offset: usize, dest: &mut [u8]) {
-    let mut i = 0;
+#[derive(Debug)]
+pub struct VectorReader {
+    /// The index of the current slice
+    ///
+    /// Set to `vector.len()` when exhausted.
+    slice_index: usize,
+    /// The index in the current slice
+    ///
+    /// E.g. the element to be read next.
+    current_slice_offset: usize,
+}
 
-    while i < dest.len() {
-        let payload_index = i + offset;
-        let next = get_sub_slice(src, payload_index);
-        let remaining = dest.len() - i;
-        if remaining > next.len() {
-            dest[i..(i + next.len())].copy_from_slice(next);
-            i += next.len();
-        } else {
-            dest[i..].copy_from_slice(&next[..remaining]);
-            return;
+impl VectorReader {
+    /// Create a new reader
+    pub fn new() -> Self {
+        VectorReader {
+            slice_index: 0,
+            current_slice_offset: 0,
         }
     }
-}
+    /// Read `dest.len()` bytes from `src` into `dest`, returning how many bytes were read
+    ///
+    /// Returns a [VectorReaderError] when the current position is out of range for `src`.
+    ///
+    /// The same `src` buffer has to be passed to subsequent calls to `read()`.
+    /// Changing the vector is undefined behaviour.
+    pub fn read(
+        &mut self,
+        src: &[&[u8]],
+        dest: &mut [u8],
+    ) -> Result<usize, VectorReaderError> {
+        let mut i = 0;
+        while i < dest.len() {
+            if self.is_exhausted(src)? {
+                return Ok(i);
+            }
 
-/// Get a slice of `vector` indexed by `offset`
-///
-/// The `offset` is the absolute byte index.
-/// The returned slice is the remaining sub slice starting at `offset`.
-///
-/// ## Panics
-///
-/// Will panic when offset is larger than the size of `vector`.
-///
-/// ## Example
-/// ```ignore
-/// # use mctp_estack::fragment::get_slice;
-/// let vector: &[&[u8]] = &[&[1, 2, 3], &[4, 5, 6]];
-///
-/// let slice = get_slice(vector, 4);
-///
-/// assert_eq!(slice, &[5, 6]);
-/// ```
-pub fn get_sub_slice<'a>(vector: &'a [&[u8]], offset: usize) -> &'a [u8] {
-    let mut i = offset;
-    for slice in vector {
-        if i >= slice.len() {
-            i -= slice.len();
-        } else {
-            return &slice[i..];
+            let slice = &src[self.slice_index][self.current_slice_offset..];
+            let n = slice.len().min(dest[i..].len());
+            dest[i..i + n].copy_from_slice(&slice[..n]);
+            i += n;
+            self.increment_index(src, n);
+        }
+        Ok(i)
+    }
+    /// Checks if `src` has been read to the end
+    ///
+    /// Returns a [VectorReaderError] when the current position is out of range for `src`.
+    ///
+    /// _Note:_ Might return a `Ok` even if the `src` vector changed between calls.
+    pub fn is_exhausted(
+        &self,
+        src: &[&[u8]],
+    ) -> Result<bool, VectorReaderError> {
+        if src.len() == self.slice_index {
+            return Ok(true);
+        }
+        // This shlould only occur if the caller passed varying vectors
+        src.get(self.slice_index).ok_or(VectorReaderError)?;
+        Ok(false)
+    }
+    /// Increment the index by `n`, panic if out ouf bounds
+    ///
+    /// If this exhausts the vector exactly, the index is incremented to `vector[vector.len()][0]`
+    fn increment_index(&mut self, vector: &[&[u8]], n: usize) {
+        let mut n = n;
+        loop {
+            if vector[self.slice_index]
+                .get(self.current_slice_offset + n)
+                .is_some()
+            {
+                // If we can index the current slice at offset + n just increment offset and return
+                self.current_slice_offset += n;
+                return;
+            } else {
+                // Substract what has been read from the current slice, then increment to next slice
+                n -=
+                    vector[self.slice_index][self.current_slice_offset..].len();
+                self.slice_index += 1;
+                self.current_slice_offset = 0;
+                if self.slice_index == vector.len() {
+                    // return when the end of the vector is reached
+                    debug_assert_eq!(n, 0);
+                    return;
+                }
+            }
         }
     }
-    panic!("offset for vector out of bounds");
 }
 
+#[derive(Debug)]
+pub struct VectorReaderError;
+
 #[cfg(test)]
 mod tests {
     #[test]
-    fn test_get_slice() {
-        use super::get_sub_slice;
-        let vector: &[&[u8]] = &[&[1, 2, 3], &[4, 5, 6], &[7, 8, 9]];
-        let slice = get_sub_slice(vector, 4);
-        assert_eq!(slice, &[5, 6]);
-        let slice = get_sub_slice(vector, 0);
-        assert_eq!(slice, &[1, 2, 3]);
-        let slice = get_sub_slice(vector, 3);
-        assert_eq!(slice, &[4, 5, 6]);
-    }
-    #[test]
-    fn test_copy_vectored() {
-        use super::copy_vectored;
+    fn test_vector_reader() {
+        use super::VectorReader;
+        let mut reader = VectorReader::new();
         let vector: &[&[u8]] = &[&[1, 2, 3], &[4, 5], &[6, 7, 8, 9]];
 
+        // Test reading a vector partially
+        let mut dest = [0; 4];
+        let n = reader.read(vector, &mut dest).unwrap();
+        assert_eq!(n, 4);
+        assert_eq!(&dest, &[1, 2, 3, 4]);
+
+        // Test reading all remaining elements into a larger than necessary destination
         let mut dest = [0; 6];
-        copy_vectored(vector, 1, &mut dest);
-        assert_eq!(&dest, &[2, 3, 4, 5, 6, 7]);
+        let n = reader.read(vector, &mut dest).unwrap();
+        assert_eq!(n, 5);
+        assert_eq!(&dest[..5], &[5, 6, 7, 8, 9]);
 
-        let mut dest = [0; 5];
-        copy_vectored(vector, 4, &mut dest);
-        assert_eq!(&dest, &[5, 6, 7, 8, 9]);
+        assert!(reader
+            .is_exhausted(vector)
+            .expect("Vector should be exhausted"));
 
-        let mut dest = [0; 9];
-        copy_vectored(vector, 0, &mut dest);
-        assert_eq!(&dest, &[1, 2, 3, 4, 5, 6, 7, 8, 9]);
+        // Test reading to end in one pass
+        let mut reader = VectorReader::new();
+        let vector: &[&[u8]] = &[&[1, 2, 3], &[4]];
 
-        let vector: &[&[u8]] = &[&[1, 2, 3]];
+        let mut dest = [0; 4];
+        let n = reader.read(vector, &mut dest).unwrap();
+        assert_eq!(n, 4);
+        assert_eq!(&dest, &[1, 2, 3, 4]);
 
-        let mut dest = [0; 1];
-        copy_vectored(vector, 2, &mut dest);
-        assert_eq!(&dest, &[3]);
+        assert!(reader
+            .is_exhausted(vector)
+            .expect("Vector should be exhausted"));
     }
 }
