deepslate 0.3.1

//! Connection handling: framing, encryption, compression, and the bidirectional
//! packet relay pipeline.

pub mod backend;
pub mod client;
pub mod forwarding;

use std::io;
use std::time::Duration;

use aes::Aes128;
use bytes::{Buf, BytesMut};
use cfb8::cipher::generic_array::GenericArray;
use cfb8::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use deepslate_protocol::codec;
use deepslate_protocol::packet::Packet;
use deepslate_protocol::packet::login::SetCompressionPacket;
use deepslate_protocol::types::ProtocolError;
use deepslate_protocol::varint;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufWriter};
use tokio::net::TcpStream;

use crate::auth::AuthError;

/// Maximum allowed read buffer size (2.5 MiB). If unprocessed data in the read
/// buffer exceeds this limit, the connection is terminated to prevent unbounded
/// memory growth from a fast or malicious sender.
///
/// Set to `MAX_FRAME_SIZE` (2 MiB, defined in `deepslate_protocol::codec`) plus
/// 512 KiB of headroom for partial next-frame data. This is tight enough that a
/// single incomplete max-size frame triggers the limit before the buffer can grow
/// much further.
const MAX_READ_BUF_SIZE: usize = 2 * 1024 * 1024 + 512 * 1024;

/// Errors that can occur during connection handling.
///
/// This replaces the lossy `ProtocolError` → `io::Error` conversion, preserving
/// error semantics so callers can classify failures (expected disconnects vs.
/// protocol bugs vs. backend failures) for appropriate log levels and metrics.
#[derive(Debug, thiserror::Error)]
pub enum ConnectionError {
    /// An I/O error from the underlying transport.
    #[error(transparent)]
    Io(#[from] io::Error),

    /// A Minecraft protocol-level error (framing, encoding, compression).
    #[error("protocol error: {0}")]
    Protocol(#[from] ProtocolError),

    /// Mojang authentication failed.
    #[error("authentication failed: {0}")]
    Auth(#[from] AuthError),

    /// The backend server rejected or dropped the connection.
    #[error("backend connection failed: {reason}")]
    BackendFailed {
        /// Human-readable description of the failure.
        reason: String,
    },

    /// A read or connection operation timed out.
    #[error("connection timed out")]
    Timeout,
}

impl ConnectionError {
    /// Returns `true` for benign disconnects that are expected during normal
    /// operation (client closed the connection, connection reset, broken pipe,
    /// unexpected EOF). These are logged at `debug` level. All other errors
    /// are considered unexpected and logged at `warn`.
    #[must_use]
    pub fn is_expected(&self) -> bool {
        match self {
            Self::Io(io_err) => matches!(
                io_err.kind(),
                io::ErrorKind::ConnectionReset
                    | io::ErrorKind::BrokenPipe
                    | io::ErrorKind::UnexpectedEof
                    | io::ErrorKind::ConnectionAborted
            ),
            Self::Protocol(_) | Self::Auth(_) | Self::BackendFailed { .. } | Self::Timeout => false,
        }
    }
}

type Aes128Cfb8Enc = cfb8::Encryptor<Aes128>;
type Aes128Cfb8Dec = cfb8::Decryptor<Aes128>;

/// AES-CFB8 cipher wrapper that supports incremental encrypt/decrypt across
/// multiple calls (the cfb8 crate's `AsyncStreamCipher::encrypt` consumes self,
/// so we use `BlockEncryptMut`/`BlockDecryptMut` directly).
struct CipherPair {
    encryptor: Aes128Cfb8Enc,
    decryptor: Aes128Cfb8Dec,
}

impl CipherPair {
    fn new(shared_secret: &[u8]) -> Self {
        let key = shared_secret.into();
        let iv = shared_secret.into();
        Self {
            encryptor: Aes128Cfb8Enc::new(key, iv),
            decryptor: Aes128Cfb8Dec::new(key, iv),
        }
    }

    /// Encrypt data in-place, one byte at a time (CFB-8 operates on single bytes).
    fn encrypt(&mut self, data: &mut [u8]) {
        let mut block = GenericArray::default();
        for byte in data.iter_mut() {
            block[0] = *byte;
            self.encryptor.encrypt_block_mut(&mut block);
            *byte = block[0];
        }
    }

    /// Decrypt data in-place, one byte at a time.
    fn decrypt(&mut self, data: &mut [u8]) {
        let mut block = GenericArray::default();
        for byte in data.iter_mut() {
            block[0] = *byte;
            self.decryptor.decrypt_block_mut(&mut block);
            *byte = block[0];
        }
    }
}

/// A Minecraft connection with framing, optional encryption, and optional compression.
pub struct MinecraftConnection<S = TcpStream>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    /// The underlying I/O stream, wrapped in a [`BufWriter`] so that multiple
    /// small `write_raw_packet` calls are coalesced into fewer `write(2)`
    /// syscalls. Callers must [`flush`](Self::flush) after a batch of writes to
    /// ensure data reaches the socket.
    stream: BufWriter<S>,
    /// Read buffer for accumulating incoming data.
    read_buf: BytesMut,
    /// AES-CFB8 cipher pair (set after encryption handshake).
    cipher: Option<CipherPair>,
    /// Compression threshold (-1 = disabled).
    compression_threshold: i32,
    /// Reusable zlib decompressor (avoids per-packet allocation).
    decompressor: libdeflater::Decompressor,
    /// Reusable zlib compressor (avoids per-packet allocation).
    compressor: libdeflater::Compressor,
    /// Reusable buffer for decompression output.
    decompress_buf: Vec<u8>,
    /// Reusable buffer for compression output.
    compress_buf: Vec<u8>,
    /// Reusable buffer for encoding typed packets (packet ID + fields).
    encode_buf: Vec<u8>,
    /// Reusable buffer for write framing output.
    write_buf: Vec<u8>,
    /// Read timeout for frame extraction.
    read_timeout: Duration,
}

impl<S: AsyncRead + AsyncWrite + Unpin> MinecraftConnection<S> {
    /// Wrap a raw stream into a Minecraft connection.
    #[must_use]
    pub fn new(
        stream: S,
        compression_level: libdeflater::CompressionLvl,
        read_timeout: Duration,
    ) -> Self {
        Self {
            stream: BufWriter::new(stream),
            read_buf: BytesMut::with_capacity(32768),
            cipher: None,
            compression_threshold: -1,
            decompressor: libdeflater::Decompressor::new(),
            compressor: libdeflater::Compressor::new(compression_level),
            decompress_buf: Vec::new(),
            compress_buf: Vec::new(),
            encode_buf: Vec::new(),
            write_buf: Vec::new(),
            read_timeout,
        }
    }

    /// Enable AES-CFB8 encryption using the shared secret as both key and IV.
    pub fn enable_encryption(&mut self, shared_secret: &[u8]) {
        self.cipher = Some(CipherPair::new(shared_secret));
    }

    /// Enable compression at the given threshold.
    pub const fn enable_compression(&mut self, threshold: i32) {
        self.compression_threshold = threshold;
    }

    /// Read a single frame from the connection.
    ///
    /// Handles decryption (if enabled) and frame extraction. Returns the raw
    /// frame data (packet ID + payload), or `None` if the connection is closed.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`] or [`ConnectionError::Protocol`].
    #[allow(clippy::large_stack_arrays)]
    pub async fn read_frame(&mut self) -> Result<Option<BytesMut>, ConnectionError> {
        loop {
            // Try to extract a frame from the buffer
            if let Some((varint_size, frame_len)) = codec::try_read_frame(&self.read_buf)? {
                // Advance past the length prefix, then split off the frame body.
                // split_to() is O(1) — it shares the underlying allocation.
                self.read_buf.advance(varint_size);
                let mut frame = self.read_buf.split_to(frame_len);

                if self.compression_threshold >= 0 {
                    let (uncompressed_size, payload) = codec::read_compressed_frame(&frame)?;
                    if uncompressed_size == 0 {
                        // A VarInt of value 0 is always exactly 1 byte (0x00).
                        // Advance past it and return the frame directly — zero
                        // copy, since frame already owns the data via split_to().
                        frame.advance(1);
                        return Ok(Some(frame));
                    }
                    // Reuse the decompressor and buffer across packets to avoid
                    // per-packet allocations. The buffer capacity stabilizes
                    // after a few packets.
                    self.decompress_buf.resize(uncompressed_size, 0);
                    self.decompressor
                        .zlib_decompress(payload, &mut self.decompress_buf)
                        .map_err(ProtocolError::from)?;
                    return Ok(Some(BytesMut::from(&self.decompress_buf[..])));
                }

                return Ok(Some(frame));
            }

            // Read more data from the network
            if let Some(cipher) = &mut self.cipher {
                // Must use a temp buffer for in-place decryption
                let mut tmp = [0u8; 32768];
                let n = self.stream.read(&mut tmp).await?;
                if n == 0 {
                    return Ok(None);
                }
                cipher.decrypt(&mut tmp[..n]);
                self.read_buf.extend_from_slice(&tmp[..n]);
            } else {
                // Zero-copy: read directly into BytesMut
                let n = self.stream.read_buf(&mut self.read_buf).await?;
                if n == 0 {
                    return Ok(None);
                }
            }

            // Guard against unbounded buffer growth. If the buffer has grown
            // past the limit without yielding a complete frame, the peer is
            // flooding us faster than we can consume.
            if self.read_buf.len() > MAX_READ_BUF_SIZE {
                return Err(ConnectionError::Protocol(
                    ProtocolError::ReadBufferOverflow {
                        size: self.read_buf.len(),
                        max: MAX_READ_BUF_SIZE,
                    },
                ));
            }
        }
    }

    /// Read a single frame from the connection with a timeout.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Timeout`] if the timeout is reached.
    #[expect(
        clippy::large_futures,
        reason = "MinecraftConnection carries large framing buffers through read timeouts"
    )]
    pub async fn read_frame_timeout(&mut self) -> Result<Option<BytesMut>, ConnectionError> {
        tokio::time::timeout(self.read_timeout, self.read_frame())
            .await
            .map_err(|_| ConnectionError::Timeout)?
    }

