Merge pull request #2 from rictorlome/encrypt

Encrypt

Cargo.lock

@@ -75,6 +75,15 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
+[[package]]
+name = "blake2"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b94ba84325db59637ffc528bbe8c7f86c02c57cff5c0e2b9b00f9a851f42f309"
+dependencies = [
+ "digest 0.10.2",
+]
+
 [[package]]
 name = "block-buffer"
 version = "0.10.2"
@@ -179,7 +188,7 @@ checksum = "0b9fdf9972b2bd6af2d913799d9ebc165ea4d2e65878e329d9c6b372c4491b61"
 dependencies = [
  "byteorder",
  "digest 0.9.0",
- "rand_core",
+ "rand_core 0.5.1",
  "subtle",
  "zeroize",
 ]
@@ -337,7 +346,18 @@ checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce"
 dependencies = [
  "cfg-if",
  "libc",
- "wasi",
+ "wasi 0.9.0+wasi-snapshot-preview1",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "418d37c8b1d42553c93648be529cb70f920d3baf8ef469b74b9638df426e0b4c"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi 0.10.2+wasi-snapshot-preview1",
 ]
 
 [[package]]
@@ -605,6 +625,12 @@ dependencies = [
  "universal-hash",
 ]
 
+[[package]]
+name = "ppv-lite86"
+version = "0.2.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872"
+
 [[package]]
 name = "proc-macro-error"
 version = "1.0.4"
