ENCRYPTION WORKS!!!!!!!!!!!

9 files changed, 163 insertions, 31 deletions

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,
         }
     }
 }
diff --git a/azalea-protocol/src/mc_buf/mod.rs b/azalea-protocol/src/mc_buf/mod.rs
index 4ecb65d1..3ba6ac3e 100755
--- a/azalea-protocol/src/mc_buf/mod.rs
+++ b/azalea-protocol/src/mc_buf/mod.rs
@@ -5,11 +5,31 @@ mod write;
 
 pub use read::{McBufReadable, McBufVarintReadable, Readable};
 pub use write::{McBufVarintWritable, McBufWritable, Writable};
+use std::ops::Deref;
 
 // const DEFAULT_NBT_QUOTA: u32 = 2097152;
 const MAX_STRING_LENGTH: u16 = 32767;
 // const MAX_COMPONENT_STRING_LENGTH: u32 = 262144;
 
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct ByteArray(Vec<u8>);
+
+impl Deref for ByteArray {
+    type Target = Vec<u8>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl From<Vec<u8>> for ByteArray {
+    fn from(vec: Vec<u8>) -> Self {
+        Self(vec)
+    }
+}
+
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/azalea-protocol/src/mc_buf/read.rs b/azalea-protocol/src/mc_buf/read.rs
index 1e031916..3d50e5aa 100755
--- a/azalea-protocol/src/mc_buf/read.rs
+++ b/azalea-protocol/src/mc_buf/read.rs
@@ -5,6 +5,7 @@ use azalea_core::{
 };
 use serde::Deserialize;
 use tokio::io::{AsyncRead, AsyncReadExt};
+use crate::mc_buf::ByteArray;
 
 use super::MAX_STRING_LENGTH;
 
@@ -14,7 +15,7 @@ pub trait Readable {
     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_byte_array(&mut self) -> Result<ByteArray, String>;
     async fn read_bytes_with_len(&mut self, n: usize) -> Result<Vec<u8>, String>;
     async fn read_bytes(&mut self) -> Result<Vec<u8>, String>;
     async fn read_utf(&mut self) -> Result<String, String>;
@@ -80,9 +81,9 @@ where
         10
     }
 
-    async fn read_byte_array(&mut self) -> Result<Vec<u8>, String> {
+    async fn read_byte_array(&mut self) -> Result<ByteArray, String> {
         let length = self.read_varint().await? as usize;
-        Ok(self.read_bytes_with_len(length).await?)
+        Ok(ByteArray(self.read_bytes_with_len(length).await?))
     }
 
     async fn read_bytes_with_len(&mut self, n: usize) -> Result<Vec<u8>, String> {
@@ -251,6 +252,17 @@ impl McBufReadable for Vec<u8> {
     }
 }
 
+
+#[async_trait]
+impl McBufReadable for ByteArray {
+    async fn read_into<R>(buf: &mut R) -> Result<Self, String>
+    where
+        R: AsyncRead + std::marker::Unpin + std::marker::Send,
+    {
+        buf.read_byte_array().await
+    }
+}
+
 // string
 #[async_trait]
 impl McBufReadable for String {
diff --git a/azalea-protocol/src/mc_buf/write.rs b/azalea-protocol/src/mc_buf/write.rs
index 05f613d8..f1362402 100755
--- a/azalea-protocol/src/mc_buf/write.rs
+++ b/azalea-protocol/src/mc_buf/write.rs
@@ -1,3 +1,5 @@
+use super::MAX_STRING_LENGTH;
+use crate::mc_buf::ByteArray;
 use async_trait::async_trait;
 use azalea_chat::component::Component;
 use azalea_core::{
@@ -6,8 +8,6 @@ use azalea_core::{
 use byteorder::{BigEndian, WriteBytesExt};
 use std::io::Write;
 
-use super::MAX_STRING_LENGTH;
-
 #[async_trait]
 pub trait Writable {
     fn write_list<F, T>(&mut self, list: &[T], writer: F) -> Result<(), std::io::Error>
@@ -193,6 +193,12 @@ impl McBufWritable for Vec<u8> {
     }
 }
 
+impl McBufWritable for ByteArray {
+    fn write_into(&self, buf: &mut Vec<u8>) -> Result<(), std::io::Error> {
+        buf.write_byte_array(&self)
+    }
+}
+
 // string
 impl McBufWritable for String {
     fn write_into(&self, buf: &mut Vec<u8>) -> Result<(), std::io::Error> {
diff --git a/azalea-protocol/src/packets/login/clientbound_hello_packet.rs b/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
index 9d0cec39..d36f1335 100755
--- a/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
+++ b/azalea-protocol/src/packets/login/clientbound_hello_packet.rs
@@ -1,14 +1,13 @@
 use std::hash::Hash;
 
-use crate::mc_buf::Readable;
-
 use super::LoginPacket;
+use crate::mc_buf::{ByteArray, Readable};
 
 #[derive(Hash, Clone, Debug)]
 pub struct ClientboundHelloPacket {
     pub server_id: String,
-    pub public_key: Vec<u8>,
-    pub nonce: Vec<u8>,
+    pub public_key: ByteArray,
+    pub nonce: ByteArray,
 }
 
 impl ClientboundHelloPacket {
diff --git a/azalea-protocol/src/packets/login/mod.rs b/azalea-protocol/src/packets/login/mod.rs
index 93c1200f..b160a28c 100755
--- a/azalea-protocol/src/packets/login/mod.rs
+++ b/azalea-protocol/src/packets/login/mod.rs
@@ -15,9 +15,9 @@ declare_state_packets!(
         0x01: serverbound_key_packet::ServerboundKeyPacket,
     },
     Clientbound => {
-        // 0x00: clientbound_login_disconnect_packet::ClientboundLoginDisconnectPacket,
-        // for some reason this is used instead of 0x00??
-        0x1a: clientbound_login_disconnect_packet::ClientboundLoginDisconnectPacket,
+        0x00: clientbound_login_disconnect_packet::ClientboundLoginDisconnectPacket,
+        // sometimes this is used for some reason?
+        // 0x1a: clientbound_login_disconnect_packet::ClientboundLoginDisconnectPacket,
 
         0x01: clientbound_hello_packet::ClientboundHelloPacket,
         0x02: clientbound_game_profile_packet::ClientboundGameProfilePacket,
diff --git a/azalea-protocol/src/packets/login/serverbound_key_packet.rs b/azalea-protocol/src/packets/login/serverbound_key_packet.rs
index f402d357..3750331f 100644
--- a/azalea-protocol/src/packets/login/serverbound_key_packet.rs
+++ b/azalea-protocol/src/packets/login/serverbound_key_packet.rs
@@ -1,10 +1,10 @@
 use super::LoginPacket;
-use crate::mc_buf::Writable;
+use crate::mc_buf::{ByteArray, Writable};
 use packet_macros::LoginPacket;
 use std::hash::Hash;
 
 #[derive(Hash, Clone, Debug, LoginPacket)]
 pub struct ServerboundKeyPacket {
-    pub shared_secret: Vec<u8>,
-    pub nonce: Vec<u8>,
+    pub shared_secret: ByteArray,
+    pub nonce: ByteArray,
 }
diff --git a/azalea-protocol/src/read.rs b/azalea-protocol/src/read.rs
index ed36d6b6..be20ac23 100755
--- a/azalea-protocol/src/read.rs
+++ b/azalea-protocol/src/read.rs
@@ -1,13 +1,17 @@
+use std::{cell::Cell, pin::Pin};
+
 use crate::{connect::PacketFlow, mc_buf::Readable, packets::ProtocolPacket};
 use async_compression::tokio::bufread::ZlibDecoder;
+use azalea_auth::encryption::Aes128Cfb;
 use tokio::io::{AsyncRead, AsyncReadExt};
 
-async fn frame_splitter<R>(stream: &mut R) -> Result<Vec<u8>, String>
+async fn frame_splitter<R: ?Sized>(mut stream: &mut R) -> Result<Vec<u8>, String>
 where
     R: AsyncRead + std::marker::Unpin + std::marker::Send,
 {
     // Packet Length
     let length_result = stream.read_varint().await;
+    println!("length_result: {:?}", length_result);
     match length_result {
         Ok(length) => {
             let mut buf = vec![0; length as usize];
@@ -17,6 +21,8 @@ where
                 .await
                 .map_err(|e| e.to_string())?;
 
+            println!("buf: {:?}", buf);
+
             Ok(buf)
         }
         Err(_) => Err("length wider than 21-bit".to_string()),
@@ -90,15 +96,61 @@ where
     Ok(decoded_buf)
 }
 
-pub async fn read_packet<P: ProtocolPacket, R>(
+struct EncryptedStream<'a, R>
+where
+    R: AsyncRead + std::marker::Unpin + std::marker::Send,
+{
+    cipher: Cell<&'a mut Option<Aes128Cfb>>,
+    stream: &'a mut Pin<&'a mut R>,
+}
+
+impl<R> AsyncRead for EncryptedStream<'_, R>
+where
+    R: AsyncRead + std::marker::Unpin + std::marker::Send,
+{
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> std::task::Poll<std::io::Result<()>> {
+        // i hate this
+        let polled = self.as_mut().stream.as_mut().poll_read(cx, buf);
+        match polled {
+            std::task::Poll::Ready(r) => {
+                println!("encrypted packet {:?}", buf.initialized_mut());
+                if let Some(cipher) = self.as_mut().cipher.get_mut() {
+                    azalea_auth::encryption::decrypt_packet(cipher, buf.initialized_mut());
+                    println!("decrypted packet {:?}", buf.initialized_mut());
+                }
+                match r {
+                    Ok(()) => std::task::Poll::Ready(Ok(())),
+                    Err(e) => panic!("{:?}", e),
+                }
+            }
+            std::task::Poll::Pending => {
+                return std::task::Poll::Pending;
+            }
+        }
+    }
+}
+
+pub async fn read_packet<'a, P: ProtocolPacket, R>(
     flow: &PacketFlow,
-    stream: &mut R,
+    stream: &'a mut R,
     compression_threshold: Option<u32>,
+    cipher: &mut Option<Aes128Cfb>,
 ) -> Result<P, String>
 where
-    R: AsyncRead + std::marker::Unpin + std::marker::Send,
+    R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync,
 {
-    let mut buf = frame_splitter(stream).await?;
+    // if we were given a cipher, decrypt the packet
+    let mut encrypted_stream = EncryptedStream {
+        cipher: Cell::new(cipher),
+        stream: &mut Pin::new(stream),
+    };
+
+    let mut buf = frame_splitter(&mut encrypted_stream).await?;
+
     if let Some(compression_threshold) = compression_threshold {
         buf = compression_decoder(&mut buf.as_slice(), compression_threshold).await?;
     }
diff --git a/azalea-protocol/src/write.rs b/azalea-protocol/src/write.rs
index 821ed6e8..e72a74b1 100755
--- a/azalea-protocol/src/write.rs
+++ b/azalea-protocol/src/write.rs
@@ -1,5 +1,6 @@
 use crate::{mc_buf::Writable, packets::ProtocolPacket, read::MAXIMUM_UNCOMPRESSED_LENGTH};
 use async_compression::tokio::bufread::ZlibEncoder;
+use azalea_auth::encryption::Aes128Cfb;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::TcpStream,
@@ -50,14 +51,22 @@ async fn compression_encoder(data: &[u8], compression_threshold: u32) -> Result<
     }
 }
 
-pub async fn write_packet<P>(packet: P, stream: &mut TcpStream, compression_threshold: Option<u32>)
-where
+pub async fn write_packet<P>(
+    packet: P,
+    stream: &mut TcpStream,
+    compression_threshold: Option<u32>,
+    cipher: &mut Option<Aes128Cfb>,
+) where
     P: ProtocolPacket + std::fmt::Debug,
 {
     let mut buf = packet_encoder(&packet).unwrap();
     if let Some(threshold) = compression_threshold {
         buf = compression_encoder(&buf, threshold).await.unwrap();
     }
+    // if we were given a cipher, encrypt the packet
+    if let Some(cipher) = cipher {
+        azalea_auth::encryption::encrypt_packet(cipher, &mut buf);
+    }
     buf = frame_prepender(&mut buf).unwrap();
     stream.write_all(&buf).await.unwrap();
 }