Improved handling of the output buffer in read_all_timeout() function.

Author: dEajL3kA

Cargo.toml

@@ -3,7 +3,7 @@
 
 [package]
 name = "mtcp-rs"
-version = "0.1.8"
+version = "0.1.9"
 edition = "2021"
 license-file = "LICENSE"
 description = "Provides a “blocking” implementation of TcpListener and TcpStream with proper timeout and cancellation support."

src/error.rs

@@ -43,7 +43,7 @@ impl Debug for TcpError {
             Self::TimedOut => write!(f, "TcpError::TimedOut"),
             Self::Incomplete => write!(f, "TcpError::Incomplete"),
             Self::TooBig => write!(f, "TcpError::TooBig"),
-            Self::Failed(error) => write!(f, "TcpError::Failed({:?})", error),
+            Self::Failed(error) => write!(f, "TcpError::Failed({error})"),
         }
     }
 }
@@ -55,7 +55,7 @@ impl Display for TcpError {
             Self::TimedOut => write!(f, "The TCP socket operation timed out!"),
             Self::Incomplete => write!(f, "The TCP socket operation is incomplete!"),
             Self::TooBig => write!(f, "The TCP socket operation aborted, data is too big!"),
-            Self::Failed(error) => write!(f, "{}", error),
+            Self::Failed(error) => write!(f, "{error}"),
         }
     }
 }

src/stream.rs

@@ -14,7 +14,7 @@ use mio::net::TcpStream as MioTcpStream;
 
 use log::warn;
 
-use crate::utilities::Timeout;
+use crate::utilities::{Timeout, BufferManager};
 use crate::{TcpConnection, TcpManager, TcpError};
 use crate::manager::TcpPollContext;
 
@@ -284,26 +284,25 @@ impl TcpStream {
         F: Fn(&[u8]) -> bool,
     {
         let chunk_size = chunk_size.unwrap_or_else(|| NonZeroUsize::new(4096).unwrap());
-        let maximum_length = maximum_length.map(|value| NonZeroUsize::new(round_up(value.get(), chunk_size)).unwrap());
-        let mut valid_length = buffer.len();
+        if maximum_length.map_or(false, |value| value < chunk_size) {
+            panic!("maximum_length must be greater than or equal to chunk_size!")
+        }
+
+        let mut buffer = BufferManager::from(buffer, maximum_length);
 
         loop {
-            adjust_buffer(buffer, valid_length, chunk_size, maximum_length)?;
-            let done = match self.read_timeout(&mut buffer[valid_length..], timeout) {
-                Ok(0) => Some(Err(TcpError::Incomplete)),
+            let spare = buffer.alloc_spare_buffer(chunk_size);
+            match self.read_timeout(spare, timeout) {
+                Ok(0) => return Err(TcpError::Incomplete),
                 Ok(count) => {
-                    valid_length += count;
-                    match fn_complete(&buffer[..valid_length]) {
-                        true => Some(Ok(())),
-                        false => None,
+                    buffer.commit(count).map_err(|_err| TcpError::TooBig)?;
+                    match fn_complete(buffer.valid_data()) {
+                        true => return Ok(()),
+                        false => {},
                     }
                 },
-                Err(error) => Some(Err(error)),
+                Err(error) => return Err(error),
             };
-            if let Some(result) = done {
-                buffer.truncate(valid_length);
-                return result;
-            }
         }
     }
 
@@ -445,25 +444,3 @@ fn into_io_result<T>(result: Result<T, TcpError>) -> IoResult<T> {
         Err(error) => Err(error.into()),
     }
 }
-
-fn adjust_buffer(buffer: &mut Vec<u8>, length: usize, chunk_size: NonZeroUsize, maximum_length: Option<NonZeroUsize>) -> Result<(), TcpError> {
-    let mut capacity = round_up(length, chunk_size);
-    while capacity <= length {
-        capacity = capacity.checked_add(chunk_size.get()).expect("Numerical overflow!")
-    }
-    if capacity > buffer.len() {
-        if maximum_length.map_or(false, |max_len| capacity > max_len.get()) {
-            return Err(TcpError::TooBig);
-        }
-        buffer.resize(capacity, 0);
-    }
-    Ok(())
-}
-
-fn round_up(value: usize, block_size: NonZeroUsize) -> usize {
-    let block_size = block_size.get();
-    match value % block_size {
-        0 => value,
-        r => value.checked_add(block_size - r).expect("Numerical overflow!"),
-    }
-}

src/utilities/buffer.rs

@@ -0,0 +1,48 @@
+/*
+ * mtcp - TcpListener/TcpStream *with* timeout/cancellation support
+ * This is free and unencumbered software released into the public domain.
+ */
+use std::io::{Result as IoResult, Error as IoError, ErrorKind};
+use std::num::NonZeroUsize;
+use std::slice::from_raw_parts_mut;
+
+pub struct BufferManager<'a> {
+    buffer: &'a mut Vec<u8>,
+    limit: usize,
+}
+
+impl<'a> BufferManager<'a> {
+    pub fn from(buffer: &'a mut Vec<u8>, limit: Option<NonZeroUsize>) -> Self {
+        Self {
+            buffer,
+            limit: limit.map_or(usize::MAX, NonZeroUsize::get),
+        }
+    }
+
+    pub fn valid_data(&self) -> &[u8] {
+        &self.buffer[..]
+    }
+
+    pub fn alloc_spare_buffer(&mut self, min_length: NonZeroUsize) -> &'a mut[u8] {
+        self.buffer.reserve(min_length.get());
+        let spare = self.buffer.spare_capacity_mut();
+        unsafe {
+            from_raw_parts_mut(spare.as_mut_ptr() as *mut u8, spare.len())
+        }
+    }
+
+    pub fn commit(&mut self, additional: usize) -> IoResult<()> {
+        if additional > 0 {
+            let new_length = self.buffer.len().checked_add(additional).expect("Numerical overflow!");
+            if new_length <= self.limit {
+                assert!(new_length <= self.buffer.capacity());
+                unsafe {
+                    self.buffer.set_len(new_length)
+                }
+            } else {
+                return Err(IoError::new(ErrorKind::OutOfMemory, "New length exceeds the limit!"))
+            }
+        }
+        Ok(())
+    }
+}

src/utilities/mod.rs

@@ -3,7 +3,9 @@
  * This is free and unencumbered software released into the public domain.
  */
 mod flag;
+mod buffer;
 mod timeout;
 
+pub(crate) use buffer::BufferManager;
 pub(crate) use flag::Flag;
 pub(crate) use timeout::Timeout;