deepslate 0.1.0

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

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

use std::io;

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::varint;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

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 {
    /// The underlying TCP stream.
    stream: TcpStream,
    /// 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 write framing output.
    write_buf: Vec<u8>,
}

impl MinecraftConnection {
    /// Wrap a raw TCP stream into a Minecraft connection.
    #[must_use]
    pub fn new(stream: TcpStream, compression_level: libdeflater::CompressionLvl) -> Self {
        Self {
            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(),
            write_buf: Vec::new(),
        }
    }

    /// 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;
    }

    /// Split into owned read/write halves for bidirectional relay.
    pub fn into_split(self) -> (MinecraftReader, MinecraftWriter) {
        let (read_half, write_half) = self.stream.into_split();
        // If encryption is enabled, clone 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,
            },
            MinecraftWriter {
                stream: write_half,
                cipher: enc_cipher,
            },
        )
    }

    /// 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 I/O errors or protocol errors.
    #[allow(clippy::large_stack_arrays)]
    pub async fn read_frame(&mut self) -> io::Result<Option<BytesMut>> {
        loop {
            // Try to extract a frame from the buffer
            if let Some((varint_size, frame_len)) =
                codec::try_read_frame(&self.read_buf).map_err(protocol_err)?
            {
                // 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 frame = self.read_buf.split_to(frame_len);

                if self.compression_threshold >= 0 {
                    let (uncompressed_size, payload) =
                        codec::read_compressed_frame(&frame).map_err(protocol_err)?;
                    if uncompressed_size == 0 {
                        // Payload sits within `frame`; copy only the sub-slice
                        // past the compression VarInt. This is one copy instead
                        // of the previous two.
                        return Ok(Some(BytesMut::from(payload)));
                    }
                    // 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(|e| {
                            protocol_err(
                                deepslate_protocol::types::ProtocolError::CompressionError(
                                    format!("zlib decompression failed: {e}"),
                                ),
                            )
                        })?;
                    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);
                }
            }
        }
    }

    /// Write a typed packet to the connection.
    ///
    /// # Errors
    ///
    /// Returns I/O errors.
    #[allow(clippy::future_not_send)]
    pub async fn write_packet<P: Packet>(&mut self, packet: &P) -> io::Result<()> {
        let packet_data = codec::encode_packet_data(P::PACKET_ID, |buf| packet.encode(buf));
        self.write_raw_packet(&packet_data).await
    }

    /// Write raw packet data (already has packet ID + payload) to the connection.
    ///
    /// # Errors
    ///
    /// Returns I/O errors.
    #[allow(
        clippy::cast_sign_loss,
        clippy::cast_possible_truncation,
        clippy::cast_possible_wrap
    )]
    pub async fn write_raw_packet(&mut self, packet_data: &[u8]) -> io::Result<()> {
        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(|e| {
                        protocol_err(deepslate_protocol::types::ProtocolError::CompressionError(
                            format!("zlib compression failed: {e}"),
                        ))
                    })?;
                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 I/O errors.
    pub async fn set_compression(&mut self, threshold: i32) -> io::Result<()> {
        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 I/O errors 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) -> io::Result<()> {
        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(())
    }
}

/// 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>,
}

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 I/O or protocol errors.
    #[allow(
        clippy::cast_possible_truncation,
        clippy::cast_possible_wrap,
        clippy::large_stack_arrays
    )]
    pub async fn read_raw_frame(&mut self) -> io::Result<Option<Vec<u8>>> {
        loop {
            if let Some((varint_size, frame_len)) =
                codec::try_read_frame(&self.read_buf).map_err(protocol_err)?
            {
                // 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);
                }
            }
        }
    }
}

/// 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 I/O errors.
    pub async fn write_raw_frame(&mut self, mut data: Vec<u8>) -> io::Result<()> {
        if let Some(cipher) = &mut self.cipher {
            cipher.encrypt(&mut data);
        }
        self.stream.write_all(&data).await
    }
}

/// Convert a protocol error into an I/O error.
fn protocol_err(e: deepslate_protocol::types::ProtocolError) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, e)
}