Tweaks

.gitignore

@@ -1 +1,2 @@
 /target
+Notes.md

src/connection.rs

@@ -1,6 +1,7 @@
 use crate::crypto::Crypt;
 use crate::message::{Message, MessageStream};
 use anyhow::{anyhow, Result};
+use bytes::Bytes;
 use futures::{SinkExt, StreamExt};
 use tokio::net::TcpStream;
 
@@ -15,24 +16,28 @@ impl Connection {
         let crypt = Crypt::new(&key);
         Connection { ms, crypt }
     }
-    pub async fn send_msg(&mut self, msg: Message) -> Result<()> {
+
-        let msg = msg.serialize()?;
+    async fn send_bytes(&mut self, bytes: Bytes) -> Result<()> {
-        let bytes = self.crypt.encrypt(msg)?;
         match self.ms.send(bytes).await {
             Ok(_) => Ok(()),
             Err(e) => Err(anyhow!(e.to_string())),
         }
     }
 
+    pub async fn send_msg(&mut self, msg: Message) -> Result<()> {
+        let msg = msg.serialize()?;
+        let bytes = self.crypt.encrypt(msg)?;
+        self.send_bytes(bytes).await
+    }
+
     pub async fn await_msg(&mut self) -> Result<Message> {
         match self.ms.next().await {
             Some(Ok(msg)) => {
                 let decrypted_bytes = self.crypt.decrypt(msg.freeze())?;
                 Message::deserialize(decrypted_bytes)
             }
-            _ => {
+            Some(Err(e)) => Err(anyhow!(e.to_string())),
-                return Err(anyhow!("No response to negotiation message"));
+            None => Err(anyhow!("Error awaiting msg")),
-            }
         }
     }
 }

src/crypto.rs

@@ -8,6 +8,7 @@ use rand::{thread_rng, Rng};
 
 pub struct Crypt {
     cipher: Aes256Gcm,
+    arr: [u8; NONCE_SIZE_IN_BYTES],
 }
 
 impl Crypt {
@@ -15,13 +16,14 @@ impl Crypt {
         let key = Key::from_slice(&key[..]);
         Crypt {
             cipher: Aes256Gcm::new(key),
+            arr: [0u8; NONCE_SIZE_IN_BYTES],
         }
     }
 
-    pub fn encrypt(&self, plaintext: Bytes) -> Result<Bytes> {
+    // Returns wire format, includes nonce as prefix
-        let mut arr = [0u8; NONCE_SIZE_IN_BYTES];
+    pub fn encrypt(&mut self, plaintext: Bytes) -> Result<Bytes> {
-        thread_rng().try_fill(&mut arr[..])?;
+        thread_rng().try_fill(&mut self.arr[..])?;
-        let nonce = Nonce::from_slice(&arr);
+        let nonce = Nonce::from_slice(&self.arr);
         match self.cipher.encrypt(nonce, plaintext.as_ref()) {
             Ok(body) => {
                 let mut buffer = BytesMut::with_capacity(NONCE_SIZE_IN_BYTES + body.len());
@@ -33,11 +35,12 @@ impl Crypt {
         }
     }
 
-    pub fn decrypt(&self, body: Bytes) -> Result<Bytes> {
+    // Accepts wire format, includes nonce as prefix
-        let mut body = body;
+    pub fn decrypt(&self, ciphertext: Bytes) -> Result<Bytes> {
-        let nonce_bytes = body.split_to(NONCE_SIZE_IN_BYTES);
+        let mut ciphertext_body = ciphertext;
+        let nonce_bytes = ciphertext_body.split_to(NONCE_SIZE_IN_BYTES);
         let nonce = Nonce::from_slice(&nonce_bytes);
-        match self.cipher.decrypt(nonce, body.as_ref()) {
+        match self.cipher.decrypt(nonce, ciphertext_body.as_ref()) {
             Ok(payload) => Ok(Bytes::from(payload)),
             Err(e) => Err(anyhow!(e.to_string())),
         }