fix(net): CMsgBuilder get buffer reference only

Author: AsakuraMizu

compio-net/src/cmsg/mod.rs

@@ -1,7 +1,5 @@
 use std::marker::PhantomData;
 
-use compio_buf::{IoBuf, IoBufMut};
-
 cfg_if::cfg_if! {
     if #[cfg(windows)] {
         #[path = "windows.rs"]
@@ -31,9 +29,9 @@ impl<'a> CMsgIter<'a> {
     /// # Safety
     ///
     /// The buffer should contain valid control messages.
-    pub unsafe fn new<B: IoBuf>(buffer: &'a B) -> Self {
+    pub unsafe fn new(buffer: &'a [u8]) -> Self {
         Self {
-            inner: sys::CMsgIter::new(buffer.as_buf_ptr(), buffer.buf_len()),
+            inner: sys::CMsgIter::new(buffer.as_ptr(), buffer.len()),
             _p: PhantomData,
         }
     }
@@ -52,59 +50,53 @@ impl<'a> Iterator for CMsgIter<'a> {
 }
 
 /// Helper to construct control message.
-pub struct CMsgBuilder<B> {
+pub struct CMsgBuilder<'a> {
     inner: sys::CMsgIter,
-    buffer: B,
     len: usize,
+    _p: PhantomData<&'a mut ()>,
 }
 
-impl<B> CMsgBuilder<B> {
-    /// Finishes building, returns the buffer and the length of the control
-    /// message.
-    pub fn build(self) -> (B, usize) {
-        (self.buffer, self.len)
+impl<'a> CMsgBuilder<'a> {
+    /// Create [`CMsgBuilder`] with the given buffer. The buffer will be zeroed
+    /// on creation.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the buffer is too short or not properly
+    /// aligned.
+    pub fn new(buffer: &'a mut [u8]) -> Self {
+        buffer.fill(0);
+        Self {
+            inner: sys::CMsgIter::new(buffer.as_ptr(), buffer.len()),
+            len: 0,
+            _p: PhantomData,
+        }
+    }
+
+    /// Finishes building, returns length of the control message.
+    pub fn finish(self) -> usize {
+        self.len
     }
 
     /// Try to append a control message entry into the buffer. If the buffer
     /// does not have enough space or is not properly aligned with the value
     /// type, returns `None`.
-    ///
-    /// # Safety
-    ///
-    /// TODO: This function may be safe? Given that the buffer is zeroed,
-    /// properly aligned and has enough space, safety conditions of all unsafe
-    /// functions involved are satisfied, except for `CMSG_*`/`wsa_cmsg_*`, as
-    /// their safety are not documented.
-    pub unsafe fn try_push<T>(&mut self, level: i32, ty: i32, value: T) -> Option<()> {
+    pub fn try_push<T>(&mut self, level: i32, ty: i32, value: T) -> Option<()> {
         if !self.inner.is_aligned::<T>() || !self.inner.is_space_enough::<T>() {
             return None;
         }
 
-        let mut cmsg = self.inner.current_mut()?;
-        cmsg.set_level(level);
-        cmsg.set_ty(ty);
-        cmsg.set_data(value);
-
-        self.inner.next();
-        self.len += sys::space_of::<T>();
-        Some(())
-    }
-}
+        // SAFETY: the buffer is zeroed and the pointer is valid and aligned
+        unsafe {
+            let mut cmsg = self.inner.current_mut()?;
+            cmsg.set_level(level);
+            cmsg.set_ty(ty);
+            cmsg.set_data(value);
 
-impl<B: IoBufMut> CMsgBuilder<B> {
-    /// Create [`CMsgBuilder`] with the given buffer. The buffer will be zeroed
-    /// on creation.
-    ///
-    /// # Panics
-    ///
-    /// This function will panic if the buffer is too short or not properly
-    /// aligned.
-    pub fn new(mut buffer: B) -> Self {
-        buffer.as_mut_slice().fill(std::mem::MaybeUninit::zeroed());
-        Self {
-            inner: sys::CMsgIter::new(buffer.as_buf_mut_ptr(), buffer.buf_len()),
-            buffer,
-            len: 0,
+            self.inner.next();
+            self.len += sys::space_of::<T>();
         }
+
+        Some(())
     }
 }

compio-net/tests/cmsg.rs

@@ -3,48 +3,45 @@ use compio_net::{CMsgBuilder, CMsgIter};
 
 #[test]
 fn test_cmsg() {
-    let buf = vec![0u8; 64];
-    let mut builder = CMsgBuilder::new(buf);
+    let mut buf = [0u8; 64];
+    let mut builder = CMsgBuilder::new(&mut buf);
 
-    unsafe {
-        builder.try_push(0, 0, ()).unwrap(); // 16 / 12
-        builder.try_push(1, 1, u32::MAX).unwrap(); // 16 + 4 + 4 / 12 + 4
-        builder.try_push(2, 2, i64::MIN).unwrap(); // 16 + 8 / 12 + 8
-    }
-    let (buf, len) = builder.build();
+    builder.try_push(0, 0, ()).unwrap(); // 16 / 12
+    builder.try_push(1, 1, u32::MAX).unwrap(); // 16 + 4 + 4 / 12 + 4
+    builder.try_push(2, 2, i64::MIN).unwrap(); // 16 + 8 / 12 + 8
+    let len = builder.finish();
     assert!(len == 64 || len == 48);
 
-    let buf = buf.slice(..len);
-    let mut iter = unsafe { CMsgIter::new(&buf) };
+    unsafe {
+        let buf = buf.slice(..len);
+        let mut iter = CMsgIter::new(&buf);
 
-    let cmsg = iter.next().unwrap();
-    assert_eq!(
-        (cmsg.level(), cmsg.ty(), unsafe { cmsg.data::<()>() }),
-        (0, 0, &())
-    );
-    let cmsg = iter.next().unwrap();
-    assert_eq!(
-        (cmsg.level(), cmsg.ty(), unsafe { cmsg.data::<u32>() }),
-        (1, 1, &u32::MAX)
-    );
-    let cmsg = iter.next().unwrap();
-    assert_eq!(
-        (cmsg.level(), cmsg.ty(), unsafe { cmsg.data::<i64>() }),
-        (2, 2, &i64::MIN)
-    );
-    assert!(iter.next().is_none());
+        let cmsg = iter.next().unwrap();
+        assert_eq!((cmsg.level(), cmsg.ty(), cmsg.data::<()>()), (0, 0, &()));
+        let cmsg = iter.next().unwrap();
+        assert_eq!(
+            (cmsg.level(), cmsg.ty(), cmsg.data::<u32>()),
+            (1, 1, &u32::MAX)
+        );
+        let cmsg = iter.next().unwrap();
+        assert_eq!(
+            (cmsg.level(), cmsg.ty(), cmsg.data::<i64>()),
+            (2, 2, &i64::MIN)
+        );
+        assert!(iter.next().is_none());
+    }
 }
 
 #[test]
 #[should_panic]
 fn invalid_buffer_length() {
-    let buf = vec![0u8; 1];
-    CMsgBuilder::new(buf);
+    let mut buf = [0u8; 1];
+    CMsgBuilder::new(&mut buf);
 }
 
 #[test]
 #[should_panic]
 fn invalid_buffer_alignment() {
-    let buf = vec![0u8; 64];
-    CMsgBuilder::new(buf.slice(1..));
+    let mut buf = [0u8; 64];
+    CMsgBuilder::new(&mut buf[1..]);
 }