fix: harden decoding and encrypted writes (2.5.0)

Cargo.toml

@@ -10,7 +10,7 @@ license = "MIT"
 repository = "https://github.com/kauri-off/mc_protocol"
 
 [workspace.dependencies]
-mc_protocol_derive = { path = "./mc_protocol_derive", version = "2.0.3" }
+mc_protocol_derive = { path = "./mc_protocol_derive", version = "2.0.4" }
 
 [workspace.lints.rust]
 unsafe_code = "forbid"
@@ -26,7 +26,7 @@ print_stderr = "warn"
 
 [package]
 name = "mc_protocol"
-version = "2.4.0"
+version = "2.5.0"
 edition.workspace = true
 rust-version.workspace = true
 authors.workspace = true

mc_protocol_derive/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mc_protocol_derive"
-version = "2.0.3"
+version = "2.0.4"
 edition.workspace = true
 rust-version.workspace = true
 authors.workspace = true

mc_protocol_derive/src/lib.rs

@@ -131,7 +131,7 @@ fn expand_packet(input: &DeriveInput) -> syn::Result<TokenStream2> {
         quote! {
             impl #impl_generics #struct_name #ty_generics #where_clause {
                 /// The numeric ID that identifies this packet on the wire.
-                pub const PACKET_ID: i32 = #expr as i32;
+                pub const PACKET_ID: i32 = (#expr) as i32;
             }
 
             #[automatically_derived]
@@ -148,9 +148,9 @@ fn expand_packet(input: &DeriveInput) -> syn::Result<TokenStream2> {
 
         #[automatically_derived]
         impl #ser_impl_generics ::mc_protocol::ser::Serialize for #struct_name #ty_generics #ser_where_clause {
-            fn serialize<W: ::std::io::Write + ::core::marker::Unpin>(
+            fn serialize<__McWriter: ::std::io::Write + ::core::marker::Unpin>(
                 &self,
-                __writer: &mut W,
+                __writer: &mut __McWriter,
             ) -> ::core::result::Result<(), ::mc_protocol::ser::SerializationError> {
                 #(#serialize_calls)*
                 ::core::result::Result::Ok(())
@@ -159,8 +159,8 @@ fn expand_packet(input: &DeriveInput) -> syn::Result<TokenStream2> {
 
         #[automatically_derived]
         impl #de_impl_generics ::mc_protocol::ser::Deserialize for #struct_name #ty_generics #de_where_clause {
-            fn deserialize<R: ::std::io::Read + ::core::marker::Unpin>(
-                __reader: &mut R,
+            fn deserialize<__McReader: ::std::io::Read + ::core::marker::Unpin>(
+                __reader: &mut __McReader,
             ) -> ::core::result::Result<Self, ::mc_protocol::ser::SerializationError> {
                 ::core::result::Result::Ok(Self {
                     #(#deserialize_calls)*

src/compression.rs

@@ -101,7 +101,7 @@ pub fn decompress_zlib(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
 ///
 /// Returns an error of kind `InvalidData` if the decompressed output would
 /// exceed `limit` — this guards against zip bombs, where a tiny compressed
-/// input expands to a huge output.
+/// input expands to a huge output — or if bytes follow the end of the zlib stream.
 pub fn decompress_zlib_limited(data: &[u8], limit: usize) -> Result<Vec<u8>, CompressionError> {
     // Room for limit + 1 bytes: producing more than `limit` means the input lied.
     let max_len = limit.saturating_add(1);
@@ -123,6 +123,9 @@ pub fn decompress_zlib_limited(data: &[u8], limit: usize) -> Result<Vec<u8>, Com
                 )));
             }
             if status == Status::StreamEnd {
+                if inflate.total_in() != data.len() as u64 {
+                    return Err(invalid_data("trailing data after zlib stream"));
+                }
                 return Ok(output);
             }
             if output.len() == output.capacity() {

src/encryption.rs

@@ -144,11 +144,17 @@ impl<R: Read> Read for Cfb8ReadHalf<R> {
 }
 
 /// A synchronous write-half that transparently encrypts outgoing bytes.
+///
+/// Input is encrypted exactly once; ciphertext the inner writer refuses stays
+/// buffered, and its error surfaces on the next `write` or `flush`. Call
+/// `flush` to confirm delivery on non-blocking or timed-out writers.
 pub struct Cfb8WriteHalf<W> {
     inner: W,
     encryptor: Cfb8Encryptor,
-    /// Scratch space for ciphertext, reused across writes.
+    /// Encrypted bytes not yet accepted by the inner writer.
     buffer: Vec<u8>,
+    /// Number of bytes at the front of `buffer` already written.
+    written: usize,
 }
 
 impl<W: Write> Cfb8WriteHalf<W> {
@@ -158,25 +164,43 @@ impl<W: Write> Cfb8WriteHalf<W> {
             inner,
             encryptor: Cfb8Encryptor::new(key)?,
             buffer: Vec::new(),
+            written: 0,
         })
     }
 
-    /// Unwrap, discarding the encryptor state.
+    /// Unwrap, discarding the encryptor state and any unflushed ciphertext.
     pub fn into_inner(self) -> W {
         self.inner
     }
+
+    fn drain_buffer(&mut self) -> io::Result<()> {
+        while self.written < self.buffer.len() {
+            match self.inner.write(&self.buffer[self.written..]) {
+                Ok(0) => return Err(io::ErrorKind::WriteZero.into()),
+                Ok(n) => self.written += n,
+                Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
+                Err(e) => return Err(e),
+            }
+        }
+        self.buffer.clear();
+        self.written = 0;
+        Ok(())
+    }
 }
 
 impl<W: Write> Write for Cfb8WriteHalf<W> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        self.buffer.clear();
+        // An error here consumes nothing, so the caller may safely retry.
+        self.drain_buffer()?;
         self.buffer.extend_from_slice(buf);
         self.encryptor.encrypt_in_place(&mut self.buffer);
-        self.inner.write_all(&self.buffer)?;
+        // `buf` is consumed from here on; a failure stays buffered and resurfaces later.
+        let _ = self.drain_buffer();
         Ok(buf.len())
     }
 
     fn flush(&mut self) -> io::Result<()> {
+        self.drain_buffer()?;
         self.inner.flush()
     }
 }
@@ -320,10 +344,8 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for AsyncCfb8WriteHalf<W> {
         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.
-        if let Poll::Ready(Err(e)) = self.poll_drain_buffer(cx) {
-            return Poll::Ready(Err(e));
-        }
+        // stays buffered, and any error resurfaces on the next write/flush/shutdown.
+        let _ = self.poll_drain_buffer(cx);
         Poll::Ready(Ok(buf.len()))
     }
 

src/lib.rs

@@ -161,6 +161,7 @@
 //! larger packets are zlib-compressed.
 //!
 //! ```rust
+//! # #[cfg(feature = "compression")] fn main() {
 //! use mc_protocol::packet::UncompressedPacket;
 //!
 //! let up = UncompressedPacket::new(0x26, vec![0u8; 512]);
@@ -169,6 +170,8 @@
 //! let raw = up.to_raw_packet_compressed(threshold).unwrap();
 //! let decoded = raw.uncompress(threshold).unwrap();
 //! assert_eq!(decoded.packet_id, 0x26);
+//! # }
+//! # #[cfg(not(feature = "compression"))] fn main() {}
 //! ```
 //!
 //! ## Encryption (requires `encryption` feature)

src/packet.rs

@@ -23,7 +23,7 @@
 use std::io::{self, Cursor, Read, Write};
 use thiserror::Error;
 
-use crate::ser::{Deserialize, SerializationError, Serialize};
+use crate::ser::{Deserialize, MAX_PREALLOC_BYTES, SerializationError, Serialize};
 use crate::varint::{VarInt, VarIntError};
 
 #[cfg(feature = "async")]
@@ -69,18 +69,29 @@ pub const MAX_PACKET_LENGTH: usize = 2_097_151;
 /// Maximum `data_length` of a compressed packet (vanilla limit, 2^23 bytes).
 pub const MAX_UNCOMPRESSED_LENGTH: usize = 8_388_608;
 
-#[cfg(feature = "compression")]
 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 {
+/// Maximum width of the frame length prefix on the wire.
+const MAX_FRAME_PREFIX_BYTES: usize = 3;
+
+fn checked_read_len((len, width): (VarInt, usize)) -> Result<usize, PacketError> {
+    if len.0 < 0 || len.0 as usize > MAX_PACKET_LENGTH || width > MAX_FRAME_PREFIX_BYTES {
         return Err(PacketError::InvalidLength(len.0 as i64));
     }
     Ok(len.0 as usize)
 }
 
+#[cfg(not(feature = "compression"))]
+fn compression_unavailable() -> PacketError {
+    io::Error::new(
+        io::ErrorKind::Unsupported,
+        "compression threshold is active but the `compression` feature is disabled",
+    )
+    .into()
+}
+
 /// 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..]);
@@ -136,13 +147,13 @@ impl RawPacket {
     /// Read a length-prefixed frame from a synchronous reader.
     ///
     /// Fails with [`PacketError::InvalidLength`] if the declared length is
-    /// negative or exceeds [`MAX_PACKET_LENGTH`].
+    /// negative, exceeds [`MAX_PACKET_LENGTH`], or its prefix is wider than 3 bytes.
     ///
     /// 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 = checked_read_len(VarInt::read_sync(reader)?)?;
-        let mut data = Vec::with_capacity(len);
+        let len = checked_read_len(VarInt::read_sync_counted(reader)?)?;
+        let mut data = Vec::with_capacity(len.min(MAX_PREALLOC_BYTES));
         reader.take(len as u64).read_to_end(&mut data)?;
         if data.len() != len {
             return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
@@ -162,7 +173,7 @@ impl RawPacket {
     /// Read a length-prefixed frame from an async reader (requires `async` feature).
     ///
     /// Fails with [`PacketError::InvalidLength`] if the declared length is
-    /// negative or exceeds [`MAX_PACKET_LENGTH`].
+    /// negative, exceeds [`MAX_PACKET_LENGTH`], or its prefix is wider than 3 bytes.
     ///
     /// 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
@@ -171,8 +182,8 @@ impl RawPacket {
     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);
+        let len = checked_read_len(VarInt::read_async_counted(reader).await?)?;
+        let mut data = Vec::with_capacity(len.min(MAX_PREALLOC_BYTES));
         while data.len() < len {
             let remaining = len - data.len();
             if reader.read_buf(&mut (&mut data).limit(remaining)).await? == 0 {
@@ -183,6 +194,8 @@ impl RawPacket {
     }
 
     /// Write this packet frame to an async writer; rejects frames over [`MAX_PACKET_LENGTH`].
+    ///
+    /// Does not flush; buffering writers need `flush().await` afterwards.
     #[cfg(feature = "async")]
     pub async fn write_async<W: AsyncWrite + Unpin>(
         &self,
@@ -219,16 +232,17 @@ impl RawPacket {
     /// first field is `data_length` (VarInt): if 0, the rest is uncompressed;
     /// otherwise it is the uncompressed length and the rest is zlib-deflated.
     ///
-    /// Returns `Err` if decompression fails, if `data_length` is negative or
-    /// exceeds [`MAX_UNCOMPRESSED_LENGTH`], or if the decompressed size does not
-    /// match the declared `data_length` (zip-bomb guard).
+    /// Returns `Err` if decompression fails, if a non-zero `data_length` is
+    /// below `threshold`, if `data_length` is negative or exceeds
+    /// [`MAX_UNCOMPRESSED_LENGTH`], or if the decompressed size does not match
+    /// the declared `data_length` (zip-bomb guard).
     #[cfg(feature = "compression")]
     pub fn uncompress(&self, threshold: Option<i32>) -> Result<UncompressedPacket, PacketError> {
         use crate::compression::decompress_zlib_limited;
 
-        if active_threshold(threshold).is_none() {
+        let Some(active) = active_threshold(threshold) else {
             return self.as_uncompressed();
-        }
+        };
 
         let mut cursor = Cursor::new(&self.data);
         let data_length = VarInt::read_sync(&mut cursor)?;
@@ -249,6 +263,10 @@ impl RawPacket {
             });
         }
 
+        if (data_length.0 as usize) < active {
+            return Err(PacketError::InvalidLength(data_length.0 as i64));
+        }
+
         // Compressed payload — never inflate past the declared size, and
         // require the declared size to be exact.
         let compressed_data = &self.data[compressed_start..];
@@ -276,13 +294,18 @@ impl RawPacket {
             return split_packet_id(self.data, offset);
         }
         #[cfg(not(feature = "compression"))]
-        let _ = threshold;
+        if active_threshold(threshold).is_some() {
+            return Err(compression_unavailable());
+        }
         split_packet_id(self.data, 0)
     }
 
-    /// Interpret frame without feature-gating (no compression).
+    /// Interpret an uncompressed frame; fails if a compression threshold is active.
     #[cfg(not(feature = "compression"))]
-    pub fn uncompress(&self, _threshold: Option<i32>) -> Result<UncompressedPacket, PacketError> {
+    pub fn uncompress(&self, threshold: Option<i32>) -> Result<UncompressedPacket, PacketError> {
+        if active_threshold(threshold).is_some() {
+            return Err(compression_unavailable());
+        }
         self.as_uncompressed()
     }
 }
@@ -370,11 +393,14 @@ impl UncompressedPacket {
     }
 
     #[cfg(not(feature = "compression"))]
-    /// Encode to raw packet (compression stub — feature not enabled).
+    /// Encode an uncompressed frame; fails if a compression threshold is active.
     pub fn to_raw_packet_compressed(
         &self,
-        _threshold: Option<i32>,
+        threshold: Option<i32>,
     ) -> Result<RawPacket, PacketError> {
+        if active_threshold(threshold).is_some() {
+            return Err(compression_unavailable());
+        }
         self.to_raw_packet()
     }