perf: fewer syscalls, copies and allocations on I/O paths (2.4.0)

Cargo.toml

@@ -26,7 +26,7 @@ print_stderr = "warn"
 
 [package]
 name = "mc_protocol"
-version = "2.3.0"
+version = "2.4.0"
 edition.workspace = true
 rust-version.workspace = true
 authors.workspace = true
@@ -47,7 +47,7 @@ workspace = true
 [features]
 default = ["async", "compression"]
 # Enable async I/O via tokio
-async = ["dep:tokio"]
+async = ["dep:tokio", "dep:bytes"]
 # Enable AES-128-CFB8 encryption support (pure Rust, RustCrypto)
 encryption = ["dep:aes", "dep:cfb8"]
 # Enable zlib packet compression support
@@ -59,6 +59,7 @@ thiserror = "2.0.21"
 uuid = "1.27.0"
 
 tokio = { version = "1.53.2", features = ["io-util", "net"], optional = true }
+bytes = { version = "1.12.1", optional = true }
 aes = { version = "0.9.3", optional = true }
 cfb8 = { version = "0.9.1", optional = true }
 flate2 = { version = "1.1.10", optional = true }

src/compression.rs

@@ -16,10 +16,27 @@
 //! assert_eq!(&decompressed, data);
 //! ```
 
-use flate2::{Compression, read::ZlibDecoder, write::ZlibEncoder};
-use std::io::{Read, Write};
+use flate2::{
+    Compress, Compression, Decompress, FlushCompress, FlushDecompress, Status,
+    bufread::ZlibDecoder, write::ZlibEncoder,
+};
+use std::cell::RefCell;
+use std::io::{self, Read, Write};
 use thiserror::Error;
 
