Skip to content

Commit

Permalink
pumpkin-protocol: remove pumpkin-config
Browse files Browse the repository at this point in the history
  • Loading branch information
Snowiiii committed Dec 26, 2024
1 parent d833582 commit a033972
Show file tree
Hide file tree
Showing 5 changed files with 89 additions and 66 deletions.
1 change: 0 additions & 1 deletion pumpkin-protocol/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ query = []

[dependencies]
pumpkin-nbt = { path = "../pumpkin-nbt" }
pumpkin-config = { path = "../pumpkin-config" }
pumpkin-macros = { path = "../pumpkin-macros" }
pumpkin-world = { path = "../pumpkin-world" }
pumpkin-core = { path = "../pumpkin-core" }
Expand Down
15 changes: 15 additions & 0 deletions pumpkin-protocol/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,21 @@ pub const MAX_PACKET_SIZE: i32 = 2097152;

pub type FixedBitSet = bytes::Bytes;

/// Represents a compression threshold.
///
/// The threshold determines the minimum size of data that should be compressed.
/// Data smaller than the threshold will not be compressed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CompressionThreshold(pub u32);

/// Represents a compression level.
///
/// The level controls the amount of compression applied to the data.
/// Higher levels generally result in higher compression ratios but also
/// increase CPU usage.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CompressionLevel(pub u32);

