From 3c3cb50c862e482dc29d0debfe7559346bd0abea Mon Sep 17 00:00:00 2001
From: 17ms <79069176+17ms@users.noreply.github.com>
Date: Wed, 29 Mar 2023 20:20:47 +0300
Subject: [PATCH] base64 encode payloads (allows the use of msg delimiters)

---
 src/comms.rs  | 43 ++++++++++++++++++++++++++++++++-----------
 src/crypto.rs | 35 +++++++++++++++++++++++++++++------
 2 files changed, 61 insertions(+), 17 deletions(-)
 mode change 100644 => 100755 src/comms.rs
 mode change 100644 => 100755 src/crypto.rs

diff --git a/src/comms.rs b/src/comms.rs
old mode 100644
new mode 100755
index 4aabc94..a8169c3
--- a/src/comms.rs
+++ b/src/comms.rs
@@ -1,5 +1,6 @@
 use crate::crypto;
 use aes_gcm::{aead::consts::U12, aes::Aes256, AesGcm};
+use base64::{engine::general_purpose, Engine};
 use rand::rngs::OsRng;
 use std::error::Error;
 use tokio::{
@@ -7,29 +8,49 @@ use tokio::{
     net::tcp::{ReadHalf, WriteHalf},
 };
 
-pub async fn send_bytes(
+pub async fn send(
     writer: &mut BufWriter<WriteHalf<'_>>,
-    enc: Option<(&mut AesGcm<Aes256, U12>, &mut OsRng)>,
+    cipher: Option<&mut AesGcm<Aes256, U12>>,
+    rng: Option<&mut OsRng>,
     data: &Vec<u8>,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let data = enc.map_or(Ok(data.clone()), |enc| {
-        crypto::aes_encrypt(data, enc.0, enc.1)
-    })?;
-    writer.write_all(&data).await?;
+    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_bytes(
+pub async fn recv(
     reader: &mut BufReader<ReadHalf<'_>>,
     cipher: Option<&mut AesGcm<Aes256, U12>>,
     buf: &mut Vec<u8>,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let recv_bytes = reader.read_until(b'\n', buf).await?;
-    *buf = cipher.map_or(Ok(buf.clone()), |c| crypto::aes_decrypt(&buf, c))?;
-    if recv_bytes == 0 {
-        todo!("ERROR: No message received or client <xyz> crashed");
+    let n = reader.read_until(b':', buf).await?;
+
+    if n == 0 {
+        todo!("error: connection closed unexpectedly");
+    }
+
+    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(())
diff --git a/src/crypto.rs b/src/crypto.rs
old mode 100644
new mode 100755
index ebfac8c..6b9e490
--- a/src/crypto.rs
+++ b/src/crypto.rs
@@ -13,6 +13,7 @@ use tokio::{
 use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
 
 const AES_NONCE_SIZE: usize = 12;
+const DH_PBK_SIZE: usize = 32;
 
 pub async fn edh(
     reader: &mut BufReader<ReadHalf<'_>>,
@@ -22,18 +23,19 @@ pub async fn edh(
 ) -> Result<SharedSecret, Box<dyn Error + Send + Sync>> {
     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_bytes(writer, None, &own_pbk.as_bytes().to_vec()).await?;
-        comms::recv_bytes(reader, None, buf).await?;
+        comms::send(writer, None, None, &msg).await?;
+        comms::recv(reader, None, buf).await?;
     } else {
-        comms::recv_bytes(reader, None, buf).await?;
-        comms::send_bytes(writer, None, &own_pbk.as_bytes().to_vec()).await?;
+        comms::recv(reader, None, buf).await?;
+        comms::send(writer, None, None, &msg).await?;
     }
 
-    let sliced_buf: [u8; 32] = buf[..32].try_into()?;
-    let recv_pbk = PublicKey::from(sliced_buf);
+    let slice: [u8; DH_PBK_SIZE] = buf[..DH_PBK_SIZE].try_into()?;
     buf.clear();
+    let recv_pbk = PublicKey::from(slice);
 
     Ok(own_sec.diffie_hellman(&recv_pbk))
 }
@@ -75,3 +77,24 @@ pub fn aes_decrypt(
 
     Ok(decrypted)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_aes() {
+        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);
+    }
+}