Skip to content

Commit

Permalink
Add Packet compression
Browse files Browse the repository at this point in the history
  • Loading branch information
Snowiiii committed Aug 8, 2024
1 parent 0b55a90 commit b4abc6b
Show file tree
Hide file tree
Showing 10 changed files with 172 additions and 11 deletions.
2 changes: 1 addition & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions pumpkin-protocol/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
8 changes: 8 additions & 0 deletions pumpkin-protocol/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down
53 changes: 49 additions & 4 deletions pumpkin-protocol/src/packet_decoder.rs
Original file line number Diff line number Diff line change
@@ -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,
};
Expand All @@ -11,7 +16,8 @@ type Cipher = cfb8::Decryptor<aes::Aes128>;
#[derive(Default)]
pub struct PacketDecoder {
buf: BytesMut,

decompress_buf: BytesMut,
compression: Option<u32>,
cipher: Option<Cipher>,
}

Expand All @@ -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)?;
Expand All @@ -64,6 +105,10 @@ impl PacketDecoder {
self.cipher = Some(cipher);
}

pub fn set_compression(&mut self, compression: Option<u32>) {
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);
Expand Down
70 changes: 68 additions & 2 deletions pumpkin-protocol/src/packet_encoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<aes::Aes128>;
Expand All @@ -11,7 +16,8 @@ type Cipher = cfb8::Encryptor<aes::Aes128>;
#[derive(Default)]
pub struct PacketEncoder {
buf: BytesMut,

compress_buf: Vec<u8>,
compression: Option<(u32, u32)>,
cipher: Option<Cipher>,
}

Expand All @@ -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 {
Expand All @@ -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()) {
Expand Down
2 changes: 0 additions & 2 deletions pumpkin/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]}

Expand Down
13 changes: 12 additions & 1 deletion pumpkin/src/client/client_packet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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);
Expand Down
6 changes: 6 additions & 0 deletions pumpkin/src/client/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
Expand Down
25 changes: 25 additions & 0 deletions pumpkin/src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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 }
Expand All @@ -38,6 +62,7 @@ impl Default for AdvancedConfiguration {
Self {
authentication: Authentication::default(),
commands: Commands::default(),
packet_compression: Compression::default(),
}
}
}
Expand Down

0 comments on commit b4abc6b

Please sign in to comment.