    /// Try to extract a frame that is already fully buffered in memory.
    ///
    /// Unlike [`read_frame`](Self::read_frame), this never performs I/O — it
    /// only inspects the internal read buffer. Returns `Ok(None)` when there is
    /// not enough data buffered for a complete frame.
    ///
    /// This is used by the relay loops to drain all buffered frames after a
    /// `select!` branch fires, writing them all into the [`BufWriter`] before a
    /// single `flush()` — coalescing many small packets into fewer syscalls.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Protocol`] on malformed frame data.
    pub fn try_read_frame(&mut self) -> Result<Option<BytesMut>, ConnectionError> {
        let Some((varint_size, frame_len)) = codec::try_read_frame(&self.read_buf)? else {
            return Ok(None);
        };

        self.read_buf.advance(varint_size);
        let mut frame = self.read_buf.split_to(frame_len);

        if self.compression_threshold >= 0 {
            let (uncompressed_size, payload) = codec::read_compressed_frame(&frame)?;
            if uncompressed_size == 0 {
                frame.advance(1);
                return Ok(Some(frame));
            }
            self.decompress_buf.resize(uncompressed_size, 0);
            self.decompressor
                .zlib_decompress(payload, &mut self.decompress_buf)
                .map_err(ProtocolError::from)?;
            return Ok(Some(BytesMut::from(&self.decompress_buf[..])));
        }

        Ok(Some(frame))
    }

