Fixed ProtocolConnection::split_conn function

Cargo.toml

@@ -12,3 +12,9 @@ tokio = { version = "1", features = ["full"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 async-trait = "0.1.75"
+rand = "0.8.5"
+# aes = "0.8.4"
+aes = "0.7"
+# rsa = "0.9.6"
+rsa = "0.6"
+pkcs8 = "0.8"

src/encrypt.rs

@@ -0,0 +1,46 @@
+// Yeahbut May 2024
+
+use rsa::PublicKey;
+use rsa::{RsaPrivateKey, RsaPublicKey, PaddingScheme, errors::Result};
+use rand::rngs::OsRng;
+use aes::Aes128;
+use aes::cipher::{
+    BlockEncrypt, BlockDecrypt, NewBlockCipher, generic_array::GenericArray};
+
+pub fn generate_rsa_keys() -> Result<RsaPrivateKey> {
+    let mut rng = OsRng;
+    let bits = 2048;
+    let private_key = RsaPrivateKey::new(&mut rng, bits)?;
+    Ok(private_key)
+}
+
+pub fn encrypt_rsa(
+    public_key: &RsaPublicKey,
+    data: &[u8; 16],
+) -> Result<Vec<u8>> {
+    let padding = PaddingScheme::new_pkcs1v15_encrypt();
+    let mut rng = OsRng;
+    public_key.encrypt(&mut rng, padding, data)
+}
+
+pub fn decrypt_rsa(
+    private_key: &RsaPrivateKey,
+    data: &[u8],
+) -> Result<Vec<u8>> {
+    let padding = PaddingScheme::new_pkcs1v15_encrypt();
+    private_key.decrypt(padding, data)
+}
+
+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(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/handshake.rs

@@ -2,8 +2,6 @@
 
 pub mod serverbound {
 
-    use tokio::net::tcp::OwnedReadHalf;
-
     use crate::mc_types::{self, Result, Packet, PacketError};
 
     pub enum HandshakeEnum {
@@ -11,8 +9,10 @@ pub mod serverbound {
     }
 
     impl HandshakeEnum {
-        pub async fn read(stream: &mut OwnedReadHalf) -> Result<Self> {
+        pub async fn read<T: mc_types::ProtocolRead>(
-            let mut data = mc_types::read_data(stream).await?;
+            conn: &mut T,
+        ) -> Result<Self> {
+            let mut data = conn.read_data().await?;
             let packet_id = mc_types::get_var_int(&mut data)?;
             if packet_id == Handshake::packet_id() {
                 return Ok(Self::Handshake(Handshake::get(&mut data)?))

src/lib.rs

@@ -4,3 +4,5 @@ pub mod mc_types;
 pub mod handshake;
 pub mod status;
 pub mod login;
+pub mod encrypt;
+pub mod play;

src/login.rs

@@ -2,8 +2,6 @@
 
 pub mod clientbound {
 
-    use tokio::net::tcp::OwnedReadHalf;
-
     use crate::mc_types::{self, Result, Packet, PacketArray, PacketError};
 
     pub enum Login {
@@ -15,8 +13,10 @@ pub mod clientbound {
     }
 
     impl Login {
-        pub async fn read(stream: &mut OwnedReadHalf) -> Result<Self> {
+        pub async fn read<T: mc_types::ProtocolRead>(
-            let mut data = mc_types::read_data(stream).await?;
+            conn: &mut T,
+        ) -> Result<Self> {
+            let mut data = conn.read_data().await?;
             let packet_id = mc_types::get_var_int(&mut data)?;
             if packet_id == Disconnect::packet_id() {
                 return Ok(Self::Disconnect(Disconnect::get(&mut data)?))
@@ -232,8 +232,6 @@ pub mod clientbound {
 
 pub mod serverbound {
 
-    use tokio::net::tcp::OwnedReadHalf;
-
     use crate::mc_types::{self, Result, Packet, PacketError};
 
     pub enum Login {
@@ -243,8 +241,10 @@ pub mod serverbound {
     }
 
     impl Login {
-        pub async fn read(stream: &mut OwnedReadHalf) -> Result<Self> {
+        pub async fn read<T: mc_types::ProtocolRead>(
-            let mut data = mc_types::read_data(stream).await?;
+            conn: &mut T,
+        ) -> Result<Self> {
+            let mut data = conn.read_data().await?;
             let packet_id = mc_types::get_var_int(&mut data)?;
             if packet_id == LoginStart::packet_id() {
                 return Ok(Self::LoginStart(LoginStart::get(&mut data)?))

src/mc_types.rs

@@ -4,9 +4,17 @@ 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 crate::login;
+use crate::encrypt;
+use crate::play::Play;
 
 pub type Result<T> = std::result::Result<T, Box<dyn Error>>;
 
@@ -21,6 +29,7 @@ pub enum PacketError {
     ValueTooLarge,
     RanOutOfBytes,
     InvalidPacketId,
+    EncryptionError,
 }
 
 impl fmt::Display for PacketError {
@@ -32,6 +41,8 @@ impl fmt::Display for PacketError {
                 write!(f, "Ran out of bytes while reading VarInt"),
             PacketError::InvalidPacketId =>
                 write!(f, "Invalid packet id"),
+            PacketError::EncryptionError =>
+                write!(f, "Encryption Error"),
         }
     }
 }
@@ -43,14 +54,360 @@ pub struct Chat {
     pub text: String,
 }
 
+#[async_trait]
+pub trait ProtocolRead {
+    async fn read_data(&mut self) -> Result<Vec<u8>>;
+}
+
+#[async_trait]
+pub trait ProtocolWrite {
+    async fn write_data(&mut self, data: &mut Vec<u8>) -> Result<()>;
+}
+
+pub struct ProtocolConnection<'a> {
+    pub stream_read: &'a mut OwnedReadHalf,
+    pub stream_write: &'a mut OwnedWriteHalf,
+    rsa_private_key: Option<RsaPrivateKey>,
+    rsa_public_key: Option<RsaPublicKey>,
+    aes_encryption_key: Option<[u8; 16]>,
+    verify_token: Option<[u8; 16]>,
+}
+
+impl<'a> ProtocolConnection<'a> {
+    pub fn new(
+        stream_read: &'a mut OwnedReadHalf,
+        stream_write: &'a mut OwnedWriteHalf,
+    ) -> Self {
+        ProtocolConnection {
+            stream_read,
+            stream_write,
+            rsa_private_key: None,
+            rsa_public_key: None,
+            aes_encryption_key: None,
+            verify_token: None,
+        }
+    }
+
+    pub fn create_encryption_request(
+        &mut self,
+        private_key: RsaPrivateKey,
+    ) -> Result<login::clientbound::EncryptionRequest> {
+        match self.rsa_private_key {
+            Some(_) => {},
+            None => {
+                let public_key = RsaPublicKey::from(&private_key);
+                let mut rng = rand::thread_rng();
+                self.rsa_private_key = Some(private_key);
+                self.rsa_public_key = Some(public_key);
+                self.verify_token = Some(rng.gen());
+            }
+        };
+        match &self.rsa_public_key {
+            Some (key) => {
+                match &self.verify_token {
+                    Some (token) =>
+                        Ok(login::clientbound::EncryptionRequest {
+                            server_id: "".to_string(),
+                            public_key: key
+                                .to_public_key_der()?
+                                .as_ref()
+                                .to_vec(),
+                            verify_token: token[0..16].to_vec(),
+                        }),
+                    None => Err(Box::new(PacketError::EncryptionError))
+                }
+            },
+            None => Err(Box::new(PacketError::EncryptionError))
+        }
+    }
+
+    pub fn handle_encryption_request(
+        &mut self,
+        request: login::clientbound::EncryptionRequest,
+    ) -> Result<login::serverbound::EncryptionResponse> {
+        self.rsa_public_key = Some(
+            RsaPublicKey::from_public_key_der(&request.public_key)?);
+        let mut rng = rand::thread_rng();
+        self.aes_encryption_key = Some(rng.gen());
+        match self.aes_encryption_key {
+            Some(key) => {
+                match &self.rsa_public_key {
+                    Some(public_key) => {
+                        Ok(login::serverbound::EncryptionResponse {
+                            shared_secret: encrypt::encrypt_rsa(
+                                public_key, &key)?,
+                            verify_token: encrypt::encrypt_rsa(
+                                public_key,
+                                request.verify_token[0..16]
+                                    .try_into()
+                                    .unwrap(),
+                            )?,
+                        })
+                    },
+                    None => Err(Box::new(PacketError::EncryptionError))
+                }
+            },
+            None => Err(Box::new(PacketError::EncryptionError))
+        }
+    }
+
+    pub fn handle_encryption_response(
+        &mut self,
+        response: login::serverbound::EncryptionResponse,
+    ) -> Result<()> {
+        match &self.verify_token {
+            Some (token) => {
+                match &self.rsa_private_key {
+                    Some (private_key) => {
+                        if &encrypt::decrypt_rsa(
+                            &private_key,
+                            response.verify_token.as_slice()
+                        )? == token {
+                            self.aes_encryption_key =
+                                Some(encrypt::decrypt_rsa(
+                                    &private_key,
+                                    response.shared_secret.as_slice()
+                                )?[0..16].try_into().unwrap());
+                            Ok(())
+                        } else {
+                            Err(Box::new(PacketError::EncryptionError))
+                        }
+                    }
+                    None => Err(Box::new(PacketError::EncryptionError))
+                }
+            }
+            None => Err(Box::new(PacketError::EncryptionError))
+        }
+    }
+
+    pub async fn forward_play(
+        &mut self,
+        other: &mut ProtocolConnection<'_>,
+    ) -> Result<()> {
+        loop {
+            let packet = Play::read(self).await?;
+            match packet {
+                Play::PlayPacket(packet) => packet.write(other).await?,
+            };
+        }
+    }
+
+    pub fn split_conn(
+        &mut self
+    ) -> Result<(WriteHaftProtocolConnection, ReadHaftProtocolConnection)> {
+        Ok((WriteHaftProtocolConnection {
+            stream_write: &mut self.stream_write,
+            aes_encryption_key: self.aes_encryption_key.clone(),
+        },
+        ReadHaftProtocolConnection {
+            stream_read: &mut self.stream_read,
+            aes_encryption_key: self.aes_encryption_key.clone(),
+        }))
+    }
+}
+
+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_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_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 {
+                        (raw_length - buffer.len() as i32) / 16
+                    } else {
+                        ((raw_length - buffer.len() as i32) / 16) + 1
+                    };
+
+                for _ in 0..length {
+                    let mut block: Vec<u8> = vec![0; 16];
+                    self.stream_read.read_exact(&mut block).await?;
+                    buffer.append(&mut block);
+                }
+
+                Ok(buffer)
+            },
+            None => {
+                let length = read_var_int_stream(
+                    self.stream_read).await? as usize;
+
+                let mut buffer: Vec<u8> = vec![0; length];
+                self.stream_read.read_exact(&mut buffer).await?;
+
+                Ok(buffer)
+            }
+        }
+    }
+}
+
+#[async_trait]
+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_encryption_key {
+            Some(aes_key) => {
+                let length =
+                    if (data.len() as i32) % 16 == 0 {
+                        (data.len() as i32) / 16
+                    } else {
+                        ((data.len() as i32) / 16) + 1
+                    };
+
+                for _ in 0..length {
+                    let mut block: Vec<u8> = out_data[0..16].to_vec();
+                    block = encrypt::encrypt_aes(
+                        &aes_key, block[0..16].try_into().unwrap());
+                    self.stream_write.write_all(&block).await?;
+                }
+
+                Ok(())
+            },
+            None => {
+                self.stream_write.write_all(&out_data).await?;
+
+                Ok(())
+            }
+        }
+    }
+}
+
+pub struct WriteHaftProtocolConnection<'a> {
+    pub stream_write: &'a mut OwnedWriteHalf,
+    aes_encryption_key: Option<[u8; 16]>,
+}
+
+impl<'a> WriteHaftProtocolConnection<'a> {
+    pub fn new(
+        stream_write: &'a mut OwnedWriteHalf,
+    ) -> Self {
+        WriteHaftProtocolConnection {
+            stream_write,
+            aes_encryption_key: None,
+        }
+    }
+}
+
+unsafe impl<'a> Send for WriteHaftProtocolConnection<'a> {}
+
+#[async_trait]
+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_encryption_key {
+            Some(aes_key) => {
+                let length =
+                    if (data.len() as i32) % 16 == 0 {
+                        (data.len() as i32) / 16
+                    } else {
+                        ((data.len() as i32) / 16) + 1
+                    };
+
+                for _ in 0..length {
+                    let mut block: Vec<u8> = out_data[0..16].to_vec();
+                    block = encrypt::encrypt_aes(
+                        &aes_key, block[0..16].try_into().unwrap());
+                    self.stream_write.write_all(&block).await?;
+                }
+
+
+                Ok(())
+            },
+            None => {
+                self.stream_write.write_all(&out_data).await?;
+
+                Ok(())
+            }
+        }
+    }
+}
+
+pub struct ReadHaftProtocolConnection<'a> {
+    pub stream_read: &'a mut OwnedReadHalf,
+    aes_encryption_key: Option<[u8; 16]>,
+}
+
+impl<'a> ReadHaftProtocolConnection<'a> {
+    pub fn new(
+        stream_read: &'a mut OwnedReadHalf,
+    ) -> Self {
+        ReadHaftProtocolConnection {
+            stream_read,
+            aes_encryption_key: None,
+        }
+    }
+
+    pub async fn forward_play<T: ProtocolWrite + Send>(
+        &mut self,
+        other: &mut T,
+    ) -> Result<()> {
+        loop {
+            let packet = Play::read(self).await?;
+            match packet {
+                Play::PlayPacket(packet) => packet.write(other).await?,
+            };
+        }
+    }
+}
+
+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_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_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 {
+                        (raw_length - buffer.len() as i32) / 16
+                    } else {
+                        ((raw_length - buffer.len() as i32) / 16) + 1
+                    };
+
+                for _ in 0..length {
+                    let mut block: Vec<u8> = vec![0; 16];
+                    self.stream_read.read_exact(&mut block).await?;
+                    buffer.append(&mut block);
+                }
+
+                Ok(buffer)
+            },
+            None => {
+                let length = read_var_int_stream(
+                    self.stream_read).await? as usize;
+
+                let mut buffer: Vec<u8> = vec![0; length];
+                self.stream_read.read_exact(&mut buffer).await?;
+
+                Ok(buffer)
+            }
+        }
+    }
+}
+
 #[async_trait]
 pub trait Packet: Sized {
     fn packet_id() -> i32;
     fn get(data: &mut Vec<u8>) -> Result<Self>;
     fn convert(&self) -> Vec<u8>;
 
-    async fn read(stream: &mut OwnedReadHalf) -> Result<Self> {
+    async fn read<T: ProtocolRead + Send>(conn: &mut T) -> Result<Self> {
-        let mut data = read_data(stream).await?;
+        let mut data = conn.read_data().await?;
         let packet_id = get_var_int(&mut data)?;
         if packet_id == Self::packet_id() {
             return Ok(Self::get(&mut data)?)
@@ -59,30 +416,11 @@ pub trait Packet: Sized {
         }
     }
 
-    async fn write(&self, stream: &mut OwnedWriteHalf) -> Result<()> {
+    async fn write<T: ProtocolWrite + Send>(&self, conn: &mut T) -> Result<()> {
-        write_data(stream, &mut self.convert()).await
+        conn.write_data(&mut self.convert()).await
     }
 }
 
-pub async fn read_data(stream: &mut OwnedReadHalf) -> Result<Vec<u8>> {
-    let length = read_var_int_stream(stream).await? as usize;
-
-    let mut buffer: Vec<u8> = vec![0; length];
-    stream.read_exact(&mut buffer).await?;
-
-    Ok(buffer)
-}
-pub async fn write_data(
-    stream: &mut OwnedWriteHalf,
-    data: &mut Vec<u8>,
-) -> Result<()> {
-    let mut out_data = convert_var_int(data.len() as i32);
-    out_data.append(data);
-
-    stream.write_all(&out_data).await?;
-
-    Ok(())
-}
 async fn read_var_int_stream(stream: &mut OwnedReadHalf) -> Result<i32> {
     let mut data: Vec<u8> = vec![];
 
@@ -100,6 +438,23 @@ async fn read_var_int_stream(stream: &mut OwnedReadHalf) -> Result<i32> {
 
     Ok(varint)
 }
+fn read_var_int_vec(stream: &mut Vec<u8>) -> Result<i32> {
+    let mut data: Vec<u8> = vec![];
+
+    loop {
+        let current_byte = stream.remove(0);
+
+        data.append(&mut vec![current_byte]);
+
+        if (current_byte & CONTINUE_BIT) == 0 {
+            break;
+        }
+    }
+
+    let varint = get_var_int(&mut data)?;
+
+    Ok(varint)
+}
 
 pub trait PacketArray: Sized {
     fn get(data: &mut Vec<u8>) -> Result<Self>;

src/play.rs

@@ -0,0 +1,39 @@
+
+use crate::mc_types::{self, Packet, Result};
+
+
+pub enum Play {
+    PlayPacket(PlayPacket),
+}
+
+impl Play {
+    pub async fn read<T: mc_types::ProtocolRead>(
+        conn: &mut T,
+    ) -> Result<Self> {
+        let mut data = conn.read_data().await?;
+        Ok(Self::PlayPacket(PlayPacket::get(&mut data)?))
+    }
+}
+
+pub struct PlayPacket {
+    pub data: Vec<u8>
+}
+
+impl Packet for PlayPacket {
+
+    fn packet_id() -> i32 {0}
+
+    fn get(data: &mut Vec<u8>) -> Result<Self> {
+        Ok(Self {
+            data: data.clone()
+        })
+    }
+
+    fn convert(&self) -> Vec<u8> {
+        let mut data: Vec<u8> = vec![];
+        data.append(&mut self.data.clone());
+
+        data
+    }
+
+}

src/status.rs

@@ -2,7 +2,6 @@
 
 pub mod clientbound {
 
-    use tokio::net::tcp::OwnedReadHalf;
     use serde::{Serialize, Deserialize};
 
     use crate::mc_types::{self, Result, Packet, PacketError};
@@ -46,8 +45,10 @@ pub mod clientbound {
     }
 
     impl StatusPackets {
-        pub async fn read(stream: &mut OwnedReadHalf) -> Result<Self> {
+        pub async fn read<T: mc_types::ProtocolRead>(
-            let mut data = mc_types::read_data(stream).await?;
+            conn: &mut T,
+        ) -> Result<Self> {
+            let mut data = conn.read_data().await?;
             let packet_id = mc_types::get_var_int(&mut data)?;
             if packet_id == Status::packet_id() {
                 return Ok(Self::Status(Status::get(&mut data)?))
@@ -125,8 +126,6 @@ pub mod clientbound {
 
 pub mod serverbound {
 
-    use tokio::net::tcp::OwnedReadHalf;
-
     use crate::mc_types::{self, Result, Packet, PacketError};
 
     pub enum StatusPackets {
@@ -135,8 +134,10 @@ pub mod serverbound {
     }
 
     impl StatusPackets {
-        pub async fn read(stream: &mut OwnedReadHalf) -> Result<Self> {
+        pub async fn read<T: mc_types::ProtocolRead>(
-            let mut data = mc_types::read_data(stream).await?;
+            conn: &mut T,
+        ) -> Result<Self> {
+            let mut data = conn.read_data().await?;
             let packet_id = mc_types::get_var_int(&mut data)?;
             if packet_id == Status::packet_id() {
                 return Ok(Self::Status(Status::get(&mut data)?))