Handshake refactored not working

src/client.rs

@@ -1,13 +1,12 @@
 use crate::conf::BUFFER_SIZE;
-use crate::crypto::handshake;
 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 blake2::{Blake2s256, Digest};
 use bytes::{Bytes, BytesMut};
 use futures::future::try_join_all;
 use futures::prelude::*;
@@ -20,13 +19,6 @@ use tokio::net::TcpStream;
 use tokio::sync::mpsc;
 use tokio_util::codec::{FramedRead, LinesCodec};
 
-fn pass_to_bytes(password: &String) -> Bytes {
-    let mut hasher = Blake2s256::new();
-    hasher.update(password.as_bytes());
-    let res = hasher.finalize();
-    BytesMut::from(&res[..]).freeze()
-}
-
 pub async fn send(file_paths: &Vec<PathBuf>, password: &String) -> Result<()> {
     // Fail early if there are problems generating file handles
     let handles = get_file_handles(file_paths).await?;
@@ -35,12 +27,9 @@ pub async fn send(file_paths: &Vec<PathBuf>, password: &String) -> Result<()> {
     let socket = TcpStream::connect("127.0.0.1:8080").await?;
 
     // Complete handshake, returning cipher used for encryption
-    let (socket, cipher) = handshake(
-        socket,
-        Bytes::from(password.to_string()),
-        pass_to_bytes(password),
-    )
-    .await?;
+    let (handshake, s1) = Handshake::from_password(password);
+    let (socket, cipher) = handshake.negotiate(socket, s1).await?;
+
     let mut stream = Message::to_stream(socket);
     // Complete file negotiation
     let handles = negotiate_files_up(handles, &mut stream, &cipher).await?;
@@ -54,12 +43,8 @@ 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 (socket, cipher) = handshake(
-        socket,
-        Bytes::from(password.to_string()),
-        pass_to_bytes(password),
-    )
-    .await?;
+    let (handshake, s1) = Handshake::from_password(password);
+    let (socket, cipher) = handshake.negotiate(socket, s1).await?;
     let mut stream = Message::to_stream(socket);
     let files = negotiate_files_down(&mut stream, &cipher).await?;
 

src/conf.rs

@@ -1 +1,3 @@
+pub const ID_SIZE: usize = 32;
+pub const HANDSHAKE_MSG_SIZE: usize = 33;
 pub const BUFFER_SIZE: usize = 1024 * 1024 * 10;

src/crypto.rs

@@ -1,42 +1,39 @@
 use crate::message::EncryptedPayload;
-
 use aes_gcm::aead::{Aead, NewAead};
 use aes_gcm::{Aes256Gcm, Key, Nonce}; // Or `Aes128Gcm`
 use anyhow::{anyhow, Result};
-use bytes::{Bytes, BytesMut};
-use rand::{thread_rng, Rng};
-use spake2::{Ed25519Group, Identity, Password, Spake2};
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::TcpStream;
+use bytes::Bytes;
 
-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)));
-}
+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[..]);

src/handshake.rs