+/// Upper bound on the output buffer reserved up front, before any data is inflated.
+const MAX_PREALLOC: usize = 16 * 1024 * 1024;
+
+thread_local! {
+    // Deflate/inflate state is large and costly to initialise, so reuse it per thread.
+    static DEFLATE: RefCell<Compress> = RefCell::new(Compress::new(Compression::default(), true));
+    static INFLATE: RefCell<Decompress> = RefCell::new(Decompress::new(true));
+}
+
+fn invalid_data(msg: impl Into<String>) -> CompressionError {
+    CompressionError::Io(io::Error::new(io::ErrorKind::InvalidData, msg.into()))
+}
+
 /// Errors that can occur during compression or decompression.
 #[derive(Debug, Error)]
 pub enum CompressionError {
@@ -39,17 +56,31 @@ impl From<CompressionError> for crate::packet::PacketError {
 ///
 /// Returns the compressed bytes.
 pub fn compress_zlib(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
-    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
-    encoder.write_all(data)?;
-    Ok(encoder.finish()?)
+    DEFLATE.with_borrow_mut(|deflate| {
+        deflate.reset();
+        let mut output = Vec::with_capacity(data.len() / 2 + 64);
+        loop {
+            let consumed = deflate.total_in() as usize;
+            let status = deflate
+                .compress_vec(&data[consumed..], &mut output, FlushCompress::Finish)
+                .map_err(|e| invalid_data(e.to_string()))?;
+            if status == Status::StreamEnd {
+                return Ok(output);
+            }
+            output.reserve(output.capacity().max(64));
+        }
+    })
 }
 
 /// Compress `data` using zlib deflate at a specific compression level.
 ///
 /// Valid levels are 0 (no compression) through 9 (best compression);
 /// values above 9 are clamped to 9.
 pub fn compress_zlib_level(data: &[u8], level: u32) -> Result<Vec<u8>, CompressionError> {
-    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::new(level.min(9)));
+    let mut encoder = ZlibEncoder::new(
+        Vec::with_capacity(data.len() / 2 + 64),
+        Compression::new(level.min(9)),
+    );
     encoder.write_all(data)?;
     Ok(encoder.finish()?)
 }
@@ -72,15 +103,33 @@ pub fn decompress_zlib(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
 /// exceed `limit` — this guards against zip bombs, where a tiny compressed
 /// input expands to a huge output.
 pub fn decompress_zlib_limited(data: &[u8], limit: usize) -> Result<Vec<u8>, CompressionError> {
-    // Read at most limit + 1 bytes: anything past `limit` means the input lied.
-    let mut decoder = ZlibDecoder::new(data).take(limit as u64 + 1);
-    let mut output = Vec::new();
-    decoder.read_to_end(&mut output)?;
-    if output.len() > limit {
-        return Err(CompressionError::Io(std::io::Error::new(
-            std::io::ErrorKind::InvalidData,
-            format!("decompressed size exceeds declared limit of {limit} bytes"),
-        )));
-    }
-    Ok(output)
+    // Room for limit + 1 bytes: producing more than `limit` means the input lied.
+    let max_len = limit.saturating_add(1);
+    INFLATE.with_borrow_mut(|inflate| {
+        inflate.reset(true);
+        let mut output = Vec::with_capacity(max_len.min(MAX_PREALLOC));
+        loop {
+            let (consumed, produced) = (inflate.total_in(), inflate.total_out());
+            let status = inflate
+                .decompress_vec(
+                    &data[consumed as usize..],
+                    &mut output,
+                    FlushDecompress::Finish,
+                )
+                .map_err(|e| invalid_data(e.to_string()))?;
+            if output.len() > limit {
+                return Err(invalid_data(format!(
+                    "decompressed size exceeds declared limit of {limit} bytes"
+                )));
+            }
+            if status == Status::StreamEnd {
+                return Ok(output);
+            }
+            if output.len() == output.capacity() {
+                output.reserve(output.len().max(64).min(max_len - output.len()));
+            } else if (consumed, produced) == (inflate.total_in(), inflate.total_out()) {
+                return Err(CompressionError::Io(io::ErrorKind::UnexpectedEof.into()));
+            }
+        }
+    })
 }

src/encryption.rs

@@ -70,9 +70,14 @@ impl Cfb8Encryptor {
     /// Encrypt a slice of plaintext bytes, returning the ciphertext.
     pub fn encrypt(&mut self, plaintext: &[u8]) -> io::Result<Vec<u8>> {
         let mut output = plaintext.to_vec();
-        self.cipher.encrypt(&mut output);
+        self.encrypt_in_place(&mut output);
         Ok(output)
     }
+
+    /// Encrypt `data` in place, without allocating.
+    pub fn encrypt_in_place(&mut self, data: &mut [u8]) {
+        self.cipher.encrypt(data);
+    }
 }
 
 /// Stateful AES-128-CFB8 decryptor.
@@ -92,9 +97,14 @@ impl Cfb8Decryptor {
     /// Decrypt a slice of ciphertext bytes, returning the plaintext.
     pub fn decrypt(&mut self, ciphertext: &[u8]) -> io::Result<Vec<u8>> {
         let mut output = ciphertext.to_vec();
-        self.cipher.decrypt(&mut output);
+        self.decrypt_in_place(&mut output);
         Ok(output)
     }
+
+    /// Decrypt `data` in place, without allocating.
+    pub fn decrypt_in_place(&mut self, data: &mut [u8]) {
+        self.cipher.decrypt(data);
+    }
 }
 
 // ---------------------------------------------------------------------------
@@ -128,7 +138,7 @@ impl<R: Read> Read for Cfb8ReadHalf<R> {
         if n == 0 {
             return Ok(0);
         }
-        self.decryptor.cipher.decrypt(&mut buf[..n]);
+        self.decryptor.decrypt_in_place(&mut buf[..n]);
         Ok(n)
     }
 }
@@ -137,6 +147,8 @@ impl<R: Read> Read for Cfb8ReadHalf<R> {
 pub struct Cfb8WriteHalf<W> {
     inner: W,
     encryptor: Cfb8Encryptor,
+    /// Scratch space for ciphertext, reused across writes.
+    buffer: Vec<u8>,
 }
 
 impl<W: Write> Cfb8WriteHalf<W> {
@@ -145,6 +157,7 @@ impl<W: Write> Cfb8WriteHalf<W> {
         Ok(Self {
             inner,
             encryptor: Cfb8Encryptor::new(key)?,
+            buffer: Vec::new(),
         })
     }
 
@@ -156,8 +169,10 @@ impl<W: Write> Cfb8WriteHalf<W> {
 
 impl<W: Write> Write for Cfb8WriteHalf<W> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        let encrypted = self.encryptor.encrypt(buf)?;
-        self.inner.write_all(&encrypted)?;
+        self.buffer.clear();
+        self.buffer.extend_from_slice(buf);
+        self.encryptor.encrypt_in_place(&mut self.buffer);
+        self.inner.write_all(&self.buffer)?;
         Ok(buf.len())
     }
 
@@ -298,9 +313,11 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for AsyncCfb8WriteHalf<W> {
             return Poll::Ready(Ok(0));
         }
 
+        // The drain above emptied the buffer, so only fresh plaintext is encrypted here.
         let this = &mut *self;
+        debug_assert!(this.buffer.is_empty());
         this.buffer.extend_from_slice(buf);
-        this.encryptor.cipher.encrypt(&mut this.buffer);
+        this.encryptor.encrypt_in_place(&mut this.buffer);
 
         // Opportunistically write; whatever the inner writer doesn't take now
         // stays buffered and is drained on the next write/flush/shutdown.

src/packet.rs

@@ -27,7 +27,7 @@ use crate::ser::{Deserialize, SerializationError, Serialize};
 use crate::varint::{VarInt, VarIntError};
 
 #[cfg(feature = "async")]
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 
 // ---------------------------------------------------------------------------
 // Errors
@@ -74,6 +74,25 @@ fn active_threshold(threshold: Option<i32>) -> Option<usize> {
     threshold.and_then(|t| usize::try_from(t).ok())
 }
 
+fn checked_read_len(len: VarInt) -> Result<usize, PacketError> {
+    if len.0 < 0 || len.0 as usize > MAX_PACKET_LENGTH {
+        return Err(PacketError::InvalidLength(len.0 as i64));
+    }
+    Ok(len.0 as usize)
+}
+
+/// Splits `[packet_id][payload]` in place, reusing the buffer for the payload.
+fn split_packet_id(mut data: Vec<u8>, offset: usize) -> Result<UncompressedPacket, PacketError> {
+    let mut cursor = Cursor::new(&data[offset..]);
+    let packet_id = VarInt::read_sync(&mut cursor)?.0;
+    let end = offset + cursor.position() as usize;
+    data.drain(..end);
+    Ok(UncompressedPacket {
+        packet_id,
+        payload: data,
+    })
+}
+
 fn checked_frame_len(len: usize) -> Result<i32, PacketError> {
     if len > MAX_PACKET_LENGTH {
         return Err(PacketError::InvalidLength(len as i64));
@@ -118,13 +137,16 @@ impl RawPacket {
     ///
     /// Fails with [`PacketError::InvalidLength`] if the declared length is
     /// negative or exceeds [`MAX_PACKET_LENGTH`].
+    ///
+    /// The length prefix is read byte by byte, so wrap unbuffered sources such
+    /// as sockets in a `BufReader`.
     pub fn read_sync<R: Read>(reader: &mut R) -> Result<Self, PacketError> {
-        let len = VarInt::read_sync(reader)?;
-        if len.0 < 0 || len.0 as usize > MAX_PACKET_LENGTH {
-            return Err(PacketError::InvalidLength(len.0 as i64));
+        let len = checked_read_len(VarInt::read_sync(reader)?)?;
+        let mut data = Vec::with_capacity(len);
+        reader.take(len as u64).read_to_end(&mut data)?;
+        if data.len() != len {
+            return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
         }
-        let mut data = vec![0u8; len.0 as usize];
-        reader.read_exact(&mut data)?;
         Ok(RawPacket { data })
     }
 
@@ -141,20 +163,28 @@ impl RawPacket {
     ///
     /// Fails with [`PacketError::InvalidLength`] if the declared length is
     /// negative or exceeds [`MAX_PACKET_LENGTH`].
+    ///
+    /// The length prefix is read byte by byte, so wrap unbuffered sources such
+    /// as sockets in a `tokio::io::BufReader`. Not cancellation safe: a frame
+    /// interrupted midway is lost and the stream is left desynchronized.
     #[cfg(feature = "async")]
-    pub async fn read_async<R: AsyncReadExt + Unpin>(reader: &mut R) -> Result<Self, PacketError> {
-        let len = VarInt::read_async(reader).await?;
-        if len.0 < 0 || len.0 as usize > MAX_PACKET_LENGTH {
-            return Err(PacketError::InvalidLength(len.0 as i64));
+    pub async fn read_async<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Self, PacketError> {
+        use bytes::BufMut;
+
+        let len = checked_read_len(VarInt::read_async(reader).await?)?;
+        let mut data = Vec::with_capacity(len);
+        while data.len() < len {
+            let remaining = len - data.len();
+            if reader.read_buf(&mut (&mut data).limit(remaining)).await? == 0 {
+                return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
+            }
         }
-        let mut data = vec![0u8; len.0 as usize];
-        reader.read_exact(&mut data).await?;
         Ok(RawPacket { data })
     }
 
     /// Write this packet frame to an async writer; rejects frames over [`MAX_PACKET_LENGTH`].
     #[cfg(feature = "async")]
-    pub async fn write_async<W: AsyncWriteExt + Unpin>(
+    pub async fn write_async<W: AsyncWrite + Unpin>(
         &self,
         writer: &mut W,
     ) -> Result<(), PacketError> {
@@ -226,14 +256,28 @@ impl RawPacket {
         if uncompressed.len() != data_length.0 as usize {
             return Err(PacketError::InvalidLength(uncompressed.len() as i64));
         }
-        let mut inner = Cursor::new(&uncompressed);
-        let packet_id = VarInt::read_sync(&mut inner)?;
-        let pos = inner.position() as usize;
-        let payload = uncompressed[pos..].to_vec();
-        Ok(UncompressedPacket {
-            packet_id: packet_id.0,
-            payload,
-        })
+        split_packet_id(uncompressed, 0)
+    }
+
+    /// Like [`uncompress`](Self::uncompress), but consumes the frame so an
+    /// uncompressed body is reused in place instead of copied.
+    pub fn into_uncompressed(
+        self,
+        threshold: Option<i32>,
+    ) -> Result<UncompressedPacket, PacketError> {
+        #[cfg(feature = "compression")]
+        if active_threshold(threshold).is_some() {
+            let mut cursor = Cursor::new(&self.data);
+            let data_length = VarInt::read_sync(&mut cursor)?.0;
+            if data_length != 0 {
+                return self.uncompress(threshold);
+            }
+            let offset = cursor.position() as usize;
+            return split_packet_id(self.data, offset);
+        }
+        #[cfg(not(feature = "compression"))]
+        let _ = threshold;
+        split_packet_id(self.data, 0)
     }
 
     /// Interpret frame without feature-gating (no compression).
@@ -352,7 +396,7 @@ impl UncompressedPacket {
 
     /// Write this packet to an async writer (uncompressed framing).
     #[cfg(feature = "async")]
-    pub async fn write_async<W: AsyncWriteExt + Unpin>(
+    pub async fn write_async<W: AsyncWrite + Unpin>(
         &self,
         writer: &mut W,
     ) -> Result<(), PacketError> {

src/prelude/mod.rs

@@ -25,7 +25,7 @@
 //! AES-128-CFB8 types.
 
 // --- Core serialization ---
-pub use crate::ser::{Deserialize, RawBytes, SerializationError, Serialize};
+pub use crate::ser::{ByteArray, Deserialize, RawBytes, SerializationError, Serialize};
 
 // --- Variable-length integers ---
 pub use crate::varint::{VarInt, VarLong};