From f4dd3a9293367fa8f3d839a184e8055b22595204 Mon Sep 17 00:00:00 2001
From: mat <github@matdoes.dev>
Date: Sun, 24 Apr 2022 22:46:41 -0500
Subject: ENCRYPTION WORKS!!!!!!!!!!!

---
 azalea-protocol/src/connect.rs | 52 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 43 insertions(+), 9 deletions(-)

(limited to 'azalea-protocol/src/connect.rs')

diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index 3d910d3a..e81bc368 100755
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -7,6 +7,7 @@ use crate::packets::status::StatusPacket;
 use crate::read::read_packet;
 use crate::write::write_packet;
 use crate::ServerIpAddress;
+use azalea_auth::encryption::Aes128Cfb;
 use tokio::net::TcpStream;
 
 pub enum PacketFlow {
@@ -25,6 +26,7 @@ pub struct GameConnection {
     /// The buffered writer
     pub stream: TcpStream,
     pub compression_threshold: Option<u32>,
+    pub cipher: Option<Aes128Cfb>,
 }
 
 pub struct StatusConnection {
@@ -38,6 +40,7 @@ pub struct LoginConnection {
     /// The buffered writer
     pub stream: TcpStream,
     pub compression_threshold: Option<u32>,
+    pub cipher: Option<Aes128Cfb>,
 }
 
 impl HandshakeConnection {
@@ -65,6 +68,7 @@ impl HandshakeConnection {
             flow: self.flow,
             stream: self.stream,
             compression_threshold: None,
+            cipher: None,
         }
     }
 
@@ -76,46 +80,69 @@ impl HandshakeConnection {
     }
 
     pub async fn read(&mut self) -> Result<HandshakePacket, String> {
-        read_packet::<HandshakePacket, _>(&self.flow, &mut self.stream, None).await
+        read_packet::<HandshakePacket, _>(&self.flow, &mut self.stream, None, &mut None).await
     }
 
     /// Write a packet to the server
     pub async fn write(&mut self, packet: HandshakePacket) {
-        write_packet(packet, &mut self.stream, None).await;
+        write_packet(packet, &mut self.stream, None, &mut None).await;
     }
 }
 
 impl GameConnection {
     pub async fn read(&mut self) -> Result<GamePacket, String> {
-        read_packet::<GamePacket, _>(&self.flow, &mut self.stream, self.compression_threshold).await
+        read_packet::<GamePacket, _>(
+            &self.flow,
+            &mut self.stream,
+            self.compression_threshold,
+            &mut self.cipher,
+        )
+        .await
     }
 
     /// Write a packet to the server
     pub async fn write(&mut self, packet: GamePacket) {
-        write_packet(packet, &mut self.stream, self.compression_threshold).await;
+        write_packet(
+            packet,
+            &mut self.stream,
+            self.compression_threshold,
+            &mut self.cipher,
+        )
+        .await;
     }
 }
 
 impl StatusConnection {
     pub async fn read(&mut self) -> Result<StatusPacket, String> {
-        read_packet::<StatusPacket, _>(&self.flow, &mut self.stream, None).await
+        read_packet::<StatusPacket, _>(&self.flow, &mut self.stream, None, &mut None).await
     }
 
     /// Write a packet to the server
     pub async fn write(&mut self, packet: StatusPacket) {
-        write_packet(packet, &mut self.stream, None).await;
+        write_packet(packet, &mut self.stream, None, &mut None).await;
     }
 }
 
 impl LoginConnection {
     pub async fn read(&mut self) -> Result<LoginPacket, String> {
-        read_packet::<LoginPacket, _>(&self.flow, &mut self.stream, self.compression_threshold)
-            .await
+        read_packet::<LoginPacket, _>(
+            &self.flow,
+            &mut self.stream,
+            self.compression_threshold,
+            &mut self.cipher,
+        )
+        .await
     }
 
     /// Write a packet to the server
     pub async fn write(&mut self, packet: LoginPacket) {
-        write_packet(packet, &mut self.stream, self.compression_threshold).await;
+        write_packet(
+            packet,
+            &mut self.stream,
+            self.compression_threshold,
+            &mut self.cipher,
+        )
+        .await;
     }
 
     pub fn set_compression_threshold(&mut self, threshold: i32) {
@@ -127,11 +154,18 @@ impl LoginConnection {
         }
     }
 
+    pub fn set_encryption_key(&mut self, key: [u8; 16]) {
+        // minecraft has a cipher decoder and encoder, i don't think it matters though?
+        let cipher = azalea_auth::encryption::create_cipher(&key);
+        self.cipher = Some(cipher);
+    }
+
     pub fn game(self) -> GameConnection {
         GameConnection {
             flow: self.flow,
             stream: self.stream,
             compression_threshold: self.compression_threshold,
+            cipher: self.cipher,
         }
     }
 }
-- 
cgit v1.2.3