add compression

Author: cyteon

src/net/packet.rs

@@ -6,8 +6,9 @@ mod packets {
 }
 
 use tokio::io::AsyncReadExt;
+use std::io::Write;
 
-use crate::net::codec::read_var;
+use crate::net::codec::{read_var, write_var};
 
 pub enum ClientPacket {
     // status state
@@ -35,13 +36,58 @@ pub enum ProtocolState {
     Play
 }
 
-pub async fn read_packet<R: AsyncReadExt + Unpin>(stream: &mut R, state: &ProtocolState) -> anyhow::Result<Option<ClientPacket>> {
+pub fn encode_packet(data: &[u8]) -> Vec<u8> {
+    let mut result = Vec::new();
+
+    if data.len() > 256 {
+        let uncompressed_len = data.len();
+
+        let mut encoder = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::fast());
+        encoder.write_all(&data).unwrap();
+
+        let compressed = encoder.finish().unwrap();
+
+        let mut inner = Vec::with_capacity(5 + compressed.len());
+        write_var(&mut inner, uncompressed_len as i32).unwrap();
+        inner.extend_from_slice(&compressed);
+
+        write_var(&mut result, inner.len() as i32).unwrap();
+        result.extend_from_slice(&inner);
+    } else {
+        let mut inner = Vec::with_capacity(5);
+        write_var(&mut inner, 0).unwrap();
+        inner.extend_from_slice(&data);
+
+        write_var(&mut result, inner.len() as i32).unwrap();
+        result.extend_from_slice(&inner);
+    }
+
+    result
+}
+
+pub async fn read_packet<R: AsyncReadExt + Unpin>(stream: &mut R, state: &ProtocolState, compression: bool) -> anyhow::Result<Option<ClientPacket>> {
     let packet_len = read_var(stream).await.map_err(|e| anyhow::anyhow!("Failed to read packet length: {}", e))?;
 
     let mut packet_buf = vec![0u8; packet_len as usize];
     stream.read_exact(&mut packet_buf).await.map_err(|e| anyhow::anyhow!("Failed to read packet data: {}", e))?;
 
     let mut cursor = std::io::Cursor::new(packet_buf);
+
+    if compression {
+        let data_len = read_var(&mut cursor).await.map_err(|e| anyhow::anyhow!("Failed to read data length: {}", e))?;
+
+        if data_len != 0 {
+            let mut compressed = Vec::new();
+            std::io::Read::read_to_end(&mut cursor, &mut compressed).map_err(|e| anyhow::anyhow!("Failed to read compressed data: {}", e))?;
+
+            let mut decoder = flate2::read::ZlibDecoder::new(&compressed[..]);
+            let mut decompressed = Vec::with_capacity(data_len as usize);
+            std::io::Read::read_to_end(&mut decoder, &mut decompressed).map_err(|e| anyhow::anyhow!("Failed to decompress data: {}", e))?;
+
+            cursor = std::io::Cursor::new(decompressed);
+        }
+    }
+
     let packet_id = read_var(&mut cursor).await.map_err(|e| anyhow::anyhow!("Failed to read packet ID: {}", e))?;
 
     if packet_id != 0x0C && packet_id != 0x1D { // these packets are spammy

src/net/packets/clientbound/chunk_data_with_light.rs

@@ -66,10 +66,6 @@ pub async fn send_chunk_data_with_light<W: tokio::io::AsyncWriteExt + Unpin>(
         packet_data.extend_from_slice(&FULLBRIGHT_ENTRY);
     }
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;
 

src/net/packets/clientbound/finish_configuration.rs

@@ -3,10 +3,6 @@ use crate::net::codec::write_var;
 pub async fn send_finish_configuration<W: tokio::io::AsyncWriteExt + Unpin>(stream: &mut W) -> anyhow::Result<()> {
     let packet_data = vec![crate::net::packet::configuration::clientbound::FINISH_CONFIGURATION as u8];
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;
 

src/net/packets/clientbound/game_event.rs

@@ -7,10 +7,6 @@ pub async fn send_game_event<W: tokio::io::AsyncWriteExt + Unpin>(stream: &mut W
     packet_data.push(event);
     packet_data.write_f32(value).await?;
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;
 

src/net/packets/clientbound/keep_alive.rs

@@ -7,10 +7,6 @@ pub async fn send_keep_alive<W: tokio::io::AsyncWriteExt + Unpin>(stream: &mut W
     let ms = chrono::Utc::now().timestamp_millis();
     packet_data.write_i64(ms).await?;
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;
 

src/net/packets/clientbound/known_packs.rs

@@ -14,10 +14,6 @@ pub async fn send_known_packs<W: tokio::io::AsyncWriteExt + Unpin>(stream: &mut
     write_var(&mut packet_data, "1.21".len() as i32)?;
     packet_data.extend_from_slice("1.21".as_bytes());
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;
 

src/net/packets/clientbound/login_play.rs

@@ -36,10 +36,6 @@ pub async fn send_login_play<W: tokio::io::AsyncWriteExt + Unpin>(stream: &mut W
     write_var(&mut packet_data, 62)?; // sea level
     packet_data.push(false as u8); // enforce secure chat
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;
 

src/net/packets/clientbound/login_success.rs

@@ -39,10 +39,6 @@ pub async fn send_login_success<W: tokio::io::AsyncWriteExt + Unpin>(stream: &mu
         }
     }
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;
 

src/net/packets/clientbound/mod.rs

@@ -15,4 +15,5 @@ pub mod plugin_message;
 pub mod set_center_chunk;
 pub mod chunk_data_with_light;
 pub mod player_info_remove;
-pub mod player_chat_message;
+pub mod player_chat_message;
+pub mod set_compression;

src/net/packets/clientbound/player_chat_message.rs

@@ -64,10 +64,6 @@ pub async fn send_player_chat_message<W: tokio::io::AsyncWriteExt + Unpin>(
 
     packet_data.write_u8(0)?; // no target name
 
-    let mut len_prefix = Vec::with_capacity(5);
-    write_var(&mut len_prefix, packet_data.len() as i32)?;
-
-    stream.write_all(&len_prefix).await?;
     stream.write_all(&packet_data).await?;
     stream.flush().await?;