upgrade aes and cfb8

3 files changed, 27 insertions, 20 deletions

diff --git a/azalea-protocol/src/connect.rs b/azalea-protocol/src/connect.rs
index e81bc368..c55f2e90 100755
--- a/azalea-protocol/src/connect.rs
+++ b/azalea-protocol/src/connect.rs
@@ -7,7 +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 azalea_auth::encryption::{Aes128CfbDec, Aes128CfbEnc};
 use tokio::net::TcpStream;
 
 pub enum PacketFlow {
@@ -26,7 +26,8 @@ pub struct GameConnection {
     /// The buffered writer
     pub stream: TcpStream,
     pub compression_threshold: Option<u32>,
-    pub cipher: Option<Aes128Cfb>,
+    pub enc_cipher: Option<Aes128CfbEnc>,
+    pub dec_cipher: Option<Aes128CfbDec>,
 }
 
 pub struct StatusConnection {
@@ -40,7 +41,8 @@ pub struct LoginConnection {
     /// The buffered writer
     pub stream: TcpStream,
     pub compression_threshold: Option<u32>,
-    pub cipher: Option<Aes128Cfb>,
+    pub enc_cipher: Option<Aes128CfbEnc>,
+    pub dec_cipher: Option<Aes128CfbDec>,
 }
 
 impl HandshakeConnection {
@@ -68,7 +70,8 @@ impl HandshakeConnection {
             flow: self.flow,
             stream: self.stream,
             compression_threshold: None,
-            cipher: None,
+            enc_cipher: None,
+            dec_cipher: None,
         }
     }
 
@@ -95,7 +98,7 @@ impl GameConnection {
             &self.flow,
             &mut self.stream,
             self.compression_threshold,
-            &mut self.cipher,
+            &mut self.dec_cipher,
         )
         .await
     }
@@ -106,7 +109,7 @@ impl GameConnection {
             packet,
             &mut self.stream,
             self.compression_threshold,
-            &mut self.cipher,
+            &mut self.enc_cipher,
         )
         .await;
     }
@@ -129,7 +132,7 @@ impl LoginConnection {
             &self.flow,
             &mut self.stream,
             self.compression_threshold,
-            &mut self.cipher,
+            &mut self.dec_cipher,
         )
         .await
     }
@@ -140,7 +143,7 @@ impl LoginConnection {
             packet,
             &mut self.stream,
             self.compression_threshold,
-            &mut self.cipher,
+            &mut self.enc_cipher,
         )
         .await;
     }
@@ -156,8 +159,9 @@ 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);
+        let (enc_cipher, dec_cipher) = azalea_auth::encryption::create_cipher(&key);
+        self.enc_cipher = Some(enc_cipher);
+        self.dec_cipher = Some(dec_cipher);
     }
 
     pub fn game(self) -> GameConnection {
@@ -165,7 +169,8 @@ impl LoginConnection {
             flow: self.flow,
             stream: self.stream,
             compression_threshold: self.compression_threshold,
-            cipher: self.cipher,
+            enc_cipher: self.enc_cipher,
+            dec_cipher: self.dec_cipher,
         }
     }
 }
diff --git a/azalea-protocol/src/read.rs b/azalea-protocol/src/read.rs
index 1e308cfb..88fe95cf 100755
--- a/azalea-protocol/src/read.rs
+++ b/azalea-protocol/src/read.rs
@@ -2,7 +2,7 @@ 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 azalea_auth::encryption::Aes128CfbDec;
 use tokio::io::{AsyncRead, AsyncReadExt};
 
 async fn frame_splitter<R: ?Sized>(mut stream: &mut R) -> Result<Vec<u8>, String>
@@ -97,7 +97,7 @@ struct EncryptedStream<'a, R>
 where
     R: AsyncRead + std::marker::Unpin + std::marker::Send,
 {
-    cipher: Cell<&'a mut Option<Aes128Cfb>>,
+    cipher: Cell<&'a mut Option<Aes128CfbDec>>,
     stream: &'a mut Pin<&'a mut R>,
 }
 
@@ -133,7 +133,7 @@ pub async fn read_packet<'a, P: ProtocolPacket, R>(
     flow: &PacketFlow,
     stream: &'a mut R,
     compression_threshold: Option<u32>,
-    cipher: &mut Option<Aes128Cfb>,
+    cipher: &mut Option<Aes128CfbDec>,
 ) -> Result<P, String>
 where
     R: AsyncRead + std::marker::Unpin + std::marker::Send + std::marker::Sync,
diff --git a/azalea-protocol/src/write.rs b/azalea-protocol/src/write.rs
index e72a74b1..e39ce18e 100755
--- a/azalea-protocol/src/write.rs
+++ b/azalea-protocol/src/write.rs
@@ -1,8 +1,9 @@
 use crate::{mc_buf::Writable, packets::ProtocolPacket, read::MAXIMUM_UNCOMPRESSED_LENGTH};
 use async_compression::tokio::bufread::ZlibEncoder;
-use azalea_auth::encryption::Aes128Cfb;
+use azalea_auth::encryption::Aes128CfbEnc;
+use std::fmt::Debug;
 use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
+    io::{AsyncReadExt, AsyncWrite, AsyncWriteExt},
     net::TcpStream,
 };
 
@@ -51,13 +52,14 @@ async fn compression_encoder(data: &[u8], compression_threshold: u32) -> Result<
     }
 }
 
-pub async fn write_packet<P>(
+pub async fn write_packet<P, W>(
     packet: P,
-    stream: &mut TcpStream,
+    stream: &mut W,
     compression_threshold: Option<u32>,
-    cipher: &mut Option<Aes128Cfb>,
+    cipher: &mut Option<Aes128CfbEnc>,
 ) where
-    P: ProtocolPacket + std::fmt::Debug,
+    P: ProtocolPacket + Debug,
+    W: AsyncWrite + Unpin + Send,
 {
     let mut buf = packet_encoder(&packet).unwrap();
     if let Some(threshold) = compression_threshold {