redesigned comms & crypto modules

This commit is contained in:
17ms 2023-05-25 02:03:49 +03:00
parent 6b3f466b3c
commit a592528fee
4 changed files with 144 additions and 150 deletions

View File

@ -1,59 +0,0 @@
use std::error::Error;
use aes_gcm::{aead::consts::U12, aes::Aes256, AesGcm};
use base64::{engine::general_purpose, Engine};
use rand::rngs::OsRng;
use tokio::{
io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter},
net::tcp::{ReadHalf, WriteHalf},
};
use crate::crypto;
pub async fn send(
writer: &mut BufWriter<WriteHalf<'_>>,
cipher: Option<&mut AesGcm<Aes256, U12>>,
rng: Option<&mut OsRng>,
data: &Vec<u8>,
) -> Result<(), Box<dyn Error + Send + Sync>> {
let enc: Vec<u8>;
if let (Some(cipher), Some(rng)) = (cipher, rng) {
enc = crypto::aes_encrypt(data, cipher, rng)?;
} else {
enc = data.clone();
}
let mut encoded = general_purpose::STANDARD_NO_PAD
.encode(enc)
.as_bytes()
.to_vec();
encoded.push(b':');
writer.write_all(&encoded).await?;
writer.flush().await?;
Ok(())
}
pub async fn recv(
reader: &mut BufReader<ReadHalf<'_>>,
cipher: Option<&mut AesGcm<Aes256, U12>>,
) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
let mut buf = Vec::new();
let n = reader.read_until(b':', &mut buf).await?;
if n == 0 {
return Err("Received 0 bytes from the socket".into());
}
buf.pop();
buf = general_purpose::STANDARD_NO_PAD.decode(&buf)?.to_vec();
if let Some(cipher) = cipher {
buf = crypto::aes_decrypt(&buf, cipher)?;
} else {
buf = buf.clone();
}
Ok(buf)
}

105
src/crypto.rs Executable file → Normal file
View File

