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

13 files changed, 208 insertions, 52 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 6c5ebad7..8be7aa31 100755
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10,13 +10,14 @@ checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
 
 [[package]]
 name = "aes"
-version = "0.8.1"
+version = "0.7.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bfe0133578c0986e1fe3dfcd4af1cc5b2dd6c3dbf534d69916ce16a2701d40ba"
+checksum = "9e8b47f52ea9bae42228d07ec09eb676433d7c4ed1ebdf0f1d1c29ed446f1ab8"
 dependencies = [
  "cfg-if",
  "cipher",
  "cpufeatures",
+ "opaque-debug",
 ]
 
 [[package]]
@@ -218,9 +219,9 @@ dependencies = [
 
 [[package]]
 name = "cfb8"
-version = "0.8.1"
+version = "0.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "014c0a0e1ad0dae6a86c082db2f9bd7fe8c2c734227047d0d8b4d4a3a094a1e1"
+checksum = "c3a4b6c43bf284e617a659ce5dc149676680530a3a4a9bb6b278d1a9ed5b229d"
 dependencies = [
  "cipher",
 ]
@@ -243,12 +244,11 @@ dependencies = [
 
 [[package]]
 name = "cipher"
-version = "0.4.3"
+version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d1873270f8f7942c191139cb8a40fd228da6c3fd2fc376d7e92d47aa14aeb59e"
+checksum = "7ee52072ec15386f770805afd189a01c8841be8696bed250fa2f13c4c0d6dfb7"
 dependencies = [
- "crypto-common",
- "inout",
+ "generic-array",
 ]
 
 [[package]]
@@ -581,15 +581,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "inout"
-version = "0.1.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a0c10553d664a4d0bcff9f4215d0aac67a639cc68ef660840afe309b807bc9f5"
-dependencies = [
- "generic-array",
-]
-
-[[package]]
 name = "instant"
 version = "0.1.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -857,6 +848,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0ab1bc2a289d34bd04a330323ac98a1b4bc82c9d9fcb1e66b63caa84da26b575"
 
 [[package]]
+name = "opaque-debug"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "624a8340c38c1b80fd549087862da4ba43e08858af025b236e509b6649fc13d5"
+
+[[package]]
 name = "packet-macros"
 version = "0.1.0"
 dependencies = [
diff --git a/azalea-auth/Cargo.toml b/azalea-auth/Cargo.toml
index 5f62991b..6279d903 100755
--- a/azalea-auth/Cargo.toml
+++ b/azalea-auth/Cargo.toml
@@ -6,8 +6,8 @@ version = "0.1.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-aes = "0.8.1"
-cfb8 = "0.8.1"
+aes = "0.7.4"
+cfb8 = "0.7.1"
 num-bigint = "^0.4.3"
 rand = {version = "^0.8.4", features = ["getrandom"]}
 rsa_public_encrypt_pkcs1 = "0.4.0"
diff --git a/azalea-auth/src/encryption.rs b/azalea-auth/src/encryption.rs
index 72739620..6f2e43a8 100644
--- a/azalea-auth/src/encryption.rs
+++ b/azalea-auth/src/encryption.rs
@@ -1,3 +1,8 @@
+use aes::{
+    cipher::{AsyncStreamCipher, NewCipher},
+    Aes128,
+};
+use cfb8::Cfb8;
 use rand::{rngs::OsRng, RngCore};
 use sha1::{Digest, Sha1};
 
@@ -24,7 +29,9 @@ fn hex_digest(digest: &[u8]) -> String {
     num_bigint::BigInt::from_signed_bytes_be(digest).to_str_radix(16)
 }
 
+#[derive(Debug)]
 pub struct EncryptResult {
+    pub secret_key: [u8; 16],
     pub encrypted_public_key: Vec<u8>,
     pub encrypted_nonce: Vec<u8>,
 }
@@ -43,11 +50,26 @@ pub fn encrypt(public_key: &[u8], nonce: &[u8]) -> Result<EncryptResult, String>
     let encrypted_nonce: Vec<u8> = rsa_public_encrypt_pkcs1::encrypt(&public_key, &nonce)?;
 
     Ok(EncryptResult {
+        secret_key,
         encrypted_public_key,
         encrypted_nonce,
     })
 }
 
+// TODO: update the aes and cfb8 crates
+pub type Aes128Cfb = Cfb8<Aes128>;
+
+pub fn create_cipher(key: &[u8]) -> Aes128Cfb {
+    Aes128Cfb::new_from_slices(&key, &key).unwrap()
+}
+
+pub fn encrypt_packet(cipher: &mut Aes128Cfb, packet: &mut [u8]) {
+    cipher.encrypt(packet);
+}
+pub fn decrypt_packet(cipher: &mut Aes128Cfb, packet: &mut [u8]) {
+    cipher.decrypt(packet);
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/azalea-client/src/connect.rs b/azalea-client/src/connect.rs
index b402200a..44eff33f 100755
--- a/azalea-client/src/connect.rs
+++ b/azalea-client/src/connect.rs
@@ -43,14 +43,19 @@ pub async fn join_server(address: &ServerAddress) -> Result<(), String> {
                 LoginPacket::ClientboundHelloPacket(p) => {
                     println!("Got encryption request {:?} {:?}", p.nonce, p.public_key);
                     let e = azalea_auth::encryption::encrypt(&p.public_key, &p.nonce).unwrap();
+
+                    // TODO: authenticate with the server here (authenticateServer)
+                    println!("Sending encryption response {:?}", e);
+
                     conn.write(
                         ServerboundKeyPacket {
-                            nonce: e.encrypted_nonce,
-                            shared_secret: e.encrypted_public_key,
+                            nonce: e.encrypted_nonce.into(),
+                            shared_secret: e.encrypted_public_key.into(),
                         }
                         .get(),
                     )
                     .await;
+                    conn.set_encryption_key(e.secret_key);
                 }
                 LoginPacket::ClientboundLoginCompressionPacket(p) => {
                     println!("Got compression request {:?}", p.compression_threshold);
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();
 }