diff --git a/src/protocol/frame/mod.rs b/src/protocol/frame/mod.rs index 2d76830d..2f76c51b 100644 --- a/src/protocol/frame/mod.rs +++ b/src/protocol/frame/mod.rs @@ -19,7 +19,7 @@ use crate::{ }; use bytes::BytesMut; use log::*; -use std::io::{Cursor, Error as IoError, ErrorKind as IoErrorKind, Read, Write}; +use std::io::{self, Cursor, Error as IoError, ErrorKind as IoErrorKind, Read, Write}; /// Read buffer size used for `FrameSocket`. const READ_BUF_LEN: usize = 128 * 1024; @@ -155,16 +155,13 @@ impl FrameCodec { } /// Read a frame from the provided stream. - pub(super) fn read_frame( + pub(super) fn read_frame( &mut self, - stream: &mut Stream, + stream: &mut impl Read, max_size: Option, unmask: bool, accept_unmasked: bool, - ) -> Result> - where - Stream: Read, - { + ) -> Result> { let max_size = max_size.unwrap_or_else(usize::max_value); let mut payload = loop { @@ -195,14 +192,8 @@ impl FrameCodec { } // Not enough data in buffer. - let reserve_len = self.header.as_ref().map(|(_, l)| *l as usize).unwrap_or(6); - self.in_buffer.reserve(reserve_len); - let mut buf = self.in_buffer.split_off(self.in_buffer.len()); - buf.resize(reserve_len.max(buf.capacity()), 0); - let size = stream.read(&mut buf)?; - buf.truncate(size); - self.in_buffer.unsplit(buf); - if size == 0 { + self.in_buffer.reserve(self.header.as_ref().map(|(_, l)| *l as usize).unwrap_or(6)); + if self.read_in(stream)? == 0 { trace!("no frame received"); return Ok(None); } @@ -230,6 +221,19 @@ impl FrameCodec { Ok(Some(frame)) } + /// Read into available `in_buffer` capacity. + fn read_in(&mut self, stream: &mut impl Read) -> io::Result { + let len = self.in_buffer.len(); + debug_assert!(self.in_buffer.capacity() > len); + // SAFETY: truncated after read so uninit bytes are never read + unsafe { + self.in_buffer.set_len(self.in_buffer.capacity()); + } + let size = stream.read(&mut self.in_buffer[len..]); + self.in_buffer.truncate(len + size.as_ref().copied().unwrap_or(0)); + size + } + /// Writes a frame into the `out_buffer`. /// If the out buffer size is over the `out_buffer_write_len` will also write /// the out buffer into the provided `stream`.