@ -5,21 +5,34 @@ use aes_gcm::{
aes::Aes256,
Aes256Gcm, AesGcm, KeyInit, Nonce,
};
use rand::{distributions::Alphanumeric, rngs::OsRng, Rng, RngCore};
use tokio::{
io::{BufReader, BufWriter},
net::tcp::{ReadHalf, WriteHalf},
};
use rand::{rngs::OsRng, RngCore};
use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
use crate::comms;
use crate::sockets::SocketHandler;
const AES_NONCE_SIZE: usize = 12;
const DH_PBK_SIZE: usize = 32;
async fn edh(
reader: &mut BufReader<ReadHalf<'_>>,
writer: &mut BufWriter<WriteHalf<'_>>,
#[derive(Clone)]
pub struct Crypto {
cipher: AesGcm<Aes256, U12>,
rng: OsRng,
}
impl Crypto {
pub async fn new(
handler: &mut SocketHandler<'_>,
go_first: bool,
) -> Result<Self, Box<dyn Error + Sync + Send>> {
let secret = Self::ecdh(handler, go_first).await?;
let cipher = Aes256Gcm::new(secret.as_bytes().into());
let rng = OsRng;
Ok(Self { cipher, rng })
}
async fn ecdh(
handler: &mut SocketHandler<'_>,
go_first: bool,
) -> Result<SharedSecret, Box<dyn Error + Send + Sync>> {
let buf: Vec<u8>;
@ -28,11 +41,11 @@ async fn edh(
let msg = own_pbk.as_bytes().to_vec();
if go_first {
comms::send(writer, None, None, &msg).await?;
buf = comms::recv(reader, None).await?;
handler.send(&msg).await?;
buf = handler.recv().await?;
} else {
buf = comms::recv(reader, None).await?;
comms::send(writer, None, None, &msg).await?;
buf = handler.recv().await?;
handler.send(&msg).await?;
}
let slice: [u8; DH_PBK_SIZE] = buf[..DH_PBK_SIZE].try_into()?;
@ -41,50 +54,37 @@ async fn edh(
Ok(own_sec.diffie_hellman(&recv_pbk))
}
pub async fn aes_cipher(
reader: &mut BufReader<ReadHalf<'_>>,
writer: &mut BufWriter<WriteHalf<'_>>,
go_first: bool,
) -> Result<AesGcm<Aes256, U12>, Box<dyn Error + Sync + Send>> {
let secret = edh(reader, writer, go_first).await?;
Ok(Aes256Gcm::new(secret.as_bytes().into()))
}
fn generate_nonce(rng: &mut impl RngCore) -> Nonce<U12> {
fn nonce(&self) -> Nonce<U12> {
let mut nonce = Nonce::default();
rng.fill_bytes(&mut nonce);
self.rng.fill_bytes(&mut nonce);
nonce
}
pub fn aes_encrypt(
data: &Vec<u8>,
cipher: &mut AesGcm<Aes256, U12>,
rng: &mut OsRng,
) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
let nonce = generate_nonce(rng);
let encrypted = match cipher.encrypt(&nonce, data.as_ref()) {
pub async fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
let nonce = self.nonce();
let encrypted = match self.cipher.encrypt(&nonce, data.as_ref()) {
Ok(data) => data,
Err(_) => return Err("AES encryption failed".into()),
Err(e) => return Err(format!("Encryption failed: {}", e).into()),
};
let mut data = nonce.to_vec();
data.extend_from_slice(&encrypted);
Ok(data)
}
pub fn aes_decrypt(
data: &[u8],
cipher: &mut AesGcm<Aes256, U12>,
) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
pub async fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
let (nonce_bytes, data) = data.split_at(AES_NONCE_SIZE);
let decrypted = match cipher.decrypt(Nonce::from_slice(nonce_bytes), data.as_ref()) {
let nonce = Nonce::from_slice(nonce_bytes);
let decrypted = match self.cipher.decrypt(nonce, data.as_ref()) {
Ok(data) => data,
Err(_) => return Err("AES decryption failed".into()),
Err(e) => return Err(format!("Decryption failed: {}", e).into()),
};
Ok(decrypted)
}
}
pub fn try_hash(path: &Path) -> Result<String, Box<dyn Error + Send + Sync>> {
let hash = sha256::try_digest(path)?;
@ -92,31 +92,4 @@ pub fn try_hash(path: &Path) -> Result<String, Box<dyn Error + Send + Sync>> {
Ok(hash)
}
pub fn keygen() -> String {
rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(8)
.map(char::from)
.collect::<String>()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn aes_implementations() {
use aes_gcm::aead;
let mut gen_rng = aead::OsRng;
let key = Aes256Gcm::generate_key(&mut gen_rng);
let mut cipher = Aes256Gcm::new(&key);
let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
let mut aes_rng = OsRng;
let enc = aes_encrypt(&data, &mut cipher, &mut aes_rng).unwrap();
let dec = aes_decrypt(&enc, &mut cipher).unwrap();
assert_eq!(data, dec);
}
}
// TODO: unit test if deemed necessary

View File

@ -1,8 +1,8 @@
pub mod cli;
pub mod common;
pub mod comms;
pub mod connector;
pub mod crypto;
pub mod listener;
pub mod parsers;
pub mod sockets;
pub mod util;

80
src/sockets.rs Normal file
View File

@ -0,0 +1,80 @@
use std::error::Error;
use base64::{engine::general_purpose, Engine};
use tokio::{
io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter},
net::{
tcp::{ReadHalf, WriteHalf},
TcpStream,
},
};
use crate::crypto::Crypto;
pub struct SocketHandler<'a> {
writer: BufWriter<WriteHalf<'a>>,
reader: BufReader<ReadHalf<'a>>,
crypto: Option<Crypto>,
}
impl<'a> SocketHandler<'a> {
pub fn new(socket: &'a mut TcpStream) -> Self {
let (reader, writer) = socket.split();
let mut reader = BufReader::new(reader);
let mut writer = BufWriter::new(writer);
Self {
writer,
reader,
crypto: None,
}
}
pub fn set_crypto(&self, crypto: Crypto) {
// setting up AES cipher requires DH key exchange in plaintext,
// meaning crypto can't be initialized at the same time as the socket handler
self.crypto = Some(crypto);
}
pub async fn sender(&mut self, data: &[u8]) -> Result<(), Box<dyn Error + Send + Sync>> {
let data = match &self.crypto {
Some(c) => c.encrypt(data).await?,
None => data.to_vec(),
};
self.send(&data).await?;
Ok(())
}
pub async fn send(&mut self, data: &[u8]) -> Result<(), Box<dyn Error + Send + Sync>> {
self.writer.write_all(data).await?;
self.writer.flush().await?;
Ok(())
}
pub async fn receiver(&mut self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
let mut buf = self.recv().await?;
buf.pop();
buf = general_purpose::STANDARD_NO_PAD.decode(&buf)?.to_vec();
let data = match &self.crypto {
Some(c) => c.decrypt(&buf).await?,
None => buf,
};
Ok(data)
}
pub async fn recv(&mut self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
let mut buf = Vec::new();
let n = self.reader.read_until(b':', &mut buf).await?;
if n == 0 {
return Err("Received 0 bytes from the socket".into());
}
Ok(buf)
}
}