Refactor handshake

src/client.rs

@@ -1,11 +1,11 @@
 use crate::conf::BUFFER_SIZE;
+use crate::crypto::Crypt;
 use crate::file::{to_size_string, FileHandle, FileInfo};
 use crate::handshake::Handshake;
 use crate::message::{
     EncryptedMessage, FileNegotiationPayload, FileTransferPayload, Message, MessageStream,
 };
 
-use aes_gcm::Aes256Gcm;
 use anyhow::{anyhow, Result};
 use bytes::{Bytes, BytesMut};
 use futures::future::try_join_all;
@@ -26,16 +26,18 @@ pub async fn send(file_paths: &Vec<PathBuf>, password: &String) -> Result<()> {
     // Establish connection to server
     let socket = TcpStream::connect("127.0.0.1:8080").await?;
 
-    // Complete handshake, returning cipher used for encryption
     let (handshake, s1) = Handshake::from_password(password);
-    let (socket, cipher) = handshake.negotiate(socket, s1).await?;
+    // Complete handshake, returning key used for encryption
+    let (socket, key) = handshake.negotiate(socket, s1).await?;
 
     let mut stream = Message::to_stream(socket);
+    let crypt = Crypt::new(&key);
     // Complete file negotiation
-    let handles = negotiate_files_up(handles, &mut stream, &cipher).await?;
+    let handles = negotiate_files_up(handles, &mut stream, &crypt).await?;
 
     // Upload negotiated files
-    upload_encrypted_files(&mut stream, handles, &cipher).await?;
+    upload_encrypted_files(&mut stream, handles, &crypt).await?;
+    println!("Done uploading.");
 
     // Exit
     Ok(())
@@ -44,11 +46,12 @@ pub async fn send(file_paths: &Vec<PathBuf>, password: &String) -> Result<()> {
 pub async fn receive(password: &String) -> Result<()> {
     let socket = TcpStream::connect("127.0.0.1:8080").await?;
     let (handshake, s1) = Handshake::from_password(password);
-    let (socket, cipher) = handshake.negotiate(socket, s1).await?;
+    let (socket, key) = handshake.negotiate(socket, s1).await?;
     let mut stream = Message::to_stream(socket);
-    let files = negotiate_files_down(&mut stream, &cipher).await?;
+    let crypt = Crypt::new(&key);
+    let files = negotiate_files_down(&mut stream, &crypt).await?;
 
-    download_files(files, &mut stream, &cipher).await?;
+    download_files(files, &mut stream, &crypt).await?;
     return Ok(());
 }
 
@@ -63,23 +66,22 @@ pub async fn get_file_handles(file_paths: &Vec<PathBuf>) -> Result<Vec<FileHandl
 pub async fn negotiate_files_up(
     file_handles: Vec<FileHandle>,
     stream: &mut MessageStream,
-    cipher: &Aes256Gcm,
+    crypt: &Crypt,
 ) -> Result<Vec<FileHandle>> {
     let files = file_handles.iter().map(|fh| fh.to_file_info()).collect();
     let msg = EncryptedMessage::FileNegotiationMessage(FileNegotiationPayload { files });
-    let server_msg = msg.to_encrypted_message(cipher)?;
+    let server_msg = msg.to_encrypted_message(crypt)?;
     println!("server_msg encrypted: {:?}", server_msg);
     stream.send(server_msg).await?;
     let reply_payload = match stream.next().await {
         Some(Ok(msg)) => match msg {
             Message::EncryptedMessage(response) => response,
-            _ => return Err(anyhow!("Expecting encrypted message back")),
         },
         _ => {
             return Err(anyhow!("No response to negotiation message"));
         }
     };
-    let plaintext_reply = EncryptedMessage::from_encrypted_message(cipher, &reply_payload)?;
+    let plaintext_reply = EncryptedMessage::from_encrypted_message(crypt, &reply_payload)?;
     let requested_paths: Vec<PathBuf> = match plaintext_reply {
         EncryptedMessage::FileNegotiationMessage(fnm) => {
             fnm.files.into_iter().map(|f| f.path).collect()
@@ -94,18 +96,17 @@ pub async fn negotiate_files_up(
 
 pub async fn negotiate_files_down(
     stream: &mut MessageStream,
-    cipher: &Aes256Gcm,
+    crypt: &Crypt,
 ) -> Result<Vec<FileInfo>> {
     let file_offer = match stream.next().await {
         Some(Ok(msg)) => match msg {
             Message::EncryptedMessage(response) => response,
-            _ => return Err(anyhow!("Expecting encrypted message back")),
         },
         _ => {
             return Err(anyhow!("No response to negotiation message"));
         }
     };
-    let plaintext_offer = EncryptedMessage::from_encrypted_message(cipher, &file_offer)?;
+    let plaintext_offer = EncryptedMessage::from_encrypted_message(crypt, &file_offer)?;
     let requested_infos: Vec<FileInfo> = match plaintext_offer {
         EncryptedMessage::FileNegotiationMessage(fnm) => fnm.files,
         _ => return Err(anyhow!("Expecting file negotiation message back")),
@@ -125,7 +126,7 @@ pub async fn negotiate_files_down(
     let msg = EncryptedMessage::FileNegotiationMessage(FileNegotiationPayload {
         files: files.clone(),
     });
-    let server_msg = msg.to_encrypted_message(cipher)?;
+    let server_msg = msg.to_encrypted_message(crypt)?;
     stream.send(server_msg).await?;
     Ok(files)
 }
@@ -133,7 +134,7 @@ pub async fn negotiate_files_down(
 pub async fn upload_encrypted_files(
     stream: &mut MessageStream,
     handles: Vec<FileHandle>,
-    cipher: &Aes256Gcm,
+    cipher: &Crypt,
 ) -> Result<()> {
     let (tx, mut rx) = mpsc::unbounded_channel::<EncryptedMessage>();
     for mut handle in handles {
@@ -188,7 +189,7 @@ pub async fn enqueue_file_chunks(
 pub async fn download_files(
     file_infos: Vec<FileInfo>,
     stream: &mut MessageStream,
-    cipher: &Aes256Gcm,
+    cipher: &Crypt,
 ) -> Result<()> {
     // for each file_info
     let mut info_handles: HashMap<PathBuf, mpsc::UnboundedSender<(u64, Bytes)>> = HashMap::new();
@@ -206,18 +207,15 @@ pub async fn download_files(
                     // println!("encrypted message received! {:?}", ec);
                     match ec {
                         EncryptedMessage::FileTransferMessage(payload) => {
-                            println!("matched file transfer message");
+                            // println!("matched file transfer message");
                             if let Some(tx) = info_handles.get(&payload.file_info.path) {
-                                println!("matched on filetype, sending to tx");
+                                // println!("matched on filetype, sending to tx");
                                 tx.send((payload.chunk_num, payload.chunk))?
                             };
                         },
                         _ => {println!("wrong msg")}
                     }
                 }
-                Some(Ok(_)) => {
-                    println!("wrong msg");
-                }
                 Some(Err(e)) => {
                     println!("Error {:?}", e);
                 }

src/conf.rs

@@ -1,3 +1,5 @@
-pub const ID_SIZE: usize = 32;
-pub const HANDSHAKE_MSG_SIZE: usize = 33;
-pub const BUFFER_SIZE: usize = 1024 * 1024 * 10;
+pub const ID_SIZE: usize = 32; // Blake256 of password
+pub const HANDSHAKE_MSG_SIZE: usize = 33; // generated by Spake2
+pub const PER_CLIENT_BUFFER: usize = 1024 * 1024; // buffer size allocated by server for each client
+pub const BUFFER_SIZE: usize = 1024 * 1024 * 1024; // chunk size for files sent over wire
+pub const NONCE_SIZE_IN_BYTES: usize = 96 / 8; // used for every encrypted message

src/connection.rs

@@ -0,0 +1,4 @@
+pub struct Connection {
+    ms: MessageStream;
+    crypt: Crypt
+}

src/crypto.rs

@@ -1,3 +1,4 @@
+use crate::conf::NONCE_SIZE_IN_BYTES;
 use crate::message::EncryptedPayload;
 use aes_gcm::aead::{Aead, NewAead};
 use aes_gcm::{Aes256Gcm, Key, Nonce}; // Or `Aes128Gcm`
@@ -6,59 +7,37 @@ use bytes::Bytes;
 
 use rand::{thread_rng, Rng};
 
-// pub async fn handshake(
-//     socket: TcpStream,
-//     password: Bytes,
-//     id: Bytes,
-// ) -> Result<(TcpStream, Aes256Gcm)> {
-//     let mut socket = socket;
-//     let (s1, outbound_msg) =
-//         Spake2::<Ed25519Group>::start_symmetric(&Password::new(password), &Identity::new(&id));
-//     println!("client - sending handshake msg");
-//     let mut handshake_msg = BytesMut::with_capacity(32 + 33);
-//     handshake_msg.extend_from_slice(&id);
-//     handshake_msg.extend_from_slice(&outbound_msg);
-//     let handshake_msg = handshake_msg.freeze();
-//     println!("client - handshake msg, {:?}", handshake_msg);
-//     println!("id: {:?}. msg: {:?}", id.clone(), outbound_msg.clone());
-//     socket.write_all(&handshake_msg).await?;
-//     let mut buffer = [0; 33];
-//     let n = socket.read_exact(&mut buffer).await?;
-//     println!("The bytes: {:?}", &buffer[..n]);
-//     let first_message = BytesMut::from(&buffer[..n]).freeze();
-//     println!("client - handshake msg responded to: {:?}", first_message);
-//     let key = match s1.finish(&first_message[..]) {
-//         Ok(key_bytes) => key_bytes,
-//         Err(e) => return Err(anyhow!(e.to_string())),
-//     };
-//     println!("Handshake successful. Key is {:?}", key);
-//     return Ok((socket, new_cipher(&key)));
-// }
-
-pub fn new_cipher(key: &Vec<u8>) -> Aes256Gcm {
-    let key = Key::from_slice(&key[..]);
-    Aes256Gcm::new(key)
+pub struct Crypt {
+    cipher: Aes256Gcm,
 }
 
-pub const NONCE_SIZE_IN_BYTES: usize = 96 / 8;
-pub fn encrypt(cipher: &Aes256Gcm, body: &Vec<u8>) -> Result<EncryptedPayload> {
-    let mut arr = [0u8; NONCE_SIZE_IN_BYTES];
-    thread_rng().try_fill(&mut arr[..])?;
-    let nonce = Nonce::from_slice(&arr);
-    let plaintext = body.as_ref();
-    match cipher.encrypt(nonce, plaintext) {
-        Ok(body) => Ok(EncryptedPayload {
-            nonce: arr.to_vec(),
-            body,
-        }),
-        Err(_) => Err(anyhow!("Encryption error")),
-    }
-}
-
-pub fn decrypt(cipher: &Aes256Gcm, payload: &EncryptedPayload) -> Result<Bytes> {
-    let nonce = Nonce::from_slice(payload.nonce.as_ref());
-    match cipher.decrypt(nonce, payload.body.as_ref()) {
-        Ok(payload) => Ok(Bytes::from(payload)),
-        Err(_) => Err(anyhow!("Decryption error")),
+impl Crypt {
+    pub fn new(key: &Vec<u8>) -> Crypt {
+        let key = Key::from_slice(&key[..]);
+        Crypt {
+            cipher: Aes256Gcm::new(key),
+        }
+    }
+
+    pub fn encrypt(&self, body: &Vec<u8>) -> Result<EncryptedPayload> {
+        let mut arr = [0u8; NONCE_SIZE_IN_BYTES];
+        thread_rng().try_fill(&mut arr[..])?;
+        let nonce = Nonce::from_slice(&arr);
+        let plaintext = body.as_ref();
+        match self.cipher.encrypt(nonce, plaintext) {
+            Ok(body) => Ok(EncryptedPayload {
+                nonce: arr.to_vec(),
+                body,
+            }),
+            Err(_) => Err(anyhow!("Encryption error")),
+        }
+    }
+
+    pub fn decrypt(&self, payload: &EncryptedPayload) -> Result<Bytes> {
+        let nonce = Nonce::from_slice(payload.nonce.as_ref());
+        match self.cipher.decrypt(nonce, payload.body.as_ref()) {
+            Ok(payload) => Ok(Bytes::from(payload)),
+            Err(_) => Err(anyhow!("Decryption error")),
+        }
     }
 }

src/handshake.rs

@@ -1,7 +1,5 @@
 use crate::conf::{HANDSHAKE_MSG_SIZE, ID_SIZE};
-use crate::crypto::new_cipher;
 
-use aes_gcm::Aes256Gcm; // Or `Aes128Gcm`
 use anyhow::{anyhow, Result};
 use blake2::{Blake2s256, Digest};
 use bytes::{Bytes, BytesMut};
@@ -29,17 +27,12 @@ impl Handshake {
 
     pub async fn from_socket(socket: TcpStream) -> Result<(Handshake, TcpStream)> {
         let mut socket = socket;
-        let mut buffer = BytesMut::with_capacity(ID_SIZE + HANDSHAKE_MSG_SIZE);
-        match socket.read_exact(&mut buffer).await? {
-            65 => Ok((Handshake::from_buffer(buffer), socket)), // magic number to catch correct capacity
-            _ => return Err(anyhow!("invalid handshake buffer pulled from socket")),
-        }
-    }
-
-    pub fn from_buffer(buffer: BytesMut) -> Handshake {
-        let mut outbound_msg = BytesMut::from(&buffer[..ID_SIZE + HANDSHAKE_MSG_SIZE]).freeze();
-        let id = outbound_msg.split_to(32);
-        Handshake { id, outbound_msg }
+        let mut buffer = [0; ID_SIZE + HANDSHAKE_MSG_SIZE];
+        let n = socket.read_exact(&mut buffer).await?;
+        println!("The bytes: {:?}", &buffer[..n]);
+        let mut outbound_msg = BytesMut::from(&buffer[..n]).freeze();
+        let id = outbound_msg.split_to(ID_SIZE);
+        Ok((Handshake { id, outbound_msg }, socket))
     }
 
     pub fn to_bytes(self) -> Bytes {
@@ -53,10 +46,11 @@ impl Handshake {
         self,
         socket: TcpStream,
         s1: spake2::Spake2<spake2::Ed25519Group>,
-    ) -> Result<(TcpStream, Aes256Gcm)> {
+    ) -> Result<(TcpStream, Vec<u8>)> {
         let mut socket = socket;
-        // println!("client - sending handshake msg");
-        socket.write_all(&self.to_bytes()).await?;
+        let bytes = &self.to_bytes();
+        println!("client - sending handshake msg= {:?}", &bytes);
+        socket.write_all(&bytes).await?;
         let mut buffer = [0; HANDSHAKE_MSG_SIZE];
         let n = socket.read_exact(&mut buffer).await?;
         let response = BytesMut::from(&buffer[..n]).freeze();
@@ -66,7 +60,7 @@ impl Handshake {
             Err(e) => return Err(anyhow!(e.to_string())),
         };
         println!("Handshake successful. Key is {:?}", key);
-        return Ok((socket, new_cipher(&key)));
+        return Ok((socket, key));
     }
 
     fn pass_to_bytes(password: &String) -> Bytes {

src/message.rs

@@ -1,7 +1,6 @@
-use crate::crypto::{decrypt, encrypt};
+use crate::crypto::Crypt;
 use crate::file::FileInfo;
 
-use aes_gcm::Aes256Gcm; // Or `Aes128Gcm`
 use anyhow::{anyhow, Result};
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
@@ -13,15 +12,7 @@ use tokio_util::codec::{Framed, LengthDelimitedCodec};
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub enum Message {
-    HandshakeMessage(HandshakePayload),
     EncryptedMessage(EncryptedPayload),
-    ErrorMessage(RuckError),
-}
-
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct HandshakePayload {
-    pub id: Bytes,
-    pub msg: Bytes,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -49,8 +40,8 @@ pub struct FileTransferPayload {
 }
 
 impl EncryptedMessage {
-    pub fn from_encrypted_message(cipher: &Aes256Gcm, payload: &EncryptedPayload) -> Result<Self> {
-        let raw = decrypt(cipher, payload)?;
+    pub fn from_encrypted_message(crypt: &Crypt, payload: &EncryptedPayload) -> Result<Self> {
+        let raw = crypt.decrypt(payload)?;
         let res = match bincode::deserialize(raw.as_ref()) {
             Ok(result) => result,
             Err(e) => {
@@ -60,7 +51,7 @@ impl EncryptedMessage {
         };
         Ok(res)
     }
-    pub fn to_encrypted_message(&self, cipher: &Aes256Gcm) -> Result<Message> {
+    pub fn to_encrypted_message(&self, crypt: &Crypt) -> Result<Message> {
         let raw = match bincode::serialize(&self) {
             Ok(result) => result,
             Err(e) => {
@@ -68,7 +59,7 @@ impl EncryptedMessage {
                 return Err(anyhow!("serialize error"));
             }
         };
-        let payload = encrypt(cipher, &raw)?;
+        let payload = crypt.encrypt(&raw)?;
         Ok(Message::EncryptedMessage(payload))
     }
 }

src/server.rs

@@ -1,4 +1,4 @@
-use crate::conf::BUFFER_SIZE;
+use crate::conf::PER_CLIENT_BUFFER;
 use crate::handshake::Handshake;
 use anyhow::{anyhow, Result};
 use bytes::{Bytes, BytesMut};
@@ -113,7 +113,7 @@ pub async fn handle_connection(
     };
     println!("Client upgraded");
     // The handshake cache should be empty for {id} at this point.
-    let mut client_buffer = BytesMut::with_capacity(BUFFER_SIZE);
+    let mut client_buffer = BytesMut::with_capacity(PER_CLIENT_BUFFER);
     loop {
         tokio::select! {
             Some(msg) = client.rx.recv() => {