    /// Flush the internal [`BufWriter`], ensuring all buffered writes reach the
    /// underlying socket.
    ///
    /// Call this after writing a batch of packets (e.g. after draining all
    /// buffered frames in a relay loop iteration) to push data to the kernel.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`].
    pub async fn flush(&mut self) -> Result<(), ConnectionError> {
        self.stream.flush().await?;
        Ok(())
    }

    /// Write a typed packet to the connection.
    ///
    /// This flushes the underlying [`BufWriter`] after writing, making it safe
    /// for conversational request/response flows (handshake, login, status).
    /// For high-throughput relay paths, use [`write_raw_packet`](Self::write_raw_packet)
    /// with an explicit [`flush`](Self::flush) after a batch of writes.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`].
    #[allow(clippy::future_not_send)]
    pub async fn write_packet<P: Packet>(&mut self, packet: &P) -> Result<(), ConnectionError> {
        codec::encode_packet_data(&mut self.encode_buf, P::PACKET_ID, |buf| packet.encode(buf));
        // Temporarily move the buffer out so we can pass it to
        // write_raw_packet without an overlapping &mut self borrow.
        // std::mem::take is O(1) (swaps three pointer-sized fields).
        let encode_buf = std::mem::take(&mut self.encode_buf);
        let result = self.write_raw_packet(&encode_buf).await;
        self.encode_buf = encode_buf;
        result?;
        self.stream.flush().await?;
        Ok(())
    }