@@ -647,13 +673,53 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "rand"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2e7573632e6454cf6b99d7aac4ccca54be06da05aca2ef7423d22d27d4d4bcd8"
+dependencies = [
+ "libc",
+ "rand_chacha",
+ "rand_core 0.6.3",
+ "rand_hc",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core 0.6.3",
+]
+
 [[package]]
 name = "rand_core"
 version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19"
 dependencies = [
- "getrandom",
+ "getrandom 0.1.16",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d34f1408f55294453790c48b2f1ebbb1c5b4b7563eb1f418bcfcfdbb06ebb4e7"
+dependencies = [
+ "getrandom 0.2.4",
+]
+
+[[package]]
+name = "rand_hc"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d51e9f596de227fda2ea6c84607f5558e196eeaf43c986b724ba4fb8fdf497e7"
+dependencies = [
+ "rand_core 0.6.3",
 ]
 
 [[package]]
@@ -672,9 +738,11 @@ dependencies = [
  "aes-gcm",
  "anyhow",
  "bincode",
+ "blake2",
  "bytes",
  "clap",
  "futures",
+ "rand",
  "serde",
  "spake2",
  "tokio",
@@ -749,7 +817,7 @@ checksum = "7139ade210556eca57dfc299ec2454846ab6be09232eb1139a36e285ae7fd48e"
 dependencies = [
  "curve25519-dalek",
  "hkdf",
- "rand_core",
+ "rand_core 0.5.1",
  "sha2",
 ]
 
@@ -897,6 +965,12 @@ version = "0.9.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519"
 
+[[package]]
+name = "wasi"
+version = "0.10.2+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6"
+
 [[package]]
 name = "winapi"
 version = "0.3.9"

Cargo.toml

@@ -7,10 +7,12 @@ edition = "2021"
 
 aes-gcm = "0.9.4"
 anyhow = "1.0"
+blake2 = "0.10.2"
 bytes = { version = "1", features = ["serde"] }
 bincode = "1.3.3"
 clap = { version = "3.0.14", features = ["derive"] }
 futures = { version = "0.3.0", features = ["thread-pool"]}
+rand = "0.8.4"
 serde = { version = "1.0", features = ["derive"] }
 spake2 = "0.3.1"
 tokio = { version = "1.16.1", features = ["full"] }

README.md

@@ -1,13 +1,17 @@
-https://github.com/mcginty/snow/tree/master/examples
+- https://github.com/mcginty/snow/tree/master/examples
-https://github.com/black-binary/snowstorm/tree/master/examples
+- https://github.com/black-binary/snowstorm/tree/master/examples
-https://cryptoservices.github.io/cryptography/protocols/2016/04/27/noise-protocol.html
+- https://cryptoservices.github.io/cryptography/protocols/2016/04/27/noise-protocol.html
 
-https://github.com/libp2p
+- https://github.com/libp2p
 
-https://docs.rs/spake2/latest/spake2/
+- https://docs.rs/spake2/latest/spake2/
-https://crates.io/crates/chacha20poly1305
+- https://crates.io/crates/chacha20poly1305
-https://briansmith.org/rustdoc/ring/aead/index.html
+- https://briansmith.org/rustdoc/ring/aead/index.html
-https://libsodium.gitbook.io/doc/secret-key_cryptography/secretstream
+- https://libsodium.gitbook.io/doc/secret-key_cryptography/secretstream
 
-https://kerkour.com/rust-symmetric-encryption-aead-benchmark/
+- https://kerkour.com/rust-symmetric-encryption-aead-benchmark/
-https://docs.rs/aes-gcm/latest/aes_gcm/
+- https://docs.rs/aes-gcm/latest/aes_gcm/
+
+
+- https://github.com/rust-lang/flate2-rs
+- https://crates.io/crates/async-compression

src/client.rs

@@ -1,51 +1,138 @@
 use crate::crypto::handshake;
-use crate::message::{Message, MessageStream};
+use crate::file::{to_size_string, FileHandle, FileInfo};
+use crate::message::{EncryptedMessage, FileNegotiationPayload, Message, MessageStream};
 
-use anyhow::Result;
+use aes_gcm::Aes256Gcm;
-use bytes::{BufMut, Bytes, BytesMut};
+use anyhow::{anyhow, Result};
+use blake2::{Blake2s256, Digest};
+use bytes::{Bytes, BytesMut};
+use futures::future::try_join_all;
 use futures::prelude::*;
 use std::path::PathBuf;
+use tokio::io;
 use tokio::net::TcpStream;
+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?;
+
+    // Establish connection to server
     let socket = TcpStream::connect("127.0.0.1:8080").await?;
     let mut stream = Message::to_stream(socket);
-    let (stream, key) = handshake(
+
+    // Complete handshake, returning cipher used for encryption
+    let (stream, cipher) = handshake(
         &mut stream,
+        true,
         Bytes::from(password.to_string()),
-        Bytes::from("id123"),
+        pass_to_bytes(password),
     )
     .await?;
 
-    return upload_encrypted_files(stream, file_paths, key).await;
+    // Complete file negotiation
+    let handles = negotiate_files_up(handles, stream, &cipher).await?;
 
-    // Send the value
+    // Upload negotiated files
-    // for path in paths.iter() {
+
-    //     let b = path.to_str().unwrap().as_bytes();
+    // Exit
-    //     let mut buf = BytesMut::with_capacity(1024);
+    Ok(())
-    //     buf.put(&b[..]);
-    //     let body = buf.freeze();
-    //     let m = Message {
-    //         key: "abc".to_string(),
-    //         from_sender: true,
-    //         body: body,
-    //     };
-    //     stream.send(m).await.unwrap();
-    // }
 }
 
 pub async fn receive(password: &String) -> Result<()> {
     let socket = TcpStream::connect("127.0.0.1:8080").await?;
     let mut stream = Message::to_stream(socket);
-    let (stream, key) = handshake(
+    let (stream, cipher) = handshake(
         &mut stream,
+        false,
         Bytes::from(password.to_string()),
-        Bytes::from("id123"),
+        pass_to_bytes(password),
     )
     .await?;
+
+    let files = negotiate_files_down(stream, &cipher).await?;
     return Ok(());
 }
 
+pub async fn get_file_handles(file_paths: &Vec<PathBuf>) -> Result<Vec<FileHandle>> {
+    let tasks = file_paths
+        .into_iter()
+        .map(|path| FileHandle::new(path.to_path_buf()));
+    let handles = try_join_all(tasks).await?;
+    Ok(handles)
+}
+
+pub async fn negotiate_files_up(
+    file_handles: Vec<FileHandle>,
+    stream: &mut MessageStream,
+    cipher: &Aes256Gcm,
+) -> 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)?;
+    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 requested_paths: Vec<PathBuf> = match plaintext_reply {
+        EncryptedMessage::FileNegotiationMessage(fnm) => {
+            fnm.files.into_iter().map(|f| f.path).collect()
+        }
+        _ => return Err(anyhow!("Expecting file negotiation message back")),
+    };
+    Ok(file_handles
+        .into_iter()
+        .filter(|fh| requested_paths.contains(&fh.path))
+        .collect())
+}
+
+pub async fn negotiate_files_down(stream: &mut MessageStream, cipher: &Aes256Gcm) -> 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 requested_infos: Vec<FileInfo> = match plaintext_offer {
+        EncryptedMessage::FileNegotiationMessage(fnm) => fnm.files,
+        _ => return Err(anyhow!("Expecting file negotiation message back")),
+    };
+    let mut stdin = FramedRead::new(io::stdin(), LinesCodec::new());
+    let mut files = vec![];
+    for path in requested_infos.into_iter() {
+        let mut reply = prompt_user_input(&mut stdin, &path).await;
+        while reply.is_none() {
+            reply = prompt_user_input(&mut stdin, &path).await;
+        }
+        match reply {
+            Some(true) => files.push(path),
+            _ => {}
+        }
+    }
+    let msg = EncryptedMessage::FileNegotiationMessage(FileNegotiationPayload { files });
+    let server_msg = msg.to_encrypted_message(cipher)?;
+    stream.send(server_msg).await?;
+    Ok(())
+}
+
 pub async fn upload_encrypted_files(
     stream: &mut MessageStream,
     file_paths: &Vec<PathBuf>,
@@ -53,3 +140,26 @@ pub async fn upload_encrypted_files(
 ) -> Result<()> {
     return Ok(());
 }
+
+pub async fn prompt_user_input(
+    stdin: &mut FramedRead<io::Stdin, LinesCodec>,
+    file_info: &FileInfo,
+) -> Option<bool> {
+    let prompt_name = file_info.path.file_name().unwrap();
+    println!(
+        "Do you want to download {:?}? It's {:?}. (Y/n)",
+        prompt_name,
+        to_size_string(file_info.size)
+    );
+    match stdin.next().await {
+        Some(Ok(line)) => match line.as_str() {
+            "" | "Y" | "y" | "yes" | "Yes" | "YES" => Some(true),
+            "N" | "n" | "NO" | "no" | "No" => Some(false),
+            _ => {
+                println!("Invalid input. Please enter one of the following characters: [YyNn]");
+                return None;
+            }
+        },
+        _ => None,
+    }
+}

src/crypto.rs

@@ -1,23 +1,29 @@
-use crate::message::{HandshakeMessage, Message, MessageStream};
+use crate::message::{EncryptedPayload, HandshakePayload, Message, MessageStream};
 
+use aes_gcm::aead::{Aead, NewAead};
+use aes_gcm::{Aes256Gcm, Key, Nonce}; // Or `Aes128Gcm`
 use anyhow::{anyhow, Result};
 use bytes::Bytes;
 use futures::prelude::*;
+use rand::{thread_rng, Rng};
 use spake2::{Ed25519Group, Identity, Password, Spake2};
 
 pub async fn handshake(
     stream: &mut MessageStream,
+    up: bool,
     password: Bytes,
     id: Bytes,
-) -> Result<(&mut MessageStream, Bytes)> {
+) -> Result<(&mut MessageStream, Aes256Gcm)> {
     let (s1, outbound_msg) =
         Spake2::<Ed25519Group>::start_symmetric(&Password::new(password), &Identity::new(&id));
-    stream
+    println!("client - sending handshake msg");
-        .send(Message::HandshakeMessage(HandshakeMessage {
+    let handshake_msg = Message::HandshakeMessage(HandshakePayload {
-            id,
+        up,
-            msg: Bytes::from(outbound_msg),
+        id,
-        }))
+        msg: Bytes::from(outbound_msg),
-        .await?;
+    });
+    println!("client - handshake msg, {:?}", handshake_msg);
+    stream.send(handshake_msg).await?;
     let first_message = match stream.next().await {
         Some(Ok(msg)) => match msg {
             Message::HandshakeMessage(response) => response.msg,
@@ -27,10 +33,39 @@ pub async fn handshake(
             return Err(anyhow!("No response to handshake message"));
         }
     };
+    println!("client - handshake msg responded to");
     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((stream, Bytes::from(key)));
+    return Ok((stream, new_cipher(&key)));
+}
+
+pub fn new_cipher(key: &Vec<u8>) -> Aes256Gcm {
+    let key = Key::from_slice(&key[..]);
+    Aes256Gcm::new(key)
+}
+
+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")),
+    }
 }

