Compare commits

..

No commits in common. "80bf8492583f32184b46d3e7a4e5616618756a63" and "4e9b60138e4577a97161409defb9624ae58fc041" have entirely different histories.

2 changed files with 26 additions and 28 deletions

View File

@ -4,7 +4,8 @@ use rsa::PublicKey;
use rsa::{RsaPrivateKey, RsaPublicKey, PaddingScheme, errors::Result}; use rsa::{RsaPrivateKey, RsaPublicKey, PaddingScheme, errors::Result};
use rand::rngs::OsRng; use rand::rngs::OsRng;
use aes::Aes128; 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> { pub fn generate_rsa_keys() -> Result<RsaPrivateKey> {
let mut rng = OsRng; let mut rng = OsRng;
@ -30,13 +31,15 @@ pub fn decrypt_rsa(
private_key.decrypt(padding, data) 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); let mut block = GenericArray::clone_from_slice(data);
cipher.encrypt_block(&mut block); cipher.encrypt_block(&mut block);
block.to_vec() 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); let mut block = GenericArray::clone_from_slice(data);
cipher.decrypt_block(&mut block); cipher.decrypt_block(&mut block);
block.to_vec() block.to_vec()

View File

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