src/encrypt.rs

@@ -4,7 +4,8 @@ use rsa::PublicKey;
 use rsa::{RsaPrivateKey, RsaPublicKey, PaddingScheme, errors::Result};
 use rand::rngs::OsRng;
 use aes::Aes128;
-use aes::cipher::{BlockEncrypt, BlockDecrypt, generic_array::GenericArray};
+use aes::cipher::{
+    BlockEncrypt, BlockDecrypt, NewBlockCipher, generic_array::GenericArray};
 
 pub fn generate_rsa_keys() -> Result<RsaPrivateKey> {
     let mut rng = OsRng;
@@ -30,13 +31,15 @@ pub fn decrypt_rsa(
     private_key.decrypt(padding, data)
 }
 
-pub fn encrypt_aes(cipher: Aes128, data: &[u8; 16]) -> Vec<u8> {
+pub fn encrypt_aes(key: &[u8; 16], data: &[u8; 16]) -> Vec<u8> {
+    let cipher = Aes128::new(GenericArray::from_slice(key));
     let mut block = GenericArray::clone_from_slice(data);
     cipher.encrypt_block(&mut block);
     block.to_vec()
 }
 
-pub fn decrypt_aes(cipher: Aes128, data: &[u8; 16]) -> Vec<u8> {
+pub fn decrypt_aes(key: &[u8; 16], data: &[u8; 16]) -> Vec<u8> {
+    let cipher = Aes128::new(GenericArray::from_slice(key));
     let mut block = GenericArray::clone_from_slice(data);
     cipher.decrypt_block(&mut block);
     block.to_vec()

src/mc_types.rs

@@ -4,14 +4,13 @@ use std::error::Error;
 use std::fmt;
 
 use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
+use tokio::net::TcpStream;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use serde::{Serialize, Deserialize};
 use async_trait::async_trait;
 use rsa::{RsaPrivateKey, RsaPublicKey};
 use rsa::pkcs8::{EncodePublicKey, DecodePublicKey};
 use rand::Rng;
-use aes::Aes128;
-use aes::cipher::{NewBlockCipher, generic_array::GenericArray};
 
 use crate::login;
 use crate::encrypt;
@@ -71,7 +70,6 @@ pub struct ProtocolConnection<'a> {
     rsa_private_key: Option<RsaPrivateKey>,
     rsa_public_key: Option<RsaPublicKey>,
     aes_encryption_key: Option<[u8; 16]>,
-    aes_cipher: Option<Aes128>,
     verify_token: Option<[u8; 16]>,
 }
 
@@ -86,7 +84,6 @@ impl<'a> ProtocolConnection<'a> {
             rsa_private_key: None,
             rsa_public_key: None,
             aes_encryption_key: None,
-            aes_cipher: None,
             verify_token: None,
         }
     }
@@ -131,9 +128,7 @@ impl<'a> ProtocolConnection<'a> {
         self.rsa_public_key = Some(
             RsaPublicKey::from_public_key_der(&request.public_key)?);
         let mut rng = rand::thread_rng();
-        let aes_key = rng.gen();
-        self.aes_encryption_key = Some(aes_key);
-        self.aes_cipher = Some(Aes128::new(GenericArray::from_slice(&aes_key)));
+        self.aes_encryption_key = Some(rng.gen());
         match self.aes_encryption_key {
             Some(key) => {
                 match &self.rsa_public_key {
@@ -202,11 +197,11 @@ impl<'a> ProtocolConnection<'a> {
     ) -> Result<(WriteHaftProtocolConnection, ReadHaftProtocolConnection)> {
         Ok((WriteHaftProtocolConnection {
             stream_write: &mut self.stream_write,
-            aes_cipher: self.aes_cipher.clone(),
+            aes_encryption_key: self.aes_encryption_key.clone(),
         },
         ReadHaftProtocolConnection {
             stream_read: &mut self.stream_read,
-            aes_cipher: self.aes_cipher.clone(),
+            aes_encryption_key: self.aes_encryption_key.clone(),
         }))
     }
 }
@@ -216,12 +211,12 @@ unsafe impl<'a> Send for ProtocolConnection<'a> {}
 #[async_trait]
 impl<'a> ProtocolRead for ProtocolConnection<'a> {
     async fn read_data(&mut self) -> Result<Vec<u8>> {
-        match self.aes_cipher.clone() {
-            Some(aes_cipher) => {
+        match self.aes_encryption_key {
+            Some(aes_key) => {
                 let mut buffer: Vec<u8> = vec![0; 16];
                 self.stream_read.read_exact(&mut buffer).await?;
                 buffer = encrypt::decrypt_aes(
-                    aes_cipher, buffer[0..16].try_into().unwrap());
+                    &aes_key, buffer[0..16].try_into().unwrap());
                 let raw_length = read_var_int_vec(&mut buffer)?;
                 let length =
                     if (raw_length - buffer.len() as i32) % 16 == 0 {
@@ -256,8 +251,8 @@ impl<'a> ProtocolWrite for ProtocolConnection<'a> {
     async fn write_data(&mut self, data: &mut Vec<u8>) -> Result<()> {
         let mut out_data = convert_var_int(data.len() as i32);
         out_data.append(data);
-        match self.aes_cipher.clone() {
-            Some(aes_cipher) => {
+        match self.aes_encryption_key {
+            Some(aes_key) => {
                 let length =
                     if (data.len() as i32) % 16 == 0 {
                         (data.len() as i32) / 16
@@ -268,7 +263,7 @@ impl<'a> ProtocolWrite for ProtocolConnection<'a> {
                 for _ in 0..length {
                     let mut block: Vec<u8> = out_data[0..16].to_vec();
                     block = encrypt::encrypt_aes(
-                        aes_cipher.clone(), block[0..16].try_into().unwrap());
+                        &aes_key, block[0..16].try_into().unwrap());
                     self.stream_write.write_all(&block).await?;
                 }
 
@@ -285,7 +280,7 @@ impl<'a> ProtocolWrite for ProtocolConnection<'a> {
 
 pub struct WriteHaftProtocolConnection<'a> {
     pub stream_write: &'a mut OwnedWriteHalf,
-    aes_cipher: Option<Aes128>,
+    aes_encryption_key: Option<[u8; 16]>,
 }
 
 impl<'a> WriteHaftProtocolConnection<'a> {
@@ -294,7 +289,7 @@ impl<'a> WriteHaftProtocolConnection<'a> {
     ) -> Self {
         WriteHaftProtocolConnection {
             stream_write,
-            aes_cipher: None,
+            aes_encryption_key: None,
         }
     }
 }
@@ -309,8 +304,8 @@ impl<'a> ProtocolWrite for WriteHaftProtocolConnection<'a> {
     ) -> Result<()> {
         let mut out_data = convert_var_int(data.len() as i32);
         out_data.append(data);
-        match self.aes_cipher.clone() {
-            Some(aes_cipher) => {
+        match self.aes_encryption_key {
+            Some(aes_key) => {
                 let length =
                     if (data.len() as i32) % 16 == 0 {
                         (data.len() as i32) / 16
@@ -321,7 +316,7 @@ impl<'a> ProtocolWrite for WriteHaftProtocolConnection<'a> {
                 for _ in 0..length {
                     let mut block: Vec<u8> = out_data[0..16].to_vec();
                     block = encrypt::encrypt_aes(
-                        aes_cipher.clone(), block[0..16].try_into().unwrap());
+                        &aes_key, block[0..16].try_into().unwrap());
                     self.stream_write.write_all(&block).await?;
                 }
 
@@ -339,7 +334,7 @@ impl<'a> ProtocolWrite for WriteHaftProtocolConnection<'a> {
 
 pub struct ReadHaftProtocolConnection<'a> {
     pub stream_read: &'a mut OwnedReadHalf,
-    aes_cipher: Option<Aes128>,
+    aes_encryption_key: Option<[u8; 16]>,
 }
 
 impl<'a> ReadHaftProtocolConnection<'a> {
@@ -348,7 +343,7 @@ impl<'a> ReadHaftProtocolConnection<'a> {
     ) -> Self {
         ReadHaftProtocolConnection {
             stream_read,
-            aes_cipher: None,
+            aes_encryption_key: None,
         }
     }
 
@@ -370,12 +365,12 @@ unsafe impl<'a> Send for ReadHaftProtocolConnection<'a> {}
 #[async_trait]
 impl<'a> ProtocolRead for ReadHaftProtocolConnection<'a> {
     async fn read_data(&mut self) -> Result<Vec<u8>> {
-        match self.aes_cipher.clone() {
-            Some(aes_cipher) => {
+        match self.aes_encryption_key {
+            Some(aes_key) => {
                 let mut buffer: Vec<u8> = vec![0; 16];
                 self.stream_read.read_exact(&mut buffer).await?;
                 buffer = encrypt::decrypt_aes(
-                    aes_cipher, buffer[0..16].try_into().unwrap());
+                    &aes_key, buffer[0..16].try_into().unwrap());
                 let raw_length = read_var_int_vec(&mut buffer)?;
                 let length =
                     if (raw_length - buffer.len() as i32) % 16 == 0 {