redesigned comms & crypto modules

src/comms.rs

@@ -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)
-}

src/crypto.rs

@@ -5,85 +5,85 @@ 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<SharedSecret, Box<dyn Error + Send + Sync>> {
+    ) -> 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>;
         let own_sec = EphemeralSecret::new(OsRng);
         let own_pbk = PublicKey::from(&own_sec);
         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()?;
         let recv_pbk = PublicKey::from(slice);
 
         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>> {
@@ -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

src/lib.rs

@@ -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;

src/sockets.rs

@@ -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)
+    }
+}