From b4abc6b297347bc510229242e90741f27fe49afa Mon Sep 17 00:00:00 2001 From: Snowiiii Date: Thu, 8 Aug 2024 16:37:40 +0200 Subject: [PATCH] Add Packet compression --- Cargo.lock | 2 +- README.md | 2 +- pumpkin-protocol/Cargo.toml | 2 + pumpkin-protocol/src/lib.rs | 8 +++ pumpkin-protocol/src/packet_decoder.rs | 53 +++++++++++++++++-- pumpkin-protocol/src/packet_encoder.rs | 70 +++++++++++++++++++++++++- pumpkin/Cargo.toml | 2 - pumpkin/src/client/client_packet.rs | 13 ++++- pumpkin/src/client/mod.rs | 6 +++ pumpkin/src/config/mod.rs | 25 +++++++++ 10 files changed, 172 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c9d4122fa..768e1fcf9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1146,7 +1146,6 @@ dependencies = [ "base64", "crossbeam-channel", "digest 0.11.0-pre.9", - "flate2", "image 0.25.2", "log", "mio", @@ -1196,6 +1195,7 @@ dependencies = [ "bytes", "cfb8", "fastnbt 2.5.0 (git+https://github.com/owengage/fastnbt.git)", + "flate2", "log", "num-derive", "num-traits", diff --git a/README.md b/README.md index b370206dd..ba2b50a25 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Pumpkin is currently under heavy development. - Login - [x] Authentication - [x] Encryption - - [ ] Packet Compression + - [x] Packet Compression - Player Configuration - [x] Registries (biome types, paintings, dimensions) - [x] Server Brand diff --git a/pumpkin-protocol/Cargo.toml b/pumpkin-protocol/Cargo.toml index 564d8de38..941bdfbfe 100644 --- a/pumpkin-protocol/Cargo.toml +++ b/pumpkin-protocol/Cargo.toml @@ -15,6 +15,8 @@ serde = { version = "1.0", features = ["derive"] } # to parse strings to json responses serde_json = "1.0" +flate2 = "1.0.30" + thiserror = "1.0.63" log = "0.4" num-traits = "0.2" diff --git a/pumpkin-protocol/src/lib.rs b/pumpkin-protocol/src/lib.rs index 02a34cc27..c84279d16 100644 --- a/pumpkin-protocol/src/lib.rs +++ b/pumpkin-protocol/src/lib.rs @@ -145,8 +145,16 @@ pub enum PacketError { DecodeID, #[error("failed to encode packet ID")] EncodeID, + #[error("failed to encode packet Length")] + EncodeLength, + #[error("failed to encode packet data")] + EncodeData, #[error("failed to write encoded packet")] EncodeFailedWrite, + #[error("failed to write into decoder")] + FailedWrite, + #[error("failed to flush decoder")] + FailedFinish, #[error("failed to write encoded packet to connection")] ConnectionWrite, #[error("packet exceeds maximum length")] diff --git a/pumpkin-protocol/src/packet_decoder.rs b/pumpkin-protocol/src/packet_decoder.rs index c53a7be08..3d3e3449a 100644 --- a/pumpkin-protocol/src/packet_decoder.rs +++ b/pumpkin-protocol/src/packet_decoder.rs @@ -1,6 +1,11 @@ use aes::cipher::{generic_array::GenericArray, BlockDecryptMut, BlockSizeUser, KeyIvInit}; use bytes::{Buf, BytesMut}; +use std::io::Write; + +use bytes::BufMut; +use flate2::write::ZlibDecoder; + use crate::{ bytebuf::ByteBuffer, PacketError, RawPacket, VarInt, VarIntDecodeError, MAX_PACKET_SIZE, }; @@ -11,7 +16,8 @@ type Cipher = cfb8::Decryptor; #[derive(Default)] pub struct PacketDecoder { buf: BytesMut, - + decompress_buf: BytesMut, + compression: Option, cipher: Option, } @@ -37,11 +43,46 @@ impl PacketDecoder { let packet_len_len = VarInt(packet_len).written_size(); let mut data; + if self.compression.is_some() { + r = &r[..packet_len as usize]; + + let data_len = VarInt::decode(&mut r).map_err(|_| PacketError::TooLong)?.0; + + if !(0..=MAX_PACKET_SIZE).contains(&data_len) { + Err(PacketError::OutOfBounds)? + } + + // Is this packet compressed? + if data_len > 0 { + debug_assert!(self.decompress_buf.is_empty()); + + self.decompress_buf.put_bytes(0, data_len as usize); + + // TODO: use libdeflater or zune-inflate? + let mut z = ZlibDecoder::new(&mut self.decompress_buf[..]); - // no compression + z.write_all(r).map_err(|_| PacketError::FailedWrite)?; + z.finish().map_err(|_| PacketError::FailedFinish)?; - self.buf.advance(packet_len_len); - data = self.buf.split_to(packet_len as usize); + let total_packet_len = VarInt(packet_len).written_size() + packet_len as usize; + + self.buf.advance(total_packet_len); + + data = self.decompress_buf.split(); + } else { + debug_assert_eq!(data_len, 0); + + let remaining_len = r.len(); + + self.buf.advance(packet_len_len + 1); + + data = self.buf.split_to(remaining_len); + } + } else { + // no compression + self.buf.advance(packet_len_len); + data = self.buf.split_to(packet_len as usize); + } r = &data[..]; let packet_id = VarInt::decode(&mut r).map_err(|_| PacketError::DecodeID)?; @@ -64,6 +105,10 @@ impl PacketDecoder { self.cipher = Some(cipher); } + pub fn set_compression(&mut self, compression: Option) { + self.compression = compression; + } + fn decrypt_bytes(cipher: &mut Cipher, bytes: &mut [u8]) { for chunk in bytes.chunks_mut(Cipher::block_size()) { let gen_arr = GenericArray::from_mut_slice(chunk); diff --git a/pumpkin-protocol/src/packet_encoder.rs b/pumpkin-protocol/src/packet_encoder.rs index 53f2116ab..9e3455e59 100644 --- a/pumpkin-protocol/src/packet_encoder.rs +++ b/pumpkin-protocol/src/packet_encoder.rs @@ -3,6 +3,11 @@ use std::io::Write; use aes::cipher::{generic_array::GenericArray, BlockEncryptMut, BlockSizeUser, KeyIvInit}; use bytes::{BufMut, BytesMut}; +use std::io::Read; + +use flate2::bufread::ZlibEncoder; +use flate2::Compression; + use crate::{bytebuf::ByteBuffer, ClientPacket, PacketError, VarInt, MAX_PACKET_SIZE}; type Cipher = cfb8::Encryptor; @@ -11,7 +16,8 @@ type Cipher = cfb8::Encryptor; #[derive(Default)] pub struct PacketEncoder { buf: BytesMut, - + compress_buf: Vec, + compression: Option<(u32, u32)>, cipher: Option, } @@ -33,8 +39,64 @@ impl PacketEncoder { let data_len = self.buf.len() - start_len; - if false { // compression + if let Some((threshold, compression_level)) = self.compression { + if data_len > threshold as usize { + let mut z = + ZlibEncoder::new(&self.buf[start_len..], Compression::new(compression_level)); + + self.compress_buf.clear(); + + let data_len_size = VarInt(data_len as i32).written_size(); + + let packet_len = data_len_size + z.read_to_end(&mut self.compress_buf).unwrap(); + + if packet_len >= MAX_PACKET_SIZE as usize { + Err(PacketError::TooLong)? + } + + drop(z); + + self.buf.truncate(start_len); + + let mut writer = (&mut self.buf).writer(); + + VarInt(packet_len as i32) + .encode(&mut writer) + .map_err(|_| PacketError::EncodeLength)?; + VarInt(data_len as i32) + .encode(&mut writer) + .map_err(|_| PacketError::EncodeData)?; + self.buf.extend_from_slice(&self.compress_buf); + } else { + let data_len_size = 1; + let packet_len = data_len_size + data_len; + + if packet_len >= MAX_PACKET_SIZE as usize { + Err(PacketError::TooLong)? + } + + let packet_len_size = VarInt(packet_len as i32).written_size(); + + let data_prefix_len = packet_len_size + data_len_size; + + self.buf.put_bytes(0, data_prefix_len); + self.buf + .copy_within(start_len..start_len + data_len, start_len + data_prefix_len); + + let mut front = &mut self.buf[start_len..]; + + VarInt(packet_len as i32) + .encode(&mut front) + .map_err(|_| PacketError::EncodeLength)?; + // Zero for no compression on this packet. + VarInt(0) + .encode(front) + .map_err(|_| PacketError::EncodeData)?; + } + + return Ok(()); } + let packet_len = data_len; if packet_len >= MAX_PACKET_SIZE as usize { @@ -59,6 +121,10 @@ impl PacketEncoder { self.cipher = Some(Cipher::new_from_slices(key, key).expect("invalid key")); } + pub fn set_compression(&mut self, compression: Option<(u32, u32)>) { + self.compression = compression; + } + pub fn take(&mut self) -> BytesMut { if let Some(cipher) = &mut self.cipher { for chunk in self.buf.chunks_mut(Cipher::block_size()) { diff --git a/pumpkin/Cargo.toml b/pumpkin/Cargo.toml index 0239c0f19..6b599f81e 100644 --- a/pumpkin/Cargo.toml +++ b/pumpkin/Cargo.toml @@ -25,8 +25,6 @@ num-bigint = "0.4.6" rsa = "0.9.6" rsa-der = "0.3.0" -flate2 = "1.0.30" - # authentication reqwest = { version = "0.12.5", features = ["json"]} diff --git a/pumpkin/src/client/client_packet.rs b/pumpkin/src/client/client_packet.rs index aaaf66c59..754103392 100644 --- a/pumpkin/src/client/client_packet.rs +++ b/pumpkin/src/client/client_packet.rs @@ -2,7 +2,7 @@ use num_traits::FromPrimitive; use pumpkin_protocol::{ client::{ config::{CFinishConfig, CKnownPacks, CRegistryData}, - login::{CEncryptionRequest, CLoginSuccess}, + login::{CEncryptionRequest, CLoginSuccess, CSetCompression}, status::{CPingResponse, CStatusResponse}, }, server::{ @@ -134,6 +134,17 @@ impl Client { unpack_textures(ele, &server.advanced_config.authentication.textures); } + // enable compression + if server.advanced_config.packet_compression.enabled { + let threshold = server + .advanced_config + .packet_compression + .compression_threshold; + let level = server.advanced_config.packet_compression.compression_level; + self.send_packet(CSetCompression::new(threshold.into())); + self.set_compression(Some((threshold, level))); + } + if let Some(profile) = self.gameprofile.as_ref().cloned() { let packet = CLoginSuccess::new(profile.id, profile.name, &profile.properties, false); self.send_packet(packet); diff --git a/pumpkin/src/client/mod.rs b/pumpkin/src/client/mod.rs index 319948b9e..daae01a85 100644 --- a/pumpkin/src/client/mod.rs +++ b/pumpkin/src/client/mod.rs @@ -111,6 +111,12 @@ impl Client { Ok(()) } + // Compression threshold, Compression level + pub fn set_compression(&mut self, compression: Option<(u32, u32)>) { + self.dec.set_compression(compression.map(|v| v.0)); + self.enc.set_compression(compression); + } + pub fn is_player(&self) -> bool { self.player.is_some() } diff --git a/pumpkin/src/config/mod.rs b/pumpkin/src/config/mod.rs index cd515490a..02acb657a 100644 --- a/pumpkin/src/config/mod.rs +++ b/pumpkin/src/config/mod.rs @@ -17,6 +17,7 @@ const CURRENT_BASE_VERSION: &str = "1.0.0"; pub struct AdvancedConfiguration { pub commands: Commands, pub authentication: Authentication, + pub packet_compression: Compression, } #[derive(Deserialize, Serialize)] @@ -26,6 +27,29 @@ pub struct Commands { // TODO: commands... } +#[derive(Deserialize, Serialize)] +// Packet compression +pub struct Compression { + // Is compression enabled ? + pub enabled: bool, + // The compression threshold used when compression is enabled + pub compression_threshold: u32, + // A value between 0..9 + // 1 = Optimize for the best speed of encoding. + // 9 = Optimize for the size of data being encoded. + pub compression_level: u32, +} + +impl Default for Compression { + fn default() -> Self { + Self { + enabled: true, + compression_threshold: 256, + compression_level: 4, + } + } +} + impl Default for Commands { fn default() -> Self { Self { use_console: true } @@ -38,6 +62,7 @@ impl Default for AdvancedConfiguration { Self { authentication: Authentication::default(), commands: Commands::default(), + packet_compression: Compression::default(), } } }