From e5b038dc91254a62595e86b0c2173d4a2562e62e Mon Sep 17 00:00:00 2001
From: 17ms <79069176+17ms@users.noreply.github.com>
Date: Thu, 23 Mar 2023 03:22:30 +0200
Subject: [PATCH] include plaintext nonce to sent packets

---
 src/comms.rs  |  6 +++---
 src/crypto.rs | 14 ++++++++++----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/comms.rs b/src/comms.rs
index a5d6844..4aabc94 100644
--- a/src/comms.rs
+++ b/src/comms.rs
@@ -1,4 +1,4 @@
-use super::crypto;
+use crate::crypto;
 use aes_gcm::{aead::consts::U12, aes::Aes256, AesGcm};
 use rand::rngs::OsRng;
 use std::error::Error;
@@ -12,10 +12,10 @@ pub async fn send_bytes(
     enc: Option<(&mut AesGcm<Aes256, U12>, &mut OsRng)>,
     data: &Vec<u8>,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let processed = enc.map_or(Ok(data.clone()), |enc| {
+    let data = enc.map_or(Ok(data.clone()), |enc| {
         crypto::aes_encrypt(data, enc.0, enc.1)
     })?;
-    writer.write_all(&processed).await?;
+    writer.write_all(&data).await?;
     writer.flush().await?;
 
     Ok(())
diff --git a/src/crypto.rs b/src/crypto.rs
index 5f16cdf..ebfac8c 100644
--- a/src/crypto.rs
+++ b/src/crypto.rs
@@ -1,4 +1,4 @@
-use super::comms;
+use crate::comms;
 use aes_gcm::{
     aead::{consts::U12, AeadMut},
     aes::Aes256,
@@ -34,6 +34,7 @@ pub async fn edh(
     let sliced_buf: [u8; 32] = buf[..32].try_into()?;
     let recv_pbk = PublicKey::from(sliced_buf);
     buf.clear();
+
     Ok(own_sec.diffie_hellman(&recv_pbk))
 }
 
@@ -46,6 +47,7 @@ pub fn aes_cipher(
 fn generate_nonce(rng: &mut impl RngCore) -> Nonce<U12> {
     let mut nonce = Nonce::default();
     rng.fill_bytes(&mut nonce);
+
     nonce
 }
 
@@ -55,8 +57,11 @@ pub fn aes_encrypt(
     rng: &mut OsRng,
 ) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
     let nonce = generate_nonce(rng);
-    let encrypted = cipher.encrypt(&nonce, data.as_ref()).unwrap(); // TODO: handle error types
-    Ok(encrypted)
+    let encrypted = cipher.encrypt(&nonce, data.as_ref()).unwrap(); // TODO: handle errors
+    let mut data = nonce.to_vec();
+    data.extend_from_slice(&encrypted);
+
+    Ok(data)
 }
 
 pub fn aes_decrypt(
@@ -66,6 +71,7 @@ pub fn aes_decrypt(
     let (nonce_bytes, data) = data.split_at(AES_NONCE_SIZE);
     let decrypted = cipher
         .decrypt(Nonce::from_slice(nonce_bytes), data.as_ref())
-        .unwrap(); // TODO: handle error types
+        .unwrap(); // TODO: handle errors
+
     Ok(decrypted)
 }