diff --git a/Cargo.lock b/Cargo.lock index c06cdd3..04b7d3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -322,7 +322,7 @@ dependencies = [ [[package]] name = "mc_protocol" -version = "2.4.0" +version = "2.5.0" dependencies = [ "aes", "bytes", @@ -338,7 +338,7 @@ dependencies = [ [[package]] name = "mc_protocol_derive" -version = "2.0.3" +version = "2.0.4" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index e56339f..8404ce2 100644 --- a/Cargo.toml +++ b/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 diff --git a/mc_protocol_derive/Cargo.toml b/mc_protocol_derive/Cargo.toml index 3a9120e..841a3f0 100644 --- a/mc_protocol_derive/Cargo.toml +++ b/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 diff --git a/mc_protocol_derive/src/lib.rs b/mc_protocol_derive/src/lib.rs index 1b6fcb3..1e02933 100644 --- a/mc_protocol_derive/src/lib.rs +++ b/mc_protocol_derive/src/lib.rs @@ -131,7 +131,7 @@ fn expand_packet(input: &DeriveInput) -> syn::Result { 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 { #[automatically_derived] impl #ser_impl_generics ::mc_protocol::ser::Serialize for #struct_name #ty_generics #ser_where_clause { - fn serialize( + 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 { #[automatically_derived] impl #de_impl_generics ::mc_protocol::ser::Deserialize for #struct_name #ty_generics #de_where_clause { - fn deserialize( - __reader: &mut R, + fn deserialize<__McReader: ::std::io::Read + ::core::marker::Unpin>( + __reader: &mut __McReader, ) -> ::core::result::Result { ::core::result::Result::Ok(Self { #(#deserialize_calls)* diff --git a/src/compression.rs b/src/compression.rs index fcbd441..10d6370 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -101,7 +101,7 @@ pub fn decompress_zlib(data: &[u8]) -> Result, 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, 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, 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() { diff --git a/src/encryption.rs b/src/encryption.rs index 72164d3..1b0b9bb 100644 --- a/src/encryption.rs +++ b/src/encryption.rs @@ -144,11 +144,17 @@ impl Read for Cfb8ReadHalf { } /// 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 { inner: W, encryptor: Cfb8Encryptor, - /// Scratch space for ciphertext, reused across writes. + /// Encrypted bytes not yet accepted by the inner writer. buffer: Vec, + /// Number of bytes at the front of `buffer` already written. + written: usize, } impl Cfb8WriteHalf { @@ -158,25 +164,43 @@ impl Cfb8WriteHalf { 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 Write for Cfb8WriteHalf { fn write(&mut self, buf: &[u8]) -> io::Result { - 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 AsyncWrite for AsyncCfb8WriteHalf { 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())) } diff --git a/src/lib.rs b/src/lib.rs index 5061838..7b3ca12 100644 --- a/src/lib.rs +++ b/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) diff --git a/src/packet.rs b/src/packet.rs index 4459438..77e5981 100644 --- a/src/packet.rs +++ b/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) -> Option { threshold.and_then(|t| usize::try_from(t).ok()) } -fn checked_read_len(len: VarInt) -> Result { - 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 { + 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, offset: usize) -> Result { 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(reader: &mut R) -> Result { - 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(reader: &mut R) -> Result { 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( &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) -> Result { 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) -> Result { + pub fn uncompress(&self, threshold: Option) -> Result { + 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, + threshold: Option, ) -> Result { + if active_threshold(threshold).is_some() { + return Err(compression_unavailable()); + } self.to_raw_packet() } diff --git a/src/ser.rs b/src/ser.rs index d264db8..ae684b8 100644 --- a/src/ser.rs +++ b/src/ser.rs @@ -10,6 +10,7 @@ //! then deserialize from a `std::io::Cursor`. use std::{ + cell::Cell, io::{self, Read, Write}, string::FromUtf8Error, }; @@ -193,7 +194,9 @@ pub fn serialize_string_with_max( max: max_utf16, }); } - VarInt(s.len() as i32).write_sync(writer)?; + let byte_len = + i32::try_from(s.len()).map_err(|_| SerializationError::InvalidLength(s.len() as i64))?; + VarInt(byte_len).write_sync(writer)?; writer.write_all(s.as_bytes())?; Ok(()) } @@ -306,7 +309,9 @@ impl Deserialize for Option { impl Serialize for Vec { /// Encoded as VarInt(length) followed by each element. fn serialize(&self, writer: &mut W) -> Result<(), SerializationError> { - VarInt(self.len() as i32).write_sync(writer)?; + let len = i32::try_from(self.len()) + .map_err(|_| SerializationError::InvalidLength(self.len() as i64))?; + VarInt(len).write_sync(writer)?; for item in self { item.serialize(writer)?; } @@ -320,6 +325,7 @@ impl Deserialize for Vec { if len < 0 || len as usize > crate::packet::MAX_UNCOMPRESSED_LENGTH { return Err(SerializationError::InvalidLength(len as i64)); } + let _scope = DecodeScope::enter(len as usize, std::mem::size_of::())?; // Cap the pre-allocation: the length prefix is attacker-controlled, // so don't reserve gigabytes up front. The vector still grows to the // full length as elements are actually read. @@ -333,7 +339,48 @@ impl Deserialize for Vec { } /// Upper bound on memory reserved from an untrusted length prefix before data arrives. -const MAX_PREALLOC_BYTES: usize = 64 * 1024; +pub(crate) const MAX_PREALLOC_BYTES: usize = 64 * 1024; + +/// Maximum nesting of `Vec` while decoding one value; guards against stack overflow. +pub const MAX_VEC_NESTING: u32 = 128; + +/// Maximum in-memory size of all `Vec` nested in one top-level `Vec`. +pub const MAX_DECODED_VEC_BYTES: usize = 32 * 1024 * 1024; + +thread_local! { + static DECODE_STATE: Cell<(u32, usize)> = const { Cell::new((0, 0)) }; +} + +/// Tracks `Vec` nesting and the shared memory budget of the current decode. +struct DecodeScope; + +impl DecodeScope { + fn enter(len: usize, elem_size: usize) -> Result { + let (depth, budget) = DECODE_STATE.get(); + let budget = if depth == 0 { + MAX_DECODED_VEC_BYTES + } else { + budget + }; + if depth >= MAX_VEC_NESTING { + return Err(io::Error::new(io::ErrorKind::InvalidData, "Vec nesting too deep").into()); + } + // Zero-sized elements still cost a loop iteration each. + let cost = len.saturating_mul(elem_size.max(1)); + let Some(budget) = budget.checked_sub(cost) else { + return Err(SerializationError::InvalidLength(len as i64)); + }; + DECODE_STATE.set((depth + 1, budget)); + Ok(Self) + } +} + +impl Drop for DecodeScope { + fn drop(&mut self) { + let (depth, budget) = DECODE_STATE.get(); + DECODE_STATE.set((depth - 1, budget)); + } +} // --------------------------------------------------------------------------- // Prefixed byte array diff --git a/src/varint.rs b/src/varint.rs index f47df3b..0f8e8a3 100644 --- a/src/varint.rs +++ b/src/varint.rs @@ -111,14 +111,21 @@ impl VarInt { /// Decode a `VarInt` from a synchronous reader. pub fn read_sync(reader: &mut R) -> Result { + Ok(Self::read_sync_counted(reader)?.0) + } + + /// Decode a `VarInt`, also returning how many bytes it occupied on the wire. + pub(crate) fn read_sync_counted(reader: &mut R) -> Result<(Self, usize), VarIntError> { let mut value: i32 = 0; let mut position: u32 = 0; + let mut count = 0; loop { // read_exact retries on ErrorKind::Interrupted and yields // UnexpectedEof when the stream ends mid-value. let mut byte_buf = [0u8; 1]; reader.read_exact(&mut byte_buf)?; + count += 1; let byte = byte_buf[0]; value |= ((byte & SEGMENT_BITS) as i32) << position; @@ -133,7 +140,7 @@ impl VarInt { } } - Ok(VarInt(value)) + Ok((VarInt(value), count)) } /// Encode this `VarInt` into a synchronous writer. @@ -148,11 +155,21 @@ impl VarInt { /// Decode a `VarInt` from an async reader (requires `async` feature). #[cfg(feature = "async")] pub async fn read_async(reader: &mut R) -> Result { + Ok(Self::read_async_counted(reader).await?.0) + } + + /// Async counterpart of [`read_sync_counted`](Self::read_sync_counted). + #[cfg(feature = "async")] + pub(crate) async fn read_async_counted( + reader: &mut R, + ) -> Result<(Self, usize), VarIntError> { let mut value: i32 = 0; let mut position: u32 = 0; + let mut count = 0; loop { let byte = reader.read_u8().await?; + count += 1; value |= ((byte & SEGMENT_BITS) as i32) << position; if byte & CONTINUE_BIT == 0 { @@ -165,7 +182,7 @@ impl VarInt { } } - Ok(VarInt(value)) + Ok((VarInt(value), count)) } /// Encode this `VarInt` into an async writer (requires `async` feature). diff --git a/tests/hardening.rs b/tests/hardening.rs new file mode 100644 index 0000000..ff04be7 --- /dev/null +++ b/tests/hardening.rs @@ -0,0 +1,235 @@ +//! Regression tests for hostile input and I/O failure handling. + +use mc_protocol::Packet; +use mc_protocol::packet::{PacketError, RawPacket, UncompressedPacket}; +use mc_protocol::ser::{Deserialize, MAX_VEC_NESTING, SerializationError}; +use mc_protocol::varint::VarInt; +use std::io::{self, Cursor}; + +#[derive(Packet, Debug)] +struct Node { + children: Vec, +} + +#[derive(Packet, Debug)] +struct Sparse { + items: Vec>, +} + +#[derive(Packet, Debug, PartialEq)] +struct Generic { + a: W, + b: R, +} + +const BASE: u8 = 0x10; + +#[derive(Packet, Debug)] +#[packet(BASE + 1)] +struct ComputedId { + x: u8, +} + +fn nested(depth: usize) -> Vec { + let mut wire = vec![1u8; depth]; + wire.push(0); + wire +} + +#[test] +fn hostile_nesting_is_rejected_on_a_small_stack() { + let handle = std::thread::Builder::new() + .stack_size(256 * 1024) + .spawn(|| { + let err = Node::deserialize(&mut Cursor::new(nested(2_000_000))).unwrap_err(); + assert!( + matches!(err, SerializationError::Io(e) if e.kind() == io::ErrorKind::InvalidData) + ); + }) + .unwrap(); + handle.join().unwrap(); +} + +#[test] +fn nesting_limit_and_state_reset() { + let limit = MAX_VEC_NESTING as usize; + Node::deserialize(&mut Cursor::new(nested(limit - 1))).unwrap(); + Node::deserialize(&mut Cursor::new(nested(limit))).unwrap_err(); + Node::deserialize(&mut Cursor::new(nested(limit - 1))).unwrap(); +} + +#[test] +fn memory_amplification_is_rejected_before_decoding() { + let mut wire = Vec::new(); + VarInt(8_000_000).write_sync(&mut wire).unwrap(); + wire.resize(wire.len() + 8_000_000, 0); + let err = Sparse::deserialize(&mut Cursor::new(wire)).unwrap_err(); + assert!(matches!(err, SerializationError::InvalidLength(8_000_000))); +} + +#[test] +fn budget_is_shared_by_nested_vecs() { + let mut wire = Vec::new(); + VarInt(16).write_sync(&mut wire).unwrap(); + for _ in 0..16 { + VarInt(200_000).write_sync(&mut wire).unwrap(); + wire.resize(wire.len() + 200_000, 0); + } + assert!(Vec::>>::deserialize(&mut Cursor::new(wire)).is_err()); +} + +#[test] +fn frame_prefix_wider_than_three_bytes_is_rejected() { + let frame = [0x81u8, 0x80, 0x80, 0x00, 0x05]; + let err = RawPacket::read_sync(&mut Cursor::new(&frame)).unwrap_err(); + assert!(matches!(err, PacketError::InvalidLength(1))); + + let canonical = [0x01u8, 0x05]; + assert_eq!( + RawPacket::read_sync(&mut Cursor::new(&canonical)) + .unwrap() + .data, + [5] + ); +} + +#[test] +fn truncated_large_frame_fails_with_eof() { + let err = RawPacket::read_sync(&mut Cursor::new(&[0xffu8, 0xff, 0x7f, 1, 2])).unwrap_err(); + assert!(matches!(err, PacketError::Io(e) if e.kind() == io::ErrorKind::UnexpectedEof)); +} + +#[test] +fn derive_accepts_w_and_r_generics() { + let value = Generic { + a: 7u8, + b: VarInt(300), + }; + let mut wire = Vec::new(); + mc_protocol::ser::Serialize::serialize(&value, &mut wire).unwrap(); + assert_eq!( + Generic::::deserialize(&mut Cursor::new(wire)).unwrap(), + value + ); +} + +#[test] +fn derive_parenthesizes_id_expression() { + assert_eq!(ComputedId::PACKET_ID, 0x11); + assert_eq!(ComputedId { x: 0 }.x, 0); +} + +#[cfg(not(feature = "compression"))] +#[test] +fn active_threshold_errors_without_compression_feature() { + let packet = UncompressedPacket::new(1, vec![2]); + assert!(packet.to_raw_packet_compressed(Some(256)).is_err()); + assert!(packet.to_raw_packet_compressed(None).is_ok()); + let raw = RawPacket::new(vec![0, 1, 2]); + assert!(raw.uncompress(Some(256)).is_err()); + assert!(raw.clone().into_uncompressed(Some(256)).is_err()); + assert_eq!(raw.uncompress(Some(-1)).unwrap().packet_id, 0); +} + +#[cfg(feature = "compression")] +mod compressed { + use super::*; + use mc_protocol::compression::{compress_zlib, decompress_zlib_limited}; + + fn frame(data_length: i32, body: &[u8]) -> RawPacket { + let mut data = Vec::new(); + VarInt(data_length).write_sync(&mut data).unwrap(); + data.extend_from_slice(body); + RawPacket::new(data) + } + + #[test] + fn declared_length_below_threshold_is_rejected() { + let raw = frame(3, &compress_zlib(&[1, 2, 3]).unwrap()); + assert!(matches!( + raw.uncompress(Some(256)), + Err(PacketError::InvalidLength(3)) + )); + assert!(matches!( + raw.clone().into_uncompressed(Some(256)), + Err(PacketError::InvalidLength(3)) + )); + assert_eq!(raw.uncompress(Some(3)).unwrap().payload, [2, 3]); + } + + #[test] + fn trailing_bytes_after_zlib_stream_are_rejected() { + let mut body = compress_zlib(&[1, 2, 3]).unwrap(); + decompress_zlib_limited(&body, 3).unwrap(); + body.push(0); + decompress_zlib_limited(&body, 3).unwrap_err(); + } + + #[test] + fn realistic_compressed_packet_still_round_trips() { + let packet = UncompressedPacket::new(9, vec![0; 4096]); + let raw = packet.to_raw_packet_compressed(Some(256)).unwrap(); + assert_eq!(raw.uncompress(Some(256)).unwrap().payload.len(), 4096); + } +} + +#[cfg(feature = "encryption")] +mod encrypted { + use super::*; + use mc_protocol::encryption::{Cfb8Decryptor, Cfb8WriteHalf}; + use std::io::Write; + + const KEY: [u8; 16] = [7; 16]; + + #[derive(Default)] + struct Flaky { + out: Vec, + failures: usize, + } + + impl Write for Flaky { + fn write(&mut self, data: &[u8]) -> io::Result { + if self.failures > 0 { + self.failures -= 1; + return Err(io::ErrorKind::WouldBlock.into()); + } + let n = data.len().min(3); + self.out.extend_from_slice(&data[..n]); + Ok(n) + } + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + fn plaintext(out: &[u8]) -> Vec { + Cfb8Decryptor::new(&KEY).unwrap().decrypt(out).unwrap() + } + + #[test] + fn would_block_keeps_cipher_and_stream_in_sync() { + let inner = Flaky { + failures: 1, + ..Default::default() + }; + let mut writer = Cfb8WriteHalf::new(inner, &KEY).unwrap(); + writer.write_all(b"hello ").unwrap(); + writer.write_all(b"world").unwrap(); + writer.flush().unwrap(); + assert_eq!(plaintext(&writer.into_inner().out), b"hello world"); + } + + #[test] + fn failed_drain_consumes_nothing_on_retry() { + let inner = Flaky { + failures: 2, + ..Default::default() + }; + let mut writer = Cfb8WriteHalf::new(inner, &KEY).unwrap(); + assert_eq!(writer.write(b"one").unwrap(), 3); + assert!(writer.write(b"two").is_err()); + assert_eq!(writer.write(b"two").unwrap(), 3); + writer.flush().unwrap(); + assert_eq!(plaintext(&writer.into_inner().out), b"onetwo"); + } +}