diff --git a/src/client.rs b/src/client.rs index 0e37446..bf15038 100644 --- a/src/client.rs +++ b/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, 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, 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) -> Result, stream: &mut MessageStream, - cipher: &Aes256Gcm, + crypt: &Crypt, ) -> Result> { 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 = 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> { 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 = 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, - cipher: &Aes256Gcm, + cipher: &Crypt, ) -> Result<()> { let (tx, mut rx) = mpsc::unbounded_channel::(); for mut handle in handles { @@ -188,7 +189,7 @@ pub async fn enqueue_file_chunks( pub async fn download_files( file_infos: Vec, stream: &mut MessageStream, - cipher: &Aes256Gcm, + cipher: &Crypt, ) -> Result<()> { // for each file_info let mut info_handles: HashMap> = 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); } diff --git a/src/conf.rs b/src/conf.rs index 9ab88ad..6fb5d0e 100644 --- a/src/conf.rs +++ b/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 diff --git a/src/connection.rs b/src/connection.rs new file mode 100644 index 0000000..413ad0b --- /dev/null +++ b/src/connection.rs @@ -0,0 +1,4 @@ +pub struct Connection { + ms: MessageStream; + crypt: Crypt +} \ No newline at end of file diff --git a/src/crypto.rs b/src/crypto.rs index 60dd513..6fb97d6 100644 --- a/src/crypto.rs +++ b/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::::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) -> 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) -> Result { - 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 { - 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) -> Crypt { + let key = Key::from_slice(&key[..]); + Crypt { + cipher: Aes256Gcm::new(key), + } + } + + pub fn encrypt(&self, body: &Vec) -> Result { + 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 { + 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")), + } } } diff --git a/src/handshake.rs b/src/handshake.rs index 1a04fa5..1b1b409 100644 --- a/src/handshake.rs +++ b/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, - ) -> Result<(TcpStream, Aes256Gcm)> { + ) -> Result<(TcpStream, Vec)> { 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 { diff --git a/src/message.rs b/src/message.rs index 10593a2..b815a1d 100644 --- a/src/message.rs +++ b/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 { - let raw = decrypt(cipher, payload)?; + pub fn from_encrypted_message(crypt: &Crypt, payload: &EncryptedPayload) -> Result { + 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 { + pub fn to_encrypted_message(&self, crypt: &Crypt) -> Result { 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)) } } diff --git a/src/server.rs b/src/server.rs index 85e8720..3509c47 100644 --- a/src/server.rs +++ b/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() => {