src/file.rs

@@ -0,0 +1,48 @@
+use anyhow::{anyhow, Result};
+use serde::{Deserialize, Serialize};
+use std::fs::Metadata;
+use std::path::PathBuf;
+use tokio::fs::File;
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct FileInfo {
+    pub path: PathBuf,
+    pub size: u64,
+}
+
+pub struct FileHandle {
+    pub file: File,
+    pub md: Metadata,
+    pub path: PathBuf,
+}
+
+impl FileHandle {
+    pub async fn new(path: PathBuf) -> Result<FileHandle> {
+        let file = File::open(&path).await?;
+        let md = file.metadata().await?;
+        let fh = FileHandle { file, md, path };
+        return Ok(fh);
+    }
+
+    pub fn to_file_info(&self) -> FileInfo {
+        FileInfo {
+            path: self.path.clone(),
+            size: self.md.len(),
+        }
+    }
+}
+
+const SUFFIX: [&'static str; 9] = ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"];
+// Stolen: https://gitlab.com/forkbomb9/human_bytes-rs/-/blob/master/src/lib.rs
+pub fn to_size_string(size: u64) -> String {
+    let size = size as f64;
+    let base = size.log10() / 1024_f64.log10();
+    let mut result = format!("{:.1}", 1024_f64.powf(base - base.floor()),)
+        .trim_end_matches(".0")
+        .to_owned();
+    // Add suffix
+    result.push(' ');
+    result.push_str(SUFFIX[base.floor() as usize]);
+
+    result
+}

src/main.rs

@@ -1,9 +1,13 @@
 mod cli;
 mod client;
 mod crypto;
+mod file;
 mod message;
 mod server;
 
+#[cfg(test)]
+mod tests;
+
 use clap::Parser;
 use cli::{Cli, Commands};
 use client::{receive, send};
@@ -16,7 +20,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     match &args.command {
         Commands::Send { paths, password } => {
             println!("Sending {:?}", paths);
-            send(&paths, password).await?;
+            send(paths, password).await?;
         }
         Commands::Receive { password } => {
             println!("Receiving password {}", password);

src/message.rs

@@ -1,20 +1,92 @@
+use crate::crypto::{decrypt, encrypt};
+use crate::file::FileInfo;
+
+use aes_gcm::Aes256Gcm; // Or `Aes128Gcm`
+use anyhow::{anyhow, Result};
+use bincode::config;
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
+use std::error::Error;
+use std::fmt;
 use tokio::net::TcpStream;
 use tokio_serde::{formats::SymmetricalBincode, SymmetricallyFramed};
 use tokio_util::codec::{Framed, LengthDelimitedCodec};
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub enum Message {
-    HandshakeMessage(HandshakeMessage),
+    HandshakeMessage(HandshakePayload),
+    EncryptedMessage(EncryptedPayload),
+    ErrorMessage(RuckError),
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct HandshakeMessage {
+pub struct HandshakePayload {
+    pub up: bool,
     pub id: Bytes,
     pub msg: Bytes,
 }
 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct EncryptedPayload {
+    pub nonce: Vec<u8>,
+    pub body: Vec<u8>,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub enum EncryptedMessage {
+    FileNegotiationMessage(FileNegotiationPayload),
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct FileNegotiationPayload {
+    pub files: Vec<FileInfo>,
+}
+
+impl EncryptedMessage {
+    pub fn from_encrypted_message(cipher: &Aes256Gcm, payload: &EncryptedPayload) -> Result<Self> {
+        let raw = decrypt(cipher, payload)?;
+        let res = match bincode::deserialize(raw.as_ref()) {
+            Ok(result) => result,
+            Err(e) => {
+                println!("deserialize error {:?}", e);
+                return Err(anyhow!("deser error"));
+            }
+        };
+        Ok(res)
+    }
+    pub fn to_encrypted_message(&self, cipher: &Aes256Gcm) -> Result<Message> {
+        let raw = match bincode::serialize(&self) {
+            Ok(result) => result,
+            Err(e) => {
+                println!("serialize error {:?}", e);
+                return Err(anyhow!("serialize error"));
+            }
+        };
+        let payload = encrypt(cipher, &raw)?;
+        Ok(Message::EncryptedMessage(payload))
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub enum RuckError {
+    NotHandshake,
+    SenderNotConnected,
+    SenderAlreadyConnected,
+    PairDisconnected,
+}
+
+impl fmt::Display for RuckError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "RuckError is here!")
+    }
+}
+
+impl Error for RuckError {
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        Some(self)
+    }
+}
+
 impl Message {
     pub fn to_stream(stream: TcpStream) -> MessageStream {
         tokio_serde::SymmetricallyFramed::new(

src/server.rs

@@ -1,66 +1,125 @@
-use crate::message::{Message, MessageStream};
+use crate::message::{HandshakePayload, Message, MessageStream, RuckError};
 
+use anyhow::{anyhow, Result};
+use bytes::Bytes;
 use futures::prelude::*;
 use std::collections::HashMap;
-use std::io;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use tokio::net::{TcpListener, TcpStream};
-use tokio::sync::{mpsc, Mutex};
+use tokio::sync::{mpsc, oneshot, Mutex};
 
 type Tx = mpsc::UnboundedSender<Message>;
 type Rx = mpsc::UnboundedReceiver<Message>;
 
 pub struct Shared {
-    rooms: HashMap<String, RoomInfo>,
+    handshakes: HashMap<Bytes, Rx>,
+    senders: HashMap<Bytes, Tx>,
+    receivers: HashMap<Bytes, Tx>,
 }
 type State = Arc<Mutex<Shared>>;
 
-struct RoomInfo {
+struct Client<'a> {
-    sender_tx: Tx,
+    up: bool,
-}
+    id: Bytes,
-
+    messages: &'a mut MessageStream,
-struct Client {
-    is_sender: bool,
-    messages: MessageStream,
     rx: Rx,
 }
 
 impl Shared {
     fn new() -> Self {
         Shared {
-            rooms: HashMap::new(),
+            handshakes: HashMap::new(),
+            senders: HashMap::new(),
+            receivers: HashMap::new(),
         }
     }
-
+    async fn relay<'a>(&self, client: &Client<'a>, message: Message) -> Result<()> {
-    // async fn broadcast(&mut self, sender: SocketAddr, message: Message) {
+        println!("in relay - got client={:?}, msg {:?}", client.id, message);
-    //     for peer in self.peers.iter_mut() {
+        match client.up {
-    //         if *peer.0 != sender {
+            true => match self.receivers.get(&client.id) {
-    //             let _ = peer.1.send(message.clone());
+                Some(tx) => {
-    //         }
+                    tx.send(message)?;
-    //     }
+                }
-    // }
+                None => {
+                    return Err(anyhow!(RuckError::PairDisconnected));
+                }
+            },
+            false => match self.senders.get(&client.id) {
+                Some(tx) => {
+                    tx.send(message)?;
+                }
+                None => {
+                    return Err(anyhow!(RuckError::PairDisconnected));
+                }
+            },
+        }
+        Ok(())
+    }
 }
 
-impl Client {
+impl<'a> Client<'a> {
-    async fn new(is_sender: bool, state: State, messages: MessageStream) -> io::Result<Client> {
+    async fn new(
+        up: bool,
+        id: Bytes,
+        state: State,
+        messages: &'a mut MessageStream,
+    ) -> Result<Client<'a>> {
         let (tx, rx) = mpsc::unbounded_channel();
-        let room_info = RoomInfo { sender_tx: tx };
+        println!("server - creating client up={:?}, id={:?}", up, id);
-        state
+        let shared = &mut state.lock().await;
-            .lock()
+        match shared.senders.get(&id) {
-            .await
+            Some(_) if up => {
-            .rooms
+                messages
-            .insert("abc".to_string(), room_info);
+                    .send(Message::ErrorMessage(RuckError::SenderAlreadyConnected))
-
+                    .await?;
+            }
+            Some(_) => {
+                println!("server - adding client to receivers");
+                shared.receivers.insert(id.clone(), tx);
+            }
+            None if up => {
+                println!("server - adding client to senders");
+                shared.senders.insert(id.clone(), tx);
+            }
+            None => {
+                messages
+                    .send(Message::ErrorMessage(RuckError::SenderNotConnected))
+                    .await?;
+            }
+        }
         Ok(Client {
-            is_sender,
+            up,
+            id,
             messages,
             rx,
         })
     }
+    async fn complete_handshake(&mut self, state: State, msg: Message) -> Result<()> {
+        match self.up {
+            true => {
+                let (tx, rx) = mpsc::unbounded_channel();
+                tx.send(msg)?;
+                state.lock().await.handshakes.insert(self.id.clone(), rx);
+            }
+            false => {
+                let shared = &mut state.lock().await;
+                if let Some(tx) = shared.senders.get(&self.id) {
+                    tx.send(msg)?;
+                }
+                if let Some(mut rx) = shared.handshakes.remove(&self.id) {
+                    drop(shared);
+                    if let Some(msg) = rx.recv().await {
+                        self.messages.send(msg).await?;
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
 }
 
-pub async fn serve() -> Result<(), Box<dyn std::error::Error>> {
+pub async fn serve() -> Result<()> {
     let addr = "127.0.0.1:8080".to_string();
     let listener = TcpListener::bind(&addr).await?;
     let state = Arc::new(Mutex::new(Shared::new()));
@@ -70,8 +129,8 @@ pub async fn serve() -> Result<(), Box<dyn std::error::Error>> {
         let state = Arc::clone(&state);
         tokio::spawn(async move {
             match handle_connection(state, stream, address).await {
-                Ok(_) => println!("ok"),
+                Ok(_) => println!("Connection complete!"),
-                Err(_) => println!("err"),
+                Err(err) => println!("Error handling connection! {:?}", err),
             }
         });
     }
@@ -81,19 +140,33 @@ pub async fn handle_connection(
     state: Arc<Mutex<Shared>>,
     socket: TcpStream,
     addr: SocketAddr,
-) -> Result<(), Box<dyn std::error::Error>> {
+) -> Result<()> {
     let mut stream = Message::to_stream(socket);
-    let first_message = match stream.next().await {
+    println!("server - new conn from {:?}", addr);
-        Some(Ok(msg)) => {
+    let handshake_payload = match stream.next().await {
-            println!("first msg: {:?}", msg);
+        Some(Ok(Message::HandshakeMessage(payload))) => payload,
-            msg
+        Some(Ok(_)) => {
+            stream
+                .send(Message::ErrorMessage(RuckError::NotHandshake))
+                .await?;
+            return Ok(());
         }
         _ => {
-            println!("no first message");
+            println!("No first message");
             return Ok(());
         }
     };
-    let mut client = Client::new(true, state.clone(), stream).await?;
+    println!("server - received msg from {:?}", addr);
+    let mut client = Client::new(
+        handshake_payload.up,
+        handshake_payload.id.clone(),
+        state.clone(),
+        &mut stream,
+    )
+    .await?;
+    client
+        .complete_handshake(state.clone(), Message::HandshakeMessage(handshake_payload))
+        .await?;
     // add client to state here
     loop {
         tokio::select! {
@@ -104,6 +177,8 @@ pub async fn handle_connection(
             result = client.messages.next() => match result {
                 Some(Ok(msg)) => {
                     println!("GOT: {:?}", msg);
+                    let state = state.lock().await;
+                    state.relay(&client, msg).await?;
                 }
                 Some(Err(e)) => {
                     println!("Error {:?}", e);

src/tests.rs

@@ -0,0 +1,11 @@
+use crate::crypto::new_cipher;
+use crate::file::FileHandle;
+
+use std::path::PathBuf;
+
+#[tokio::test]
+async fn test_file_handle_nonexistent_file() {
+    let pb = PathBuf::new();
+    let fh = FileHandle::new(pb).await;
+    assert!(fh.is_err());
+}