Optimized aes/cfb8

Cargo.toml

@@ -16,3 +16,6 @@ rand = "0.8.5"
 aes = "0.7"
 rsa = "0.6"
 pkcs8 = "0.8"
+crypto = { version = "0.5.1", features = ["digest"] }
+sha1 = "0.10.6"
+num-bigint = "0.4.5"

src/encrypt.rs

@@ -8,9 +8,9 @@ use aes::cipher::{BlockEncrypt, BlockDecrypt, generic_array::GenericArray};
 
 #[derive(Clone)]
 pub struct McCipher {
-    key: [u8; 16],
+    pub(crate) key: [u8; 16],
-    state_en: [u8; 16],
+    state_en: u128,
-    state_de: [u8; 16],
+    state_de: u128,
 }
 
 impl McCipher {
@@ -19,8 +19,8 @@ impl McCipher {
         let aes_key: [u8; 16] = rng.gen();
         Self {
             key: aes_key.clone(),
-            state_en: aes_key.clone(),
+            state_en: u128::from_be_bytes(aes_key),
-            state_de: aes_key.clone(),
+            state_de: u128::from_be_bytes(aes_key),
         }
     }
 
@@ -39,12 +39,12 @@ impl McCipher {
             .as_slice()[0..16].try_into().unwrap();
         Ok(Self {
             key: aes_key.clone(),
-            state_en: aes_key.clone(),
+            state_en: u128::from_be_bytes(aes_key),
-            state_de: aes_key.clone(),
+            state_de: u128::from_be_bytes(aes_key),
         })
     }
 
-    pub fn encrypt_aes(&self, data: Vec<u8>) -> Vec<u8> {
+    pub fn encrypt_aes(&mut self, data: Vec<u8>) -> Vec<u8> {
         let mut out_data = vec![0; data.len()];
         for i in 0..data.len() {
             out_data[i] = self.encrypt_block(data[i]);
@@ -52,7 +52,7 @@ impl McCipher {
         out_data
     }
 
-    pub fn decrypt_aes(&self, data: Vec<u8>) -> Vec<u8> {
+    pub fn decrypt_aes(&mut self, data: Vec<u8>) -> Vec<u8> {
         let mut out_data = vec![0; data.len()];
         for i in 0..data.len() {
             out_data[i] = self.decrypt_block(data[i]);
@@ -60,28 +60,23 @@ impl McCipher {
         out_data
     }
 
-    fn shift_left(mut arr: [u8; 16], new: u8) {
+    fn encrypt_block(&mut self, data: u8) -> u8 {
-        for i in 1..arr.len() {
-            arr[i] = arr[i - 1];
-        }
-        arr[0] = new;
-    }
-
-    fn encrypt_block(&self, data: u8) -> u8 {
         let cipher = Aes128::new(GenericArray::from_slice(&self.key));
-        let mut block = GenericArray::clone_from_slice(&self.state_en);
+        let mut block = GenericArray::clone_from_slice(
+            &self.state_en.to_be_bytes());
         cipher.encrypt_block(&mut block);
-        let data = data ^ block[15];
+        let data = data ^ block[0];
-        Self::shift_left(self.state_en, data);
+        self.state_en = (self.state_en << 8) + (data as u128);
         data
     }
 
-    fn decrypt_block(&self, data: u8) -> u8 {
+    fn decrypt_block(&mut self, data: u8) -> u8 {
         let cipher = Aes128::new(GenericArray::from_slice(&self.key));
-        let mut block = GenericArray::clone_from_slice(&self.state_en);
+        let mut block = GenericArray::clone_from_slice(
+            &self.state_de.to_be_bytes());
         cipher.decrypt_block(&mut block);
-        Self::shift_left(self.state_de, data);
+        self.state_de = (self.state_de << 8) + (data as u128);
-        let data = data ^ block[15];
+        let data = data ^ block[0];
         data
     }
 }

src/mc_types.rs

@@ -10,6 +10,9 @@ use async_trait::async_trait;
 use rsa::{RsaPrivateKey, RsaPublicKey};
 use rsa::pkcs8::{EncodePublicKey, DecodePublicKey};
 use rand::Rng;
+use crypto::digest::Digest;
+use sha1::Sha1;
+use num_bigint::BigInt;
 
 use crate::login;
 use crate::encrypt::{self, McCipher};
@@ -73,6 +76,7 @@ pub struct ProtocolConnection<'a> {
     rsa_public_key: Option<RsaPublicKey>,
     aes_cipher: Option<McCipher>,
     verify_token: Option<[u8; 16]>,
+    server_id: String
 }
 
 impl<'a> ProtocolConnection<'a> {
@@ -87,6 +91,7 @@ impl<'a> ProtocolConnection<'a> {
             rsa_public_key: None,
             aes_cipher: None,
             verify_token: None,
+            server_id: "".to_string(),
         }
     }
 
@@ -109,7 +114,7 @@ impl<'a> ProtocolConnection<'a> {
                 match &self.verify_token {
                     Some (token) =>
                         Ok(login::clientbound::EncryptionRequest {
-                            server_id: "".to_string(),
+                            server_id: self.server_id.clone(),
                             public_key: key
                                 .to_public_key_der()?
                                 .as_ref()
@@ -127,6 +132,7 @@ impl<'a> ProtocolConnection<'a> {
         &mut self,
         request: login::clientbound::EncryptionRequest,
     ) -> Result<login::serverbound::EncryptionResponse> {
+        self.server_id = request.server_id;
         self.rsa_public_key = Some(
             RsaPublicKey::from_public_key_der(&request.public_key)?);
         self.aes_cipher = Some(McCipher::create());
@@ -205,6 +211,23 @@ impl<'a> ProtocolConnection<'a> {
             aes_cipher: self.aes_cipher.clone(),
         }))
     }
+
+    pub async fn server_id_hash(&self) -> Result<String> {
+        let hash_data = match &self.aes_cipher {
+            Some(aes_cipher) => match &self.rsa_public_key {
+                Some(key) => [
+                    self.server_id.as_bytes(),
+                    &aes_cipher.key,
+                    key.to_public_key_der()?.as_ref(),
+                ].concat(),
+                None => return Err(Box::new(PacketError::EncryptionError))
+            },
+            None => return Err(Box::new(PacketError::EncryptionError))
+        };
+        let hash = BigInt::from_signed_bytes_be(
+            &Sha1::digest(hash_data)).to_str_radix(16);
+        Ok(hash)
+    }
 }
 
 unsafe impl<'a> Send for ProtocolConnection<'a> {}
@@ -212,7 +235,7 @@ 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 {
+        match &mut self.aes_cipher {
             Some(aes_cipher) => {
                 let length = read_var_int_stream_encrypted(
                     self.stream_read, aes_cipher).await? as usize;
@@ -239,7 +262,7 @@ 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 {
+        match &mut self.aes_cipher {
             Some(aes_cipher) => {
                 self.stream_write.write_all(
                     &aes_cipher.encrypt_aes(out_data)).await?;
@@ -278,7 +301,7 @@ impl<'a> ProtocolWrite for WriteHaftProtocolConnection<'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 {
+        match &mut self.aes_cipher {
             Some(aes_cipher) => {
                 self.stream_write.write_all(
                     &aes_cipher.encrypt_aes(out_data)).await?;
@@ -327,7 +350,7 @@ 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 {
+        match &mut self.aes_cipher {
             Some(aes_cipher) => {
                 let length = read_var_int_stream_encrypted(
                     self.stream_read, aes_cipher).await? as usize;
@@ -389,17 +412,17 @@ async fn read_var_int_stream(stream: &mut OwnedReadHalf) -> Result<i32> {
 }
 async fn read_var_int_stream_encrypted(
     stream: &mut OwnedReadHalf,
-    cipher: &McCipher,
+    cipher: &mut McCipher,
 ) -> Result<i32> {
     let mut data: Vec<u8> = vec![];
 
     loop {
         let encrypted_byte = stream.read_u8().await?;
-        let mut current_byte = cipher.decrypt_aes(vec![encrypted_byte]);
+        let current_byte = cipher.decrypt_aes(vec![encrypted_byte])[0];
 
-        data.append(&mut current_byte);
+        data.append(&mut vec![current_byte]);
 
-        if (current_byte[0] & CONTINUE_BIT) == 0 {
+        if (current_byte & CONTINUE_BIT) == 0 {
             break;
         }
     }