diff --git a/.gitignore b/.gitignore index ea8c4bf..d03e6b9 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ /target +Notes.md \ No newline at end of file diff --git a/src/connection.rs b/src/connection.rs index 3bca2a2..aef11cb 100644 --- a/src/connection.rs +++ b/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()?; - let bytes = self.crypt.encrypt(msg)?; + + async fn send_bytes(&mut self, bytes: Bytes) -> Result<()> { 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 { match self.ms.next().await { Some(Ok(msg)) => { let decrypted_bytes = self.crypt.decrypt(msg.freeze())?; Message::deserialize(decrypted_bytes) } - _ => { - return Err(anyhow!("No response to negotiation message")); - } + Some(Err(e)) => Err(anyhow!(e.to_string())), + None => Err(anyhow!("Error awaiting msg")), } } } diff --git a/src/crypto.rs b/src/crypto.rs index fc78e80..ae66d03 100644 --- a/src/crypto.rs +++ b/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 { - let mut arr = [0u8; NONCE_SIZE_IN_BYTES]; - thread_rng().try_fill(&mut arr[..])?; - let nonce = Nonce::from_slice(&arr); + // Returns wire format, includes nonce as prefix + pub fn encrypt(&mut self, plaintext: Bytes) -> Result { + thread_rng().try_fill(&mut self.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 { - let mut body = body; - let nonce_bytes = body.split_to(NONCE_SIZE_IN_BYTES); + // Accepts wire format, includes nonce as prefix + pub fn decrypt(&self, ciphertext: Bytes) -> Result { + 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())), }