try to implement compression

18 files changed, 175 insertions, 124 deletions

diff --git a/azalea-protocol/Cargo.toml b/azalea-protocol/Cargo.toml
index 9b5c3c07..272816e2 100644
--- a/azalea-protocol/Cargo.toml
+++ b/azalea-protocol/Cargo.toml
@@ -6,6 +6,7 @@ version = "0.1.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+async-compression = {version = "^0.3.8", features = ["tokio", "zlib"]}
 async-recursion = "^0.3.2"
 async-trait = "0.1.51"
 azalea-auth = {path = "../azalea-auth"}
diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index f6dd9fe7..cc06eec3 100644
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -24,6 +24,7 @@ pub struct GameConnection {
     pub flow: PacketFlow,
     /// The buffered writer
     pub stream: TcpStream,
+    pub compression_threshold: Option<u32>,
 }
 
 pub struct StatusConnection {
@@ -36,6 +37,7 @@ pub struct LoginConnection {
     pub flow: PacketFlow,
     /// The buffered writer
     pub stream: TcpStream,
+    pub compression_threshold: Option<u32>,
 }
 
 impl HandshakeConnection {
@@ -62,6 +64,7 @@ impl HandshakeConnection {
         LoginConnection {
             flow: self.flow,
             stream: self.stream,
+            compression_threshold: None,
         }
     }
 
@@ -73,7 +76,7 @@ impl HandshakeConnection {
     }
 
     pub async fn read(&mut self) -> Result<HandshakePacket, String> {
-        read_packet::<HandshakePacket>(&self.flow, &mut self.stream).await
+        read_packet::<HandshakePacket>(&self.flow, &mut self.stream, None).await
     }
 
     /// Write a packet to the server
@@ -84,7 +87,7 @@ impl HandshakeConnection {
 
 impl GameConnection {
     pub async fn read(&mut self) -> Result<GamePacket, String> {
-        read_packet::<GamePacket>(&self.flow, &mut self.stream).await
+        read_packet::<GamePacket>(&self.flow, &mut self.stream, self.compression_threshold).await
     }
 
     /// Write a packet to the server
@@ -95,7 +98,7 @@ impl GameConnection {
 
 impl StatusConnection {
     pub async fn read(&mut self) -> Result<StatusPacket, String> {
-        read_packet::<StatusPacket>(&self.flow, &mut self.stream).await
+        read_packet::<StatusPacket>(&self.flow, &mut self.stream, None).await
     }
 
     /// Write a packet to the server
@@ -106,11 +109,28 @@ impl StatusConnection {
 
 impl LoginConnection {
     pub async fn read(&mut self) -> Result<LoginPacket, String> {
-        read_packet::<LoginPacket>(&self.flow, &mut self.stream).await
+        read_packet::<LoginPacket>(&self.flow, &mut self.stream, self.compression_threshold).await
     }
 
     /// Write a packet to the server
     pub async fn write(&mut self, packet: LoginPacket) {
         write_packet(packet, &mut self.stream).await;
     }
+
+    pub fn set_compression_threshold(&mut self, threshold: i32) {
+        // if you pass a threshold of 0 or less, compression is disabled
+        if threshold > 0 {
+            self.compression_threshold = Some(threshold as u32);
+        } else {
+            self.compression_threshold = None;
+        }
+    }
+
+    pub fn game(self) -> GameConnection {
+        GameConnection {
+            flow: self.flow,
+            stream: self.stream,
+            compression_threshold: self.compression_threshold,
+        }
+    }
 }
diff --git a/azalea-protocol/src/mc_buf.rs b/azalea-protocol/src/mc_buf.rs
index 0a47f637..84c602ec 100644
--- a/azalea-protocol/src/mc_buf.rs
+++ b/azalea-protocol/src/mc_buf.rs
@@ -134,6 +134,8 @@ impl Writable for Vec<u8> {
 pub trait Readable {
     async fn read_int_id_list(&mut self) -> Result<Vec<i32>, String>;
     async fn read_varint(&mut self) -> Result<i32, String>;
+    fn get_varint_size(&mut self, value: i32) -> u8;
+    fn get_varlong_size(&mut self, value: i32) -> u8;
     async fn read_byte_array(&mut self) -> Result<Vec<u8>, String>;
     async fn read_bytes(&mut self, n: usize) -> Result<Vec<u8>, String>;
     async fn read_utf(&mut self) -> Result<String, String>;
@@ -173,6 +175,26 @@ where
         Ok(ans)
     }
 
+    fn get_varint_size(&mut self, value: i32) -> u8 {
+        for i in 1..5 {
+            if (value & -1 << i * 7) != 0 {
+                continue;
+            }
+            return i;
+        }
+        return 5;
+    }
+
+    fn get_varlong_size(&mut self, value: i32) -> u8 {
+        for i in 1..10 {
+            if (value & -1 << i * 7) != 0 {
+                continue;
+            }
+            return i;
+        }
+        return 10;
+    }
+
     async fn read_byte_array(&mut self) -> Result<Vec<u8>, String> {
         let length = self.read_varint().await? as usize;
         Ok(self.read_bytes(length).await?)
diff --git a/azalea-protocol/src/packets/game/ClientboundLoginPacket.rs b/azalea-protocol/src/packets/game/ClientboundLoginPacket.rs
new file mode 100644
index 00000000..ec869faa
--- /dev/null
+++ b/azalea-protocol/src/packets/game/ClientboundLoginPacket.rs
@@ -0,0 +1,53 @@
+use super::GamePacket;
+use crate::mc_buf::{Readable, Writable};
+use azalea_core::resource_location::ResourceLocation;
+use std::hash::Hash;
+use tokio::io::BufReader;
+
+#[derive(Hash, Clone, Debug)]
+pub struct ClientboundLoginPacket {
+	// private final int playerId;
+    // private final boolean hardcore;
+    // private final GameType gameType;
+    // @Nullable
+    // private final GameType previousGameType;
+    // private final Set<ResourceKey<Level>> levels;
+    // private final RegistryAccess.RegistryHolder registryHolder;
+    // private final DimensionType dimensionType;
+    // private final ResourceKey<Level> dimension;
+    // private final long seed;
+    // private final int maxPlayers;
+    // private final int chunkRadius;
+    // private final int simulationDistance;
+    // private final boolean reducedDebugInfo;
+    // private final boolean showDeathScreen;
+    // private final boolean isDebug;
+    // private final boolean isFlat;
+	
+}
+
+impl ClientboundLoginPacket {
+    pub fn get(self) -> GamePacket {
+        GamePacket::ClientboundLoginPacket(self)
+    }
+
+    pub fn write(&self, buf: &mut Vec<u8>) {
+        buf.write_varint(self.transaction_id as i32).unwrap();
+        buf.write_utf(self.identifier.to_string().as_str()).unwrap();
+        buf.write_bytes(&self.data).unwrap();
+    }
+
+    pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
+        buf: &mut T,
+    ) -> Result<GamePacket, String> {
+        let transaction_id = buf.read_varint().await? as u32;
+        let identifier = ResourceLocation::new(&buf.read_utf().await?)?;
+        let data = buf.read_bytes(1048576).await?;
+        Ok(ClientboundLoginPacket {
+            transaction_id,
+            identifier,
+            data,
+        }
+        .get())
+    }
+}
diff --git a/azalea-protocol/src/packets/game/mod.rs b/azalea-protocol/src/packets/game/mod.rs
index a3ef2541..0391a6b1 100644
--- a/azalea-protocol/src/packets/game/mod.rs
+++ b/azalea-protocol/src/packets/game/mod.rs
@@ -22,7 +22,7 @@ impl ProtocolPacket for GamePacket {
     async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
         _id: u32,
         flow: &PacketFlow,
-        _buf: &mut BufReader<T>,
+        _buf: &mut T,
     ) -> Result<GamePacket, String>
     where
         Self: Sized,
diff --git a/azalea-protocol/src/packets/handshake/client_intention_packet.rs b/azalea-protocol/src/packets/handshake/client_intention_packet.rs
index 868626b3..5b50c7cc 100644
--- a/azalea-protocol/src/packets/handshake/client_intention_packet.rs
+++ b/azalea-protocol/src/packets/handshake/client_intention_packet.rs
@@ -28,7 +28,7 @@ impl ClientIntentionPacket {
     }
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        _buf: &mut BufReader<T>,
+        _buf: &mut T,
     ) -> Result<HandshakePacket, String> {
         Err("ClientIntentionPacket::parse not implemented".to_string())
         // Ok(ClientIntentionPacket {}.get())
diff --git a/azalea-protocol/src/packets/handshake/mod.rs b/azalea-protocol/src/packets/handshake/mod.rs
index 01010e1e..70e1a90d 100644
--- a/azalea-protocol/src/packets/handshake/mod.rs
+++ b/azalea-protocol/src/packets/handshake/mod.rs
@@ -33,7 +33,7 @@ impl ProtocolPacket for HandshakePacket {
     async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
         id: u32,
         flow: &PacketFlow,
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<HandshakePacket, String>
     where
         Self: Sized,
diff --git a/azalea-protocol/src/packets/login/clientbound_custom_query_packet.rs b/azalea-protocol/src/packets/login/clientbound_custom_query_packet.rs
index 2c66bfa3..ed9820ef 100644
--- a/azalea-protocol/src/packets/login/clientbound_custom_query_packet.rs
+++ b/azalea-protocol/src/packets/login/clientbound_custom_query_packet.rs
@@ -23,7 +23,7 @@ impl ClientboundCustomQueryPacket {
     }
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<LoginPacket, String> {
         let transaction_id = buf.read_varint().await? as u32;
         let identifier = ResourceLocation::new(&buf.read_utf().await?)?;
diff --git a/azalea-protocol/src/packets/login/clientbound_game_profile_packet.rs b/azalea-protocol/src/packets/login/clientbound_game_profile_packet.rs
index 04ba5369..da51f115 100644
--- a/azalea-protocol/src/packets/login/clientbound_game_profile_packet.rs
+++ b/azalea-protocol/src/packets/login/clientbound_game_profile_packet.rs
@@ -23,7 +23,7 @@ impl ClientboundGameProfilePacket {
     }
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<LoginPacket, String> {
         let uuid = Uuid::from_int_array([
             buf.read_int().await? as u32,
diff --git a/azalea-protocol/src/packets/login/clientbound_hello_packet.rs b/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
index 36a48706..46ca1301 100644
--- a/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
+++ b/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
@@ -22,7 +22,7 @@ impl ClientboundHelloPacket {
     }
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<LoginPacket, String> {
         let server_id = buf.read_utf_with_len(20).await?;
         let public_key = buf.read_byte_array().await?;
diff --git a/azalea-protocol/src/packets/login/clientbound_login_compression_packet.rs b/azalea-protocol/src/packets/login/clientbound_login_compression_packet.rs
index 53c6c9e1..e5009985 100644
--- a/azalea-protocol/src/packets/login/clientbound_login_compression_packet.rs
+++ b/azalea-protocol/src/packets/login/clientbound_login_compression_packet.rs
@@ -20,7 +20,7 @@ impl ClientboundLoginCompressionPacket {
     }
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<LoginPacket, String> {
         let compression_threshold = buf.read_varint().await?;
 
diff --git a/azalea-protocol/src/packets/login/mod.rs b/azalea-protocol/src/packets/login/mod.rs
index 7fee684a..377a285a 100644
--- a/azalea-protocol/src/packets/login/mod.rs
+++ b/azalea-protocol/src/packets/login/mod.rs
@@ -51,7 +51,7 @@ impl ProtocolPacket for LoginPacket {
     async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
         id: u32,
         flow: &PacketFlow,
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<LoginPacket, String>
     where
         Self: Sized,
diff --git a/azalea-protocol/src/packets/login/serverbound_hello_packet.rs b/azalea-protocol/src/packets/login/serverbound_hello_packet.rs
index 32a6dadc..011cc590 100644
--- a/azalea-protocol/src/packets/login/serverbound_hello_packet.rs
+++ b/azalea-protocol/src/packets/login/serverbound_hello_packet.rs
@@ -20,7 +20,7 @@ impl ServerboundHelloPacket {
     }
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        _buf: &mut BufReader<T>,
+        _buf: &mut T,
     ) -> Result<LoginPacket, String> {
         Err("ServerboundHelloPacket::read not implemented".to_string())
     }
diff --git a/azalea-protocol/src/packets/mod.rs b/azalea-protocol/src/packets/mod.rs
index a074b570..08c94509 100644
--- a/azalea-protocol/src/packets/mod.rs
+++ b/azalea-protocol/src/packets/mod.rs
@@ -38,109 +38,10 @@ where
     async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
         id: u32,
         flow: &PacketFlow,
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<Self, String>
     where
         Self: Sized;
 
     fn write(&self, buf: &mut Vec<u8>);
 }
-
-// impl Packet {
-//     fn get_inner_packet(&self) -> &dyn PacketTrait {
-//         match self {
-//             Packet::ClientIntentionPacket(packet) => packet,
-//             Packet::ServerboundStatusRequestPacket(packet) => packet,
-//             Packet::ClientboundStatusResponsePacket(packet) => packet,
-//             Packet::ServerboundHelloPacket(packet) => packet,
-//             Packet::ClientboundHelloPacket(packet) => packet,
-//         }
-//     }
-
-//     pub fn id(&self) -> u32 {
-//         match self {
-//             Packet::ClientIntentionPacket(_packet) => 0x00,
-//             Packet::ServerboundStatusRequestPacket(_packet) => 0x00,
-//             Packet::ClientboundStatusResponsePacket(_packet) => 0x00,
-//             Packet::ServerboundHelloPacket(_packet) => 0x00,
-//             Packet::ClientboundHelloPacket(_packet) => 0x01,
-//         }
-//     }
-
-//     /// Read a packet by its id, ConnectionProtocol, and flow
-//     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-//         id: u32,
-//         protocol: &ConnectionProtocol,
-//         flow: &PacketFlow,
-//         buf: &mut BufReader<T>,
-//     ) -> Result<Packet, String> {
-//         match protocol {
-//             ConnectionProtocol::Handshake => match flow {
-//                 PacketFlow::ClientToServer => match id {
-//                     0x00 => Ok(
-//                         handshake::client_intention_packet::ClientIntentionPacket::read(buf).await?,
-//                     ),
-//                     _ => Err(format!("Unknown ClientToServer handshake packet id: {}", id)),
-//                 }
-//                 PacketFlow::ServerToClient => Err("ServerToClient handshake packets not implemented".to_string()),
-//             },
-
-//             ConnectionProtocol::Game => Err("Game protocol not implemented yet".to_string()),
-
-//             ConnectionProtocol::Status => match flow {
-//                 PacketFlow::ServerToClient => match id {
-//                     0x00 => Ok(
-//                         status::clientbound_status_response_packet::ClientboundStatusResponsePacket
-//                             ::read(buf)
-//                             .await?,
-//                     ),
-//                     _ => Err(format!("Unknown ServerToClient status packet id: {}", id)),
-//                 },
-//                 PacketFlow::ClientToServer => match id {
-//                     0x00 => Ok(
-//                         status::serverbound_status_request_packet::ServerboundStatusRequestPacket
-//                             ::read(buf)
-//                             .await?,
-//                     ),
-//                     _ => Err(format!("Unknown ClientToServer status packet id: {}", id)),
-//                 },
-//             },
-
-//             ConnectionProtocol::Login => match flow {
-//                 PacketFlow::ServerToClient => match id {
-//                     0x01 => Ok(
-//                         login::clientbound_hello_packet::ClientboundHelloPacket::read(buf).await?,
-//                     ),
-//                     _ => Err(format!("Unknown ServerToClient login packet id: {}", id)),
-//                 },
-//                 PacketFlow::ClientToServer => match id {
-//                     0x00 => Ok(
-//                         login::serverbound_hello_packet::ServerboundHelloPacket::read(buf).await?,
-//                     ),
-//                     _ => Err(format!("Unknown ClientToServer login packet id: {}", id)),
-//                 },
-//             },
-//         }
-//     }
-
-//     pub fn write(&self, buf: &mut Vec<u8>) {
-//         self.get_inner_packet().write(buf);
-//     }
-// }
-
-// #[async_trait]
-// pub trait PacketTrait
-// where
-//     Self: Sized,
-// {
-//     /// Return a version of the packet that you can actually use for stuff
-//     fn get(self) -> dyn ProtocolPacket;
-
-//     fn write(&self, buf: &mut Vec<u8>);
-
-//     async fn read<T: AsyncRead + std::marker::Unpin + std::marker::Send, P: ProtocolPacket>(
-//         buf: &mut BufReader<T>,
-//     ) -> Result<P, String>
-//     where
-//         Self: Sized;
-// }
diff --git a/azalea-protocol/src/packets/status/clientbound_status_response_packet.rs b/azalea-protocol/src/packets/status/clientbound_status_response_packet.rs
index 1d8a3aa4..35f913ff 100644
--- a/azalea-protocol/src/packets/status/clientbound_status_response_packet.rs
+++ b/azalea-protocol/src/packets/status/clientbound_status_response_packet.rs
@@ -43,7 +43,7 @@ impl ClientboundStatusResponsePacket {
     pub fn write(&self, _buf: &mut Vec<u8>) {}
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<StatusPacket, String> {
         let status_string = buf.read_utf().await?;
         let status_json: Value =
diff --git a/azalea-protocol/src/packets/status/mod.rs b/azalea-protocol/src/packets/status/mod.rs
index ac6a34e1..9531111a 100644
--- a/azalea-protocol/src/packets/status/mod.rs
+++ b/azalea-protocol/src/packets/status/mod.rs
@@ -41,7 +41,7 @@ impl ProtocolPacket for StatusPacket {
     async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
         id: u32,
         flow: &PacketFlow,
-        buf: &mut BufReader<T>,
+        buf: &mut T,
     ) -> Result<StatusPacket, String>
     where
         Self: Sized,
diff --git a/azalea-protocol/src/packets/status/serverbound_status_request_packet.rs b/azalea-protocol/src/packets/status/serverbound_status_request_packet.rs
index 6a58da1f..dce9b93a 100644
--- a/azalea-protocol/src/packets/status/serverbound_status_request_packet.rs
+++ b/azalea-protocol/src/packets/status/serverbound_status_request_packet.rs
@@ -16,7 +16,7 @@ impl ServerboundStatusRequestPacket {
     }
 
     pub async fn read<T: tokio::io::AsyncRead + std::marker::Unpin + std::marker::Send>(
-        _buf: &mut BufReader<T>,
+        _buf: &mut T,
     ) -> Result<StatusPacket, String> {
         Err("ServerboundStatusRequestPacket::read not implemented".to_string())
     }
diff --git a/azalea-protocol/src/read.rs b/azalea-protocol/src/read.rs
index 704774b8..3e63ccc6 100644
--- a/azalea-protocol/src/read.rs
+++ b/azalea-protocol/src/read.rs
@@ -1,10 +1,14 @@
-use tokio::{io::BufReader, net::TcpStream};
-
 use crate::{connect::PacketFlow, mc_buf::Readable, packets::ProtocolPacket};
+use async_compression::tokio::bufread::ZlibDecoder;
+use tokio::{
+    io::{AsyncReadExt, BufReader},
+    net::TcpStream,
+};
 
 pub async fn read_packet<P: ProtocolPacket>(
     flow: &PacketFlow,
     stream: &mut TcpStream,
+    compression_threshold: Option<u32>,
 ) -> Result<P, String> {
     // what this does:
     // 1. reads the first 5 bytes, probably only some of this will be used to get the packet length
@@ -15,14 +19,64 @@ pub async fn read_packet<P: ProtocolPacket>(
     // the first thing minecraft sends us is the length as a varint, which can be up to 5 bytes long
     let mut buf = BufReader::with_capacity(4 * 1024 * 1024, stream);
 
-    let _packet_size = buf.read_varint().await?;
+    // Packet Length
+    let packet_size = buf.read_varint().await?;
+
+    // if there's no compression, we can just read the rest of the packet normally
+    if compression_threshold.is_none() {
+        // then, minecraft tells us the packet id as a varint
+        let packet_id = buf.read_varint().await?;
+
+        // if we recognize the packet id, parse it
+
+        println!("reading uncompressed packet id: {}", packet_id);
+        let packet = P::read(packet_id.try_into().unwrap(), flow, &mut buf).await?;
+
+        return Ok(packet);
+    }
+
+    println!("compressed packet size: {}", packet_size);
+
+    // there's compression
+    // Data Length
+    let data_size = buf.read_varint().await?;
+    println!("data size: {}", data_size);
+
+    // this packet has no compression
+    if data_size == 0 {
+        // Packet ID
+        let packet_id = buf.read_varint().await?;
+        println!(
+            "reading compressed packet without compression packet id: {}",
+            packet_id
+        );
+        let packet = P::read(packet_id.try_into().unwrap(), flow, &mut buf).await?;
+        return Ok(packet);
+    }
+
+    // this packet has compression
+    let packet_size_varint_size = buf.get_varint_size(packet_size);
+
+    let mut compressed_data = vec![0; packet_size as usize - packet_size_varint_size as usize];
+    buf.read_exact(compressed_data.as_mut_slice())
+        .await
+        .expect("Not enough compressed data");
 
-    // then, minecraft tells us the packet id as a varint
-    let packet_id = buf.read_varint().await?;
+    let mut z = ZlibDecoder::new(compressed_data.as_slice());
 
-    // if we recognize the packet id, parse it
+    // Packet ID
+    let packet_id = z.read_varint().await.unwrap();
+    println!("reading compressed packet id: {}", packet_id);
 
-    let packet = P::read(packet_id.try_into().unwrap(), flow, &mut buf).await?;
+    if let Ok(packet) = P::read(packet_id as u32, flow, &mut z).await {
+        Ok(packet)
+    } else {
+        // read the rest of the bytes
+        let packet_id_varint_size = z.get_varint_size(packet_id);
+        let mut buf = vec![0; packet_size as usize - packet_id_varint_size as usize];
+        z.read_exact(buf.as_mut_slice()).await.unwrap();
+        println!("{:?}", buf);
 
-    Ok(packet)
+        Err(format!("Error on packet id: {}", packet_id))
+    }
 }