    /// Encode a packet from a raw packet ID and encode closure, write it, and
    /// flush.
    ///
    /// This is a convenience wrapper around [`encode_packet_data`](codec::encode_packet_data)
    /// and [`write_raw_packet`](Self::write_raw_packet) that uses the
    /// connection's reusable encode buffer. Useful for call sites that build
    /// packets from dynamic packet IDs rather than typed [`Packet`] structs.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`] or [`ConnectionError::Protocol`].
    #[allow(clippy::future_not_send)]
    pub async fn encode_and_write_packet(
        &mut self,
        packet_id: i32,
        encode_fn: impl FnOnce(&mut Vec<u8>),
    ) -> Result<(), ConnectionError> {
        codec::encode_packet_data(&mut self.encode_buf, packet_id, encode_fn);
        let encode_buf = std::mem::take(&mut self.encode_buf);
        let result = self.write_raw_packet(&encode_buf).await;
        self.encode_buf = encode_buf;
        result?;
        self.stream.flush().await?;
        Ok(())
    }

    /// Write raw packet data (already has packet ID + payload) to the connection.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`] or [`ConnectionError::Protocol`].
    #[allow(
        clippy::cast_sign_loss,
        clippy::cast_possible_truncation,
        clippy::cast_possible_wrap
    )]
    pub async fn write_raw_packet(&mut self, packet_data: &[u8]) -> Result<(), ConnectionError> {
        self.write_buf.clear();

        if self.compression_threshold >= 0 {
            if packet_data.len() >= self.compression_threshold as usize {
                // Reuse the compressor and buffer across packets to avoid
                // per-packet allocations. The buffer capacity stabilizes
                // after a few packets.
                let max_size = self.compressor.zlib_compress_bound(packet_data.len());
                self.compress_buf.resize(max_size, 0);
                let actual_size = self
                    .compressor
                    .zlib_compress(packet_data, &mut self.compress_buf)
                    .map_err(ProtocolError::from)?;
                codec::write_compressed_frame(
                    &mut self.write_buf,
                    packet_data.len() as i32,
                    &self.compress_buf[..actual_size],
                );
            } else {
                codec::write_compressed_frame(&mut self.write_buf, 0, packet_data);
            }
        } else {
            codec::write_frame(&mut self.write_buf, packet_data);
        }

        // Encrypt the frame if encryption is enabled
        if let Some(cipher) = &mut self.cipher {
            cipher.encrypt(&mut self.write_buf);
        }

        self.stream.write_all(&self.write_buf).await?;
        Ok(())
    }

    /// Send a `SetCompression` packet and enable compression on this connection.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`].
    pub async fn set_compression(&mut self, threshold: i32) -> Result<(), ConnectionError> {
        self.write_packet(&SetCompressionPacket { threshold })
            .await?;
        self.compression_threshold = threshold;
        Ok(())
    }

    /// Gracefully shut down the connection.
    ///
    /// Flushes any buffered data, sends a TCP FIN, and drains the receive
    /// buffer so that unread client data does not cause the kernel to send a
    /// TCP RST when the socket is dropped. Without the drain, the kernel's
    /// `tcp_close` sees unread data and resets the connection, which causes
    /// the peer to discard any packets we just sent (e.g. a disconnect
    /// message).
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`] from flush or shutdown. Errors during
    /// the drain phase are silently ignored (the disconnect packet has already
    /// been sent).
    pub async fn shutdown(&mut self) -> Result<(), ConnectionError> {
        self.stream.flush().await?;
        self.stream.shutdown().await?;

        // Drain unread data from the receive buffer to prevent TCP RST.
        // After shutdown(Write), the client will eventually see our FIN and
        // stop sending. We read until EOF or a timeout, whichever comes first.
        let drain = async {
            let mut buf = [0u8; 1024];
            loop {
                match self.stream.read(&mut buf).await {
                    Ok(0) | Err(_) => break,
                    Ok(_) => {}
                }
            }
        };
        let _ = tokio::time::timeout(std::time::Duration::from_secs(5), drain).await;

        Ok(())
    }
}

impl MinecraftConnection<TcpStream> {
    /// Split into owned read/write halves for bidirectional relay.
    ///
    /// This is only available for `TcpStream`-backed connections because the
    /// resulting [`MinecraftReader`] and [`MinecraftWriter`] use the lock-free
    /// `OwnedReadHalf`/`OwnedWriteHalf` types from `TcpStream::into_split()`.
    ///
    /// # Panics
    ///
    /// Debug-asserts that the internal [`BufWriter`] has no buffered data. The
    /// caller must [`flush`](Self::flush) before splitting.
    pub fn into_split(self) -> (MinecraftReader, MinecraftWriter) {
        let inner = self.stream.into_inner();
        let (read_half, write_half) = inner.into_split();
        // If encryption is enabled, split the cipher state for each half.
        // This is safe because CFB-8 encrypt and decrypt operate on independent state.
        let (enc_cipher, dec_cipher) = if let Some(cipher) = self.cipher {
            (
                Some(EncryptCipher(cipher.encryptor)),
                Some(DecryptCipher(cipher.decryptor)),
            )
        } else {
            (None, None)
        };
        (
            MinecraftReader {
                stream: read_half,
                read_buf: self.read_buf,
                cipher: dec_cipher,
                read_timeout: self.read_timeout,
            },
            MinecraftWriter {
                stream: write_half,
                cipher: enc_cipher,
            },
        )
    }
}

/// Encryption-only cipher wrapper for the write half.
struct EncryptCipher(Aes128Cfb8Enc);

impl EncryptCipher {
    fn encrypt(&mut self, data: &mut [u8]) {
        let mut block = GenericArray::default();
        for byte in data.iter_mut() {
            block[0] = *byte;
            self.0.encrypt_block_mut(&mut block);
            *byte = block[0];
        }
    }
}

/// Decryption-only cipher wrapper for the read half.
struct DecryptCipher(Aes128Cfb8Dec);

impl DecryptCipher {
    fn decrypt(&mut self, data: &mut [u8]) {
        let mut block = GenericArray::default();
        for byte in data.iter_mut() {
            block[0] = *byte;
            self.0.decrypt_block_mut(&mut block);
            *byte = block[0];
        }
    }
}

/// Read half of a split Minecraft connection.
pub struct MinecraftReader {
    stream: tokio::net::tcp::OwnedReadHalf,
    read_buf: BytesMut,
    cipher: Option<DecryptCipher>,
    read_timeout: Duration,
}

impl MinecraftReader {
    /// Read a single raw frame (re-framed: length prefix + inner data).
    ///
    /// Returns the complete wire bytes for one packet, ready to be forwarded.
    /// Returns `None` on connection close.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`] or [`ConnectionError::Protocol`].
    #[allow(
        clippy::cast_possible_truncation,
        clippy::cast_possible_wrap,
        clippy::large_stack_arrays
    )]
    pub async fn read_raw_frame(&mut self) -> Result<Option<Vec<u8>>, ConnectionError> {
        loop {
            if let Some((varint_size, frame_len)) = codec::try_read_frame(&self.read_buf)? {
                // Re-frame: write a fresh length prefix + the frame body.
                let mut out = Vec::with_capacity(varint_size + frame_len);
                varint::write_var_int(&mut out, frame_len as i32);
                // Advance past the length prefix, then split off the frame
                // body using O(1) split_to instead of copying from Bytes.
                self.read_buf.advance(varint_size);
                let frame = self.read_buf.split_to(frame_len);
                out.extend_from_slice(&frame);
                return Ok(Some(out));
            }

            if let Some(cipher) = &mut self.cipher {
                let mut tmp = [0u8; 32768];
                let n = self.stream.read(&mut tmp).await?;
                if n == 0 {
                    return Ok(None);
                }
                cipher.decrypt(&mut tmp[..n]);
                self.read_buf.extend_from_slice(&tmp[..n]);
            } else {
                let n = self.stream.read_buf(&mut self.read_buf).await?;
                if n == 0 {
                    return Ok(None);
                }
            }

            // Guard against unbounded buffer growth. If the buffer has grown
            // past the limit without yielding a complete frame, the peer is
            // flooding us faster than we can consume.
            if self.read_buf.len() > MAX_READ_BUF_SIZE {
                return Err(ConnectionError::Protocol(
                    ProtocolError::ReadBufferOverflow {
                        size: self.read_buf.len(),
                        max: MAX_READ_BUF_SIZE,
                    },
                ));
            }
        }
    }

    /// Read a single raw frame from the connection with a timeout.
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Timeout`] if the timeout is reached.
    #[expect(
        clippy::large_futures,
        reason = "MinecraftReader carries large framing buffers through read timeouts"
    )]
    pub async fn read_raw_frame_timeout(&mut self) -> Result<Option<Vec<u8>>, ConnectionError> {
        tokio::time::timeout(self.read_timeout, self.read_raw_frame())
            .await
            .map_err(|_| ConnectionError::Timeout)?
    }
}

/// Write half of a split Minecraft connection.
pub struct MinecraftWriter {
    pub(crate) stream: tokio::net::tcp::OwnedWriteHalf,
    cipher: Option<EncryptCipher>,
}

impl MinecraftWriter {
    /// Write a pre-framed raw packet (as received from the other side).
    ///
    /// # Errors
    ///
    /// Returns [`ConnectionError::Io`].
    pub async fn write_raw_frame(&mut self, mut data: Vec<u8>) -> Result<(), ConnectionError> {
        if let Some(cipher) = &mut self.cipher {
            cipher.encrypt(&mut data);
        }
        self.stream.write_all(&data).await?;
        Ok(())
    }
}

#[cfg(test)]
#[expect(
    clippy::large_futures,
    reason = "MinecraftConnection carries large framing buffers; acceptable in tests"
)]
mod tests {
    use std::time::Duration;

    use deepslate_protocol::codec;
    use deepslate_protocol::packet::Packet;
    use deepslate_protocol::packet::login::SetCompressionPacket;

    use super::MinecraftConnection;

    /// Helper: create a `MinecraftConnection` pair connected via an in-memory
    /// duplex stream. The `buf_size` controls the internal buffer of the duplex
    /// channel (not the connection's read buffer).
    fn duplex_pair(
        buf_size: usize,
    ) -> (
        MinecraftConnection<tokio::io::DuplexStream>,
        MinecraftConnection<tokio::io::DuplexStream>,
    ) {
        let (a, b) = tokio::io::duplex(buf_size);
        let compression = libdeflater::CompressionLvl::new(1).expect("valid level");
        let timeout = Duration::from_secs(5);
        (
            MinecraftConnection::new(a, compression, timeout),
            MinecraftConnection::new(b, compression, timeout),
        )
    }

    /// Decode a packet from a raw frame (packet ID + payload).
    fn decode_frame<P: Packet>(frame: &[u8]) -> P {
        let mut cursor = frame;
        let packet_id = codec::read_packet_id(&mut cursor).expect("valid packet id");
        assert_eq!(packet_id, P::PACKET_ID, "unexpected packet id");
        P::decode(&mut cursor).expect("valid packet")
    }

    #[tokio::test]
    async fn write_then_read_frame_round_trips() {
        let (mut writer, mut reader) = duplex_pair(8192);

        let packet = SetCompressionPacket { threshold: 256 };
        writer.write_packet(&packet).await.unwrap();
        drop(writer);

        let frame = reader
            .read_frame()
            .await
            .unwrap()
            .expect("expected a frame");
        let decoded: SetCompressionPacket = decode_frame(&frame);
        assert_eq!(decoded, packet);

        // After the writer is dropped, the next read should return None (EOF).
        assert!(reader.read_frame().await.unwrap().is_none());
    }

    #[tokio::test]
    async fn round_trip_with_compression() {
        let (mut writer, mut reader) = duplex_pair(8192);

        // Enable compression with a threshold of 0 (compress everything).
        writer.enable_compression(0);
        reader.enable_compression(0);

        let packet = SetCompressionPacket { threshold: 512 };
        writer.write_packet(&packet).await.unwrap();
        drop(writer);

        let frame = reader
            .read_frame()
            .await
            .unwrap()
            .expect("expected a frame");
        let decoded: SetCompressionPacket = decode_frame(&frame);
        assert_eq!(decoded, packet);
    }

    #[tokio::test]
    async fn round_trip_with_encryption() {
        let (mut writer, mut reader) = duplex_pair(8192);

        let shared_secret = b"0123456789abcdef"; // 16 bytes for AES-128
        writer.enable_encryption(shared_secret);
        reader.enable_encryption(shared_secret);

        let packet = SetCompressionPacket { threshold: 128 };
        writer.write_packet(&packet).await.unwrap();
        drop(writer);

        let frame = reader
            .read_frame()
            .await
            .unwrap()
            .expect("expected a frame");
        let decoded: SetCompressionPacket = decode_frame(&frame);
        assert_eq!(decoded, packet);
    }

    #[tokio::test]
    async fn round_trip_with_compression_and_encryption() {
        let (mut writer, mut reader) = duplex_pair(8192);

        let shared_secret = b"fedcba9876543210";
        writer.enable_encryption(shared_secret);
        reader.enable_encryption(shared_secret);
        writer.enable_compression(0);
        reader.enable_compression(0);

        let packet = SetCompressionPacket { threshold: 1024 };
        writer.write_packet(&packet).await.unwrap();
        drop(writer);

        let frame = reader
            .read_frame()
            .await
            .unwrap()
            .expect("expected a frame");
        let decoded: SetCompressionPacket = decode_frame(&frame);
        assert_eq!(decoded, packet);
    }

    #[tokio::test]
    async fn read_frame_returns_none_on_eof() {
        let (writer, mut reader) = duplex_pair(8192);
        drop(writer);

        let result = reader.read_frame().await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn read_frame_timeout_returns_error_when_idle() {
        let timeout = Duration::from_millis(50);
        let (a, b) = tokio::io::duplex(8192);
        let compression = libdeflater::CompressionLvl::new(1).expect("valid level");
        let mut reader = MinecraftConnection::new(a, compression, timeout);

        // Keep the other end alive but send nothing.
        let _writer = b;

        let result = reader.read_frame_timeout().await;
        assert!(
            matches!(result, Err(super::ConnectionError::Timeout)),
            "expected Timeout, got {result:?}"
        );
    }

    #[tokio::test]
    async fn multiple_packets_round_trip() {
        let (mut writer, mut reader) = duplex_pair(8192);

        let packets = [
            SetCompressionPacket { threshold: 0 },
            SetCompressionPacket { threshold: 256 },
            SetCompressionPacket { threshold: -1 },
        ];

        for packet in &packets {
            writer.write_packet(packet).await.unwrap();
        }
        drop(writer);

        for expected in &packets {
            let frame = reader
                .read_frame()
                .await
                .unwrap()
                .expect("expected a frame");
            let decoded: SetCompressionPacket = decode_frame(&frame);
            assert_eq!(&decoded, expected);
        }

        assert!(reader.read_frame().await.unwrap().is_none());
    }

    #[tokio::test]
    async fn read_frame_rejects_oversized_read_buffer() {
        use std::io;
        use std::pin::Pin;
        use std::task::{Context, Poll};

        use deepslate_protocol::types::ProtocolError;
        use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

        /// Mock stream that yields data from an internal buffer in small
        /// chunks, simulating a peer that sends faster than we consume.
        struct FloodStream {
            data: Vec<u8>,
            pos: usize,
        }

        impl AsyncRead for FloodStream {
            fn poll_read(
                mut self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
                buf: &mut ReadBuf<'_>,
            ) -> Poll<io::Result<()>> {
                let remaining = &self.data[self.pos..];
                if remaining.is_empty() {
                    return Poll::Ready(Ok(())); // EOF
                }
                let n = remaining.len().min(buf.remaining());
                buf.put_slice(&remaining[..n]);
                self.pos += n;
                Poll::Ready(Ok(()))
            }
        }

        impl AsyncWrite for FloodStream {
            fn poll_write(
                self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
                buf: &[u8],
            ) -> Poll<io::Result<usize>> {
                Poll::Ready(Ok(buf.len()))
            }
            fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                Poll::Ready(Ok(()))
            }
            fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                Poll::Ready(Ok(()))
            }
        }

        // Build a payload: a VarInt declaring a 2 MiB frame, followed by
        // enough filler to exceed MAX_READ_BUF_SIZE (2.5 MiB). The frame
        // never completes (needs 2 MiB of body, we provide > 2.5 MiB total
        // including the header), so try_read_frame keeps returning None and
        // the buffer grows until the pre-read check fires.
        let frame_len: i32 = 2 * 1024 * 1024; // MAX_FRAME_SIZE
        let mut data = Vec::new();
        deepslate_protocol::varint::write_var_int(&mut data, frame_len);
        // Pad to well beyond MAX_READ_BUF_SIZE. The frame needs `frame_len`
        // bytes of body to complete, so at 2.5 MiB of total buffer the check
        // fires before we reach the 2 MiB body threshold.
        data.resize(super::MAX_READ_BUF_SIZE + 64 * 1024, 0xAB);

        let stream = FloodStream { data, pos: 0 };
        let compression = libdeflater::CompressionLvl::new(1).expect("valid level");
        let timeout = Duration::from_secs(5);
        let mut conn = MinecraftConnection::new(stream, compression, timeout);

        let result = conn.read_frame().await;
        assert!(
            matches!(
                result,
                Err(super::ConnectionError::Protocol(
                    ProtocolError::ReadBufferOverflow { .. }
                ))
            ),
            "expected ReadBufferOverflow, got {result:?}"
        );
    }
}