@@ -0,0 +1,78 @@
+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};
+use spake2::{Ed25519Group, Identity, Password, Spake2};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::TcpStream;
+
+pub struct Handshake {
+    pub id: Bytes,
+    pub outbound_msg: Bytes,
+}
+
+impl Handshake {
+    pub fn from_password(pw: &String) -> (Handshake, spake2::Spake2<spake2::Ed25519Group>) {
+        let password = Bytes::from(pw.to_string());
+        let id = Handshake::pass_to_bytes(&pw);
+        let (s1, outbound_msg) =
+            Spake2::<Ed25519Group>::start_symmetric(&Password::new(&password), &Identity::new(&id));
+        let mut buffer = BytesMut::with_capacity(HANDSHAKE_MSG_SIZE);
+        buffer.extend_from_slice(&outbound_msg[..HANDSHAKE_MSG_SIZE]);
+        let outbound_msg = buffer.freeze();
+        let handshake = Handshake { id, outbound_msg };
+        (handshake, s1)
+    }
+
+    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 }
+    }
+
+    pub fn to_bytes(self) -> Bytes {
+        let mut buffer = BytesMut::with_capacity(ID_SIZE + HANDSHAKE_MSG_SIZE);
+        buffer.extend_from_slice(&self.id);
+        buffer.extend_from_slice(&self.outbound_msg);
+        buffer.freeze()
+    }
+
+    pub async fn negotiate(
+        self,
+        socket: TcpStream,
+        s1: spake2::Spake2<spake2::Ed25519Group>,
+    ) -> Result<(TcpStream, Aes256Gcm)> {
+        let mut socket = socket;
+        // println!("client - sending handshake msg");
+        socket.write_all(&self.to_bytes()).await?;
+        let mut buffer = [0; HANDSHAKE_MSG_SIZE];
+        let n = socket.read_exact(&mut buffer).await?;
+        let response = BytesMut::from(&buffer[..n]).freeze();
+        // println!("client - handshake msg, {:?}", response);
+        let key = match s1.finish(&response[..]) {
+            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)));
+    }
+
+    fn pass_to_bytes(password: &String) -> Bytes {
+        let mut hasher = Blake2s256::new();
+        hasher.update(password.as_bytes());
+        let res = hasher.finalize();
+        BytesMut::from(&res[..]).freeze()
+    }
+}

src/main.rs

@@ -3,6 +3,7 @@ mod client;
 mod conf;
 mod crypto;
 mod file;
+mod handshake;
 mod message;
 mod server;
 

src/server.rs

@@ -1,4 +1,5 @@
 use crate::conf::BUFFER_SIZE;
+use crate::handshake::Handshake;
 use anyhow::{anyhow, Result};
 use bytes::{Bytes, BytesMut};
 use std::collections::HashMap;
@@ -48,12 +49,7 @@ impl Client {
         Ok(client)
     }
 
-    async fn upgrade(
-        client: Client,
-        state: State,
-        id: Bytes,
-        handshake_msg: Bytes,
-    ) -> Result<StapledClient> {
+    async fn upgrade(client: Client, state: State, handshake: Handshake) -> Result<StapledClient> {
         let mut client = client;
         let peer_tx = match client.peer_tx {
             // Receiver - already stapled at creation
@@ -65,13 +61,13 @@ impl Client {
                         client.socket.write_all(&msg[..]).await?
                     }
                 }
-                match state.lock().await.handshake_cache.remove(&id) {
+                match state.lock().await.handshake_cache.remove(&handshake.id) {
                     Some(peer_tx) => peer_tx,
                     None => return Err(anyhow!("Connection not stapled")),
                 }
             }
         };
-        peer_tx.send(handshake_msg)?;
+        peer_tx.send(handshake.outbound_msg)?;
         Ok(StapledClient {
             socket: client.socket,
             rx: client.rx,
@@ -103,17 +99,11 @@ pub async fn handle_connection(
     _addr: SocketAddr,
 ) -> Result<()> {
     socket.readable().await?;
-    let mut socket = socket;
-    // let mut handshake_buffer = BytesMut::with_capacity(55);
-    let mut buffer = [0; 32 + 33];
-    let n = socket.read_exact(&mut buffer).await?;
-    println!("The bytes: {:?}", &buffer[..n]);
-    let mut msg = BytesMut::from(&buffer[..n]).freeze();
-    let id = msg.split_to(32);
-    println!("New client with id={:?}, msg={:?}", id.clone(), msg.clone());
+    let (handshake, socket) = Handshake::from_socket(socket).await?;
+    let id = handshake.id.clone();
     let client = Client::new(id.clone(), state.clone(), socket).await?;
     println!("Client created");
-    let mut client = match Client::upgrade(client, state.clone(), id.clone(), msg).await {
+    let mut client = match Client::upgrade(client, state.clone(), handshake).await {
         Ok(client) => client,
         Err(err) => {
             // Clear handshake cache if staple is unsuccessful