Add LimitedBuf

src/io/mod.rs

@@ -35,6 +35,15 @@
 //!
 //! [multishot reads]: crate::fd::AsyncFd::multishot_read
 //!
+//! ## Other Buffer types
+//!
+//! If you hit an error that complains that "`&'0 T` must implement `Buf`, for
+//! any lifetime `'0`, but `Buf` is actually implemented for the type `&'static
+//! T`", [`StaticBuf`] is buffer to use to work around this.
+//!
+//! [`LimitedBuf`] exist to limit the amount of bytes read from or written to a
+//! buffer.
+//!
 //! # Working with Standard I/O Streams
 //!
 //! The [`stdin`], [`stdout`] and [`stderr`] functions provide handles to
@@ -71,7 +80,7 @@ pub(crate) use std::io::*;
 pub use read_buf::{ReadBuf, ReadBufPool};
 
 #[doc(inline)]
-pub use traits::{Buf, BufMut, BufMutSlice, BufSlice, IoMutSlice, IoSlice, StaticBuf};
+pub use traits::{Buf, BufMut, BufMutSlice, BufSlice, IoMutSlice, IoSlice, LimitedBuf, StaticBuf};
 
 /// Create a function and type to wraps standard {in,out,error}.
 macro_rules! stdio {

src/io/traits.rs

@@ -135,6 +135,14 @@ pub unsafe trait BufMut: 'static {
         written
     }
 
+    /// Limit the amount of bytes written to this buffer.
+    fn limit(self, limit: u32) -> LimitedBuf<Self>
+    where
+        Self: Sized,
+    {
+        LimitedBuf::new(self, limit)
+    }
+
     /// **Not part of the public API**.
     /// Do **not** implement this.
     #[doc(hidden)]
@@ -481,6 +489,14 @@ pub unsafe trait Buf: 'static {
             slice::from_raw_parts(ptr, len as usize)
         }
     }
+
+    /// Limit the amount of bytes read from this buffer.
+    fn limit(self, limit: u32) -> LimitedBuf<Self>
+    where
+        Self: Sized,
+    {
+        LimitedBuf::new(self, limit)
+    }
 }
 
 // SAFETY: `Vec<u8>` manages the allocation of the bytes, so as long as it's
@@ -888,3 +904,64 @@ unsafe impl Buf for StaticBuf {
         Buf::as_slice(&self.0)
     }
 }
+
+/// Limited buffer.
+///
+/// This buffer limits the amount of bytes read from or written to this buffer.
+///
+/// # Notes
+///
+/// This does not work work with a [`ReadBuf`] from ([`ReadBufPool`]).
+///
+/// [`ReadBuf`]: crate::io::ReadBuf
+/// [`ReadBufPool`]: crate::io::ReadBufPool
+pub struct LimitedBuf<B> {
+    buf: B,
+    limit: u32,
+}
+
+impl<B> LimitedBuf<B> {
+    /// Create a new limited buffer.
+    pub const fn new(buf: B, limit: u32) -> LimitedBuf<B> {
+        LimitedBuf { buf, limit }
+    }
+
+    /// Returns the internal buffer.
+    pub fn into_inner(self) -> B {
+        self.buf
+    }
+}
+
+unsafe impl<B: BufMut> BufMut for LimitedBuf<B> {
+    unsafe fn parts_mut(&mut self) -> (*mut u8, u32) {
+        // SAFETY: reposibilities lie with the caller.
+        let (ptr, len) = unsafe { self.buf.parts_mut() };
+        (ptr, min(len, self.limit))
+    }
+
+    unsafe fn set_init(&mut self, n: usize) {
+        // SAFETY: reposibilities lie with the caller.
+        unsafe { self.buf.set_init(n) }
+        self.limit = self.limit.saturating_sub(n as u32);
+    }
+
+    fn spare_capacity(&self) -> u32 {
+        min(self.buf.spare_capacity(), self.limit)
+    }
+}
+
+unsafe impl<B: Buf> Buf for LimitedBuf<B> {
+    unsafe fn parts(&self) -> (*const u8, u32) {
+        // SAFETY: reposibilities lie with the caller.
+        let (ptr, len) = unsafe { self.buf.parts() };
+        (ptr, min(len, self.limit))
+    }
+
+    fn len(&self) -> usize {
+        min(self.buf.len(), self.limit as usize)
+    }
+
+    fn is_empty(&self) -> bool {
+        self.buf.is_empty() || self.limit == 0
+    }
+}

tests/functional/io.rs

@@ -5,7 +5,7 @@ use std::os::fd::FromRawFd;
 use a10::fs::{self, OpenOptions};
 #[cfg(any(target_os = "android", target_os = "linux"))]
 use a10::io::Splice;
-use a10::io::{BufMut, Close, ReadBufPool, Stderr, Stdout, stderr, stdout};
+use a10::io::{Buf, BufMut, Close, ReadBufPool, Stderr, Stdout, stderr, stdout};
 use a10::pipe::pipe;
 use a10::{AsyncFd, Extract, SubmissionQueue};
 
@@ -466,6 +466,29 @@ async fn open_read_pipe(expected: &'static [u8], sq: SubmissionQueue) -> AsyncFd
     fd
 }
 
+#[test]
+fn limited_buf_buf() {
+    let mut buf = BufMut::limit(Vec::with_capacity(32), 4);
+    assert_eq!(buf.spare_capacity(), 4);
+    buf.extend_from_slice(b"Hello!");
+    assert_eq!(buf.spare_capacity(), 0);
+    assert_eq!(buf.has_spare_capacity(), false);
+    let buf = buf.into_inner();
+    assert_eq!(buf, b"Hell");
+    assert_eq!(buf.capacity(), 32);
+
+    let buf = Buf::limit(buf, 2);
+    assert_eq!(buf.len(), 2);
+    assert_eq!(buf.is_empty(), false);
+    assert_eq!(buf.as_slice(), b"He");
+    let buf = buf.into_inner();
+
+    let buf = Buf::limit(buf, 0);
+    assert_eq!(buf.len(), 0);
+    assert_eq!(buf.is_empty(), true);
+    assert_eq!(buf.as_slice(), b"");
+}
+
 /// Macro to run a code block with all buffer kinds.
 macro_rules! all_bufs {
     (