#[derive(Debug, PartialEq, Clone, Copy)]
pub enum ConnectionState {
HandShake,
Expand Down
17 changes: 9 additions & 8 deletions pumpkin-protocol/src/packet_decoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,7 @@ pub struct PacketDecoder {
buf: BytesMut,
decompress_buf: BytesMut,
cipher: Option<Cipher>,
compression: bool,
decompressor: Decompressor,
decompressor: Option<Decompressor>,
}

// Manual implementation of Default trait for PacketDecoder
Expand All @@ -29,8 +28,7 @@ impl Default for PacketDecoder {
buf: BytesMut::new(),
decompress_buf: BytesMut::new(),
cipher: None,
compression: false,
decompressor: Decompressor::new(),
decompressor: None,
}
}
}
Expand Down Expand Up @@ -58,7 +56,7 @@ impl PacketDecoder {
let packet_len_len = VarInt(packet_len).written_size();

let mut data;
if self.compression {
if let Some(decompressor) = &mut self.decompressor {
r = &r[..packet_len as usize];

let data_len = VarInt::decode(&mut r)
Expand All @@ -77,8 +75,7 @@ impl PacketDecoder {
self.decompress_buf.resize(data_len as usize, 0);

// Perform decompression using libdeflater
let decompressed_size = self
.decompressor
let decompressed_size = decompressor
.zlib_decompress(r, &mut self.decompress_buf)
.map_err(PacketDecodeError::from)?;

Expand Down Expand Up @@ -135,7 +132,11 @@ impl PacketDecoder {

/// Sets ZLib Decompression
pub fn set_compression(&mut self, compression: bool) {
self.compression = compression;
if compression {
self.decompressor = Some(Decompressor::new());
} else {
self.decompressor = None
}
}

fn decrypt_bytes(cipher: &mut Cipher, bytes: &mut [u8]) {
Expand Down
111 changes: 56 additions & 55 deletions pumpkin-protocol/src/packet_encoder.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
use aes::cipher::{generic_array::GenericArray, BlockEncryptMut, BlockSizeUser, KeyIvInit};
use bytes::{BufMut, BytesMut};
use pumpkin_config::compression::CompressionInfo;
use thiserror::Error;

use libdeflater::{CompressionLvl, Compressor};

use crate::{codec::Codec, ClientPacket, VarInt, MAX_PACKET_SIZE};
use crate::{
codec::Codec, ClientPacket, CompressionLevel, CompressionThreshold, VarInt, MAX_PACKET_SIZE,
};

type Cipher = cfb8::Encryptor<aes::Aes128>;

Expand All @@ -16,8 +17,8 @@ pub struct PacketEncoder {
buf: BytesMut,
compress_buf: Vec<u8>,
cipher: Option<Cipher>,
compression_threshold: Option<u32>,
compressor: Compressor, // Reuse the compressor for all packets
// compression and compression threshold
compression: Option<(Compressor, CompressionThreshold)>,
}

// Manual implementation of Default trait for PacketEncoder
Expand All @@ -28,8 +29,7 @@ impl Default for PacketEncoder {
buf: BytesMut::with_capacity(1024),
compress_buf: Vec::with_capacity(1024),
cipher: None,
compression_threshold: None,
compressor: Compressor::new(CompressionLvl::fastest()), // init compressor with fastest compression level
compression: None, // init compressor with fastest compression level
}
}
}
Expand Down Expand Up @@ -73,24 +73,22 @@ impl PacketEncoder {
packet.write(&mut self.buf);
let data_len = self.buf.len() - start_len;

if let Some(compression_threshold) = self.compression_threshold {
if data_len > compression_threshold as usize {
if let Some((compressor, compression_threshold)) = &mut self.compression {
if data_len > compression_threshold.0 as usize {
// Get the data to compress
let data_to_compress = &self.buf[start_len..];

// Clear the compression buffer
self.compress_buf.clear();

// Compute the maximum size of compressed data
let max_compressed_size =
self.compressor.zlib_compress_bound(data_to_compress.len());
let max_compressed_size = compressor.zlib_compress_bound(data_to_compress.len());

// Ensure compress_buf has enough capacity
self.compress_buf.resize(max_compressed_size, 0);

// Compress the data
let compressed_size = self
.compressor
let compressed_size = compressor
.zlib_compress(data_to_compress, &mut self.compress_buf)
.map_err(|e| PacketEncodeError::CompressionFailed(e.to_string()))?;

Expand Down Expand Up @@ -166,26 +164,30 @@ impl PacketEncoder {
}
}

/// Enables or disables Zlib compression with the given options.
/// Enables or disables Zlib compression.
///
/// If `compression` is `Some`, compression is enabled with the provided
/// options. If `compression` is `None`, compression is disabled.
pub fn set_compression(&mut self, compression: Option<CompressionInfo>) {
// Reset the compressor with the new compression level
if let Some(compression) = &compression {
self.compression_threshold = Some(compression.threshold);
let compression_level = compression.level as i32;
let level = match CompressionLvl::new(compression_level) {
Ok(level) => level,
Err(error) => {
log::error!("Invalid compression level {:?}", error);
return;
}
};
self.compressor = Compressor::new(level);
} else {
self.compression_threshold = None;
/// If `compression` is `Some`, compression is enabled with the given `threshold`
/// for triggering compression and the specified `level`. If `compression` is
/// `None`, compression is disabled.
///
/// # Errors
///
/// Returns an `CompressionLevelError` if an invalid compression level is provided.
pub fn set_compression(
&mut self,
compression: Option<(CompressionThreshold, CompressionLevel)>,
) -> Result<(), CompressionLevelError> {
match compression {
Some((threshold, level)) => {
let level =
CompressionLvl::new(level.0 as i32).map_err(|_| CompressionLevelError)?;
self.compression = Some((Compressor::new(level), threshold));
}
None => {
self.compression = None;
}
}
Ok(())
}

/// Encrypts the data in the internal buffer and returns it as a `BytesMut`.
Expand All @@ -209,6 +211,10 @@ impl PacketEncoder {
}
}

#[derive(Error, Debug)]
#[error("Invalid compression Level")]
pub struct CompressionLevelError;

/// Errors that can occur during packet encoding.
#[derive(Error, Debug)]
pub enum PacketEncodeError {
Expand Down Expand Up @@ -268,15 +274,15 @@ mod tests {
/// Helper function to build a packet with optional compression and encryption
fn build_packet_with_encoder<T: ClientPacket>(
packet: &T,
compression_info: Option<CompressionInfo>,
compression_info: Option<(CompressionThreshold, CompressionLevel)>,
key: Option<&[u8; 16]>,
) -> BytesMut {
let mut encoder = PacketEncoder::default();

if let Some(compression) = compression_info {
encoder.set_compression(Some(compression));
encoder.set_compression(Some(compression)).unwrap();
} else {
encoder.set_compression(None);
encoder.set_compression(None).unwrap();
}

if let Some(key) = key {
Expand Down Expand Up @@ -328,14 +334,12 @@ mod tests {
// Create a CStatusResponse packet
let packet = CStatusResponse::new("{\"description\": \"A Minecraft Server\"}");

// Compression threshold is set to 0 to force compression
let compression_info = CompressionInfo {
threshold: 0,
level: 6, // Standard compression level
};

// Build the packet with compression enabled
let packet_bytes = build_packet_with_encoder(&packet, Some(compression_info), None);
let packet_bytes = build_packet_with_encoder(
&packet,
Some((CompressionThreshold(0), CompressionLevel(6))),
None,
);

// Decode the packet manually to verify correctness
let mut buffer = &packet_bytes[..];
Expand Down Expand Up @@ -418,18 +422,16 @@ mod tests {
// Create a CStatusResponse packet
let packet = CStatusResponse::new("{\"description\": \"A Minecraft Server\"}");

// Compression threshold is set to 0 to force compression
let compression_info = CompressionInfo {
threshold: 0,
level: 6, // Standard compression level
};

// Encryption key and IV (IV is the same as key in this case)
let key = [0x01u8; 16]; // Example key

// Build the packet with both compression and encryption enabled
let mut packet_bytes =
build_packet_with_encoder(&packet, Some(compression_info), Some(&key));
// Compression threshold is set to 0 to force compression
let mut packet_bytes = build_packet_with_encoder(
&packet,
Some((CompressionThreshold(0), CompressionLevel(6))),
Some(&key),
);

// Decrypt the packet
decrypt_aes128(&mut packet_bytes, &key, &key);
Expand Down Expand Up @@ -567,14 +569,13 @@ mod tests {
// Create a CStatusResponse packet with small payload
let packet = CStatusResponse::new("Hi");

// Compression threshold is set to a value higher than payload length
let compression_info = CompressionInfo {
threshold: 10,
level: 6, // Standard compression level
};

// Build the packet with compression enabled
let packet_bytes = build_packet_with_encoder(&packet, Some(compression_info), None);
// Compression threshold is set to a value higher than payload length
let packet_bytes = build_packet_with_encoder(
&packet,
Some((CompressionThreshold(10), CompressionLevel(6))),
None,
);

// Decode the packet manually to verify that it was not compressed
let mut buffer = &packet_bytes[..];
Expand Down
11 changes: 9 additions & 2 deletions pumpkin/src/net/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,8 @@ use pumpkin_protocol::{
},
status::{SStatusPingRequest, SStatusRequest},
},
ClientPacket, ConnectionState, Property, RawPacket, ServerPacket,
ClientPacket, CompressionLevel, CompressionThreshold, ConnectionState, Property, RawPacket,
ServerPacket,
};
use serde::Deserialize;
use sha1::Digest;
Expand Down Expand Up @@ -224,7 +225,13 @@ impl Client {
/// * `compression`: An optional `CompressionInfo` struct containing the compression threshold and compression level.
pub async fn set_compression(&self, compression: Option<CompressionInfo>) {
self.dec.lock().await.set_compression(compression.is_some());
self.enc.lock().await.set_compression(compression);
self.enc
.lock()
.await
.set_compression(
compression.map(|s| (CompressionThreshold(s.threshold), CompressionLevel(s.level))),
)
.unwrap_or_else(|_| log::warn!("invalid compression level"));
}

/// Sends a clientbound packet to the connected client.
Expand Down

0 comments on commit a033972

Please sign in to comment.