Merge pull request #4 from rictorlome/speed-up

Speed up

.gitignore

@@ -1 +1,2 @@
 /target
+Notes.md

Cargo.lock

@@ -2,6 +2,12 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "adler"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
+
 [[package]]
 name = "aead"
 version = "0.4.3"
@@ -162,6 +168,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "crc32fast"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "crypto-common"
 version = "0.1.2"
@@ -214,28 +229,13 @@ dependencies = [
 ]
 
 [[package]]
-name = "educe"
+name = "flate2"
-version = "0.4.18"
+version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f86b50932a01e7ec5c06160492ab660fb19b6bb2a7878030dd6cd68d21df9d4d"
+checksum = "f82b0f4c27ad9f8bfd1f3208d882da2b09c301bc1c828fd3a00d0216d2fbbff6"
 dependencies = [
- "enum-ordinalize",
+ "crc32fast",
- "proc-macro2",
+ "miniz_oxide",
- "quote",
- "syn",
-]
-
-[[package]]
-name = "enum-ordinalize"
-version = "3.1.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0b166c9e378360dd5a6666a9604bb4f54ae0cac39023ffbac425e917a2a04fef"
-dependencies = [
- "num-bigint",
- "num-traits",
- "proc-macro2",
- "quote",
- "syn",
 ]
 
 [[package]]
@@ -464,6 +464,15 @@ version = "2.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "308cc39be01b73d0d18f82a0e7b2a3df85245f84af96fdddc5d202d27e47b86a"
 
+[[package]]
+name = "miniz_oxide"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f5c75688da582b8ffc1f1799e9db273f32133c49e048f614d22ec3256773ccc"
+dependencies = [
+ "adler",
+]
+
 [[package]]
 name = "mio"
 version = "0.7.14"
@@ -495,36 +504,6 @@ dependencies = [
  "winapi",
 ]
 
-[[package]]
-name = "num-bigint"
-version = "0.4.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f93ab6289c7b344a8a9f60f88d80aa20032336fe78da341afc91c8a2341fc75f"
-dependencies = [
- "autocfg",
- "num-integer",
- "num-traits",
-]
-
-[[package]]
-name = "num-integer"
-version = "0.1.44"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d2cc698a63b549a70bc047073d2949cce27cd1c7b0a4a862d08a8031bc2801db"
-dependencies = [
- "autocfg",
- "num-traits",
-]
-
-[[package]]
-name = "num-traits"
-version = "0.2.14"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9a64b1ec5cda2586e284722486d802acf1f7dbdc623e2bfc57e65ca1cd099290"
-dependencies = [
- "autocfg",
-]
-
 [[package]]
 name = "num_cpus"
 version = "1.13.1"
@@ -581,26 +560,6 @@ dependencies = [
  "winapi",
 ]
 
-[[package]]
-name = "pin-project"
-version = "1.0.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "58ad3879ad3baf4e44784bc6a718a8698867bb991f8ce24d1bcbe2cfb4c3a75e"
-dependencies = [
- "pin-project-internal",
-]
-
-[[package]]
-name = "pin-project-internal"
-version = "1.0.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "744b6f092ba29c3650faf274db506afd39944f48420f6c86b17cfe0ee1cb36bb"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn",
-]
-
 [[package]]
 name = "pin-project-lite"
 version = "0.2.8"
@@ -741,13 +700,12 @@ dependencies = [
  "blake2",
  "bytes",
  "clap",
+ "flate2",
  "futures",
  "rand",
  "serde",
  "spake2",
  "tokio",
- "tokio-serde",
- "tokio-stream",
  "tokio-util",
 ]
 
@@ -889,32 +847,6 @@ dependencies = [
  "syn",
 ]
 
-[[package]]
-name = "tokio-serde"
-version = "0.8.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "911a61637386b789af998ee23f50aa30d5fd7edcec8d6d3dedae5e5815205466"
-dependencies = [
- "bincode",
- "bytes",
- "educe",
- "futures-core",
- "futures-sink",
- "pin-project",
- "serde",
-]
-
-[[package]]
-name = "tokio-stream"
-version = "0.1.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "50145484efff8818b5ccd256697f36863f587da82cf8b409c53adf1e840798e3"
-dependencies = [
- "futures-core",
- "pin-project-lite",
- "tokio",
-]
-
 [[package]]
 name = "tokio-util"
 version = "0.6.9"

Cargo.toml

@@ -11,11 +11,10 @@ blake2 = "0.10.2"
 bytes = { version = "1", features = ["serde"] }
 bincode = "1.3.3"
 clap = { version = "3.0.14", features = ["derive"] }
+flate2 = "1.0"
 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"] }
-tokio-serde = { version = "~0.8", features = ["bincode"] }
-tokio-stream = { version = "~0.1.6", features = ["net"]}
 tokio-util = { version = "0.6.3", features = ["full"]}

README.md

@@ -1,30 +1,35 @@
-- https://github.com/mcginty/snow/tree/master/examples
+# ruck
-- https://github.com/black-binary/snowstorm/tree/master/examples
-- https://cryptoservices.github.io/cryptography/protocols/2016/04/27/noise-protocol.html
 
-- https://github.com/libp2p
+`ruck` is a command line tool used for hosting relay servers and sending end-to-end encrypted files between clients. It was heavily inspired by [croc](https://github.com/schollz/croc), one of the easiest ways to send files between peers. This document describes the protocol `ruck` uses to support this functionality.
 
-- https://docs.rs/spake2/latest/spake2/
+### Version
-- https://crates.io/crates/chacha20poly1305
-- https://briansmith.org/rustdoc/ring/aead/index.html
-- https://libsodium.gitbook.io/doc/secret-key_cryptography/secretstream
 
-- https://kerkour.com/rust-symmetric-encryption-aead-benchmark/
+This document refers to version `0.1.0` of `ruck` as defined by the `Cargo.toml` file.
-- https://docs.rs/aes-gcm/latest/aes_gcm/
 
-- https://github.com/rust-lang/flate2-rs
+## Server
-- https://crates.io/crates/async-compression
 
-- https://pdos.csail.mit.edu/papers/chord:sigcomm01/chord_sigcomm.pdf
+The server in `ruck` exposes a TCP port.
-- https://pdos.csail.mit.edu/6.824/schedule.html
+Its only functions are to staple connections and shuttle bytes between stapled connections.
-- https://flylib.com/books/en/2.292.1.105/1/
+The first 32 bytes sent over the wire from a new client are used as its unique identifier.
-- https://en.wikipedia.org/wiki/Two_Generals%27_Problem
+When a new client joins, if the server has another open connection with the same identifier, the connections are then stapled.
+The clients have some mechanism for agreeing on these identifiers, however, from the server's perspective it doesn't matter how they agree.
 
-- https://kobzol.github.io/rust/ci/2021/05/07/building-rust-binaries-in-ci-that-work-with-older-glibc.html
+Once the connection is stapled, all bytes are piped across until a client disconnects or times out.
+The time out is set to remove idle connections.
+The server does nothing else with the bytes, so the clients are free to end-to-end encrypt their messages.
+For this reason, updates to the `ruck` protocol do not typically necessitate server redeployments.
 
-## Todo
+## Client
 
-- [ ] Use tracing
+There are two types of clients - `send` and `receive` clients.
-- [ ] Add progress bars
+Out of band, the clients agree on a relay server and password, from which they can derive the 32 byte identifier used by the server to staple their connections.
-- [ ] Exit happily when transfer is complete
+Clients have the option of using the single-use, automatically generated passwords which `ruck` supplies by default.
-- [ ] Compress files
+Using the passwords per the [Spake2](https://docs.rs/spake2/0.3.1/spake2/) handshake algorithm, clients generate a symmetric key with which to encrypt their subsequent messages.
+Once the handshake is complete, `send` and `receive` negotiate and exchange files per the following:
+
+- `send` offers a list of files and waits.
+- `receive` specifies which bytes it wants from these files.
+- `send` sends the specified bytes and waits.
+- `receive` sends heartbeats with progress updates.
+- `send` hangs up once the heartbeats stop or received a successful heartbeat.
+- `receive` hangs up once the downloads are complete.

src/client.rs

@@ -1,295 +1,107 @@
-use crate::conf::BUFFER_SIZE;
+use crate::connection::Connection;
-use crate::crypto::handshake;
+use crate::file::{ChunkHeader, FileHandle, FileOffer, StdFileHandle};
-use crate::file::{to_size_string, FileHandle, FileInfo};
+use crate::handshake::Handshake;
-use crate::message::{
+use crate::message::{FileOfferPayload, FileRequestPayload, Message};
-    EncryptedMessage, FileNegotiationPayload, FileTransferPayload, Message, MessageStream,
+use crate::ui::prompt_user_for_file_confirmation;
-};
 
-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::*;
-use std::collections::HashMap;
 use std::ffi::OsStr;
 use std::path::PathBuf;
 use tokio::fs::File;
-use tokio::io::{self, AsyncReadExt, AsyncWriteExt};
-use tokio::net::TcpStream;
-use tokio::sync::mpsc;
-use tokio_util::codec::{FramedRead, LinesCodec};
 
-fn pass_to_bytes(password: &String) -> Bytes {
+use tokio::net::TcpStream;
-    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?;
+    let handles = FileHandle::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);
 
-    // Complete handshake, returning cipher used for encryption
+    let (handshake, s1) = Handshake::from_password(password);
-    let (stream, cipher) = handshake(
+    // Complete handshake, returning key used for encryption
-        &mut stream,
+    let (socket, key) = handshake.negotiate(socket, s1).await?;
-        Bytes::from(password.to_string()),
-        pass_to_bytes(password),
-    )
-    .await?;
 
-    // Complete file negotiation
+    let mut connection = Connection::new(socket, key);
-    let handles = negotiate_files_up(handles, stream, &cipher).await?;
+    // Offer files, wait for requested file response
+    let requested_chunks = offer_files(&mut connection, &handles).await?;
 
     // Upload negotiated files
-    upload_encrypted_files(stream, handles, &cipher).await?;
+    let std_file_handles = FileHandle::to_stds(handles, requested_chunks).await;
+    connection.upload_files(std_file_handles).await?;
+    println!("Done uploading.");
 
     // Exit
     Ok(())
 }
 
 pub async fn receive(password: &String) -> Result<()> {
+    // Establish connection to server
     let socket = TcpStream::connect("127.0.0.1:8080").await?;
-    let mut stream = Message::to_stream(socket);
+    let (handshake, s1) = Handshake::from_password(password);
-    let (stream, cipher) = handshake(
+    // Complete handshake, returning key used for encryption
-        &mut stream,
+    let (socket, key) = handshake.negotiate(socket, s1).await?;
-        Bytes::from(password.to_string()),
+    let mut connection = Connection::new(socket, key);
-        pass_to_bytes(password),
+    // Wait for offered files, respond with desired files
-    )
+    let desired_files = request_specific_files(&mut connection).await?;
-    .await?;
+    // Create files
-
+    let std_file_handles = create_files(desired_files).await?;
-    let files = negotiate_files_down(stream, &cipher).await?;
+    // Download them
-
+    connection.download_files(std_file_handles).await?;
-    download_files(files, stream, &cipher).await?;
     return Ok(());
 }
 
-pub async fn get_file_handles(file_paths: &Vec<PathBuf>) -> Result<Vec<FileHandle>> {
+pub async fn offer_files(
-    let tasks = file_paths
+    conn: &mut Connection,
-        .into_iter()
+    file_handles: &Vec<FileHandle>,
-        .map(|path| FileHandle::new(path.to_path_buf()));
+) -> Result<Vec<ChunkHeader>> {
-    let handles = try_join_all(tasks).await?;
+    // Collect file offer
-    Ok(handles)
+    let files = file_handles.iter().map(|fh| fh.to_file_offer()).collect();
-}
+    let msg = Message::FileOffer(FileOfferPayload { files });
-
+    // Send file offer
-pub async fn negotiate_files_up(
+    conn.send_msg(msg).await?;
-    file_handles: Vec<FileHandle>,
+    // Wait for reply
-    stream: &mut MessageStream,
+    let reply = conn.await_msg().await?;
-    cipher: &Aes256Gcm,
+    // Return requested chunks
-) -> Result<Vec<FileHandle>> {
+    match reply {
-    let files = file_handles.iter().map(|fh| fh.to_file_info()).collect();
+        Message::FileRequest(file_request_payload) => Ok(file_request_payload.chunks),
-    let msg = EncryptedMessage::FileNegotiationMessage(FileNegotiationPayload { files });
+        _ => Err(anyhow!("Expecting file request message back")),
-    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<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 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 file_info in requested_infos.into_iter() {
-        let mut reply = prompt_user_input(&mut stdin, &file_info).await;
-        while reply.is_none() {
-            reply = prompt_user_input(&mut stdin, &file_info).await;
-        }
-        match reply {
-            Some(true) => files.push(file_info),
-            _ => {}
-        }
     }
-    let msg = EncryptedMessage::FileNegotiationMessage(FileNegotiationPayload {
+}
-        files: files.clone(),
+
+pub async fn request_specific_files(conn: &mut Connection) -> Result<Vec<FileOffer>> {
+    // Wait for offer message
+    let offer_message = conn.await_msg().await?;
+    let offered_files: Vec<FileOffer> = match offer_message {
+        Message::FileOffer(file_offer_payload) => file_offer_payload.files,
+        _ => return Err(anyhow!("Expecting file offer message")),
+    };
+    // Prompt user for confirmation of files
+    let desired_files = prompt_user_for_file_confirmation(offered_files).await;
+    let file_request_msg = Message::FileRequest(FileRequestPayload {
+        chunks: desired_files
+            .iter()
+            .map(|file| ChunkHeader {
+                id: file.id,
+                start: 0,
+            })
+            .collect(),
     });
-    let server_msg = msg.to_encrypted_message(cipher)?;
+    conn.send_msg(file_request_msg).await?;
-    stream.send(server_msg).await?;
+    Ok(desired_files)
-    Ok(files)
 }
 
-pub async fn upload_encrypted_files(
+pub async fn create_files(desired_files: Vec<FileOffer>) -> Result<Vec<StdFileHandle>> {
-    stream: &mut MessageStream,
+    let mut v = Vec::new();
-    handles: Vec<FileHandle>,
+    for desired_file in desired_files {
-    cipher: &Aes256Gcm,
+        let filename = desired_file
-) -> Result<()> {
+            .path
-    let (tx, mut rx) = mpsc::unbounded_channel::<EncryptedMessage>();
+            .file_name()
-    for mut handle in handles {
+            .unwrap_or(OsStr::new("random.txt"));
-        let txc = tx.clone();
+        let file = File::create(filename).await?;
-        tokio::spawn(async move {
+        let std_file_handle = StdFileHandle::new(desired_file.id, file, 0).await?;
-            let _ = enqueue_file_chunks(&mut handle, txc).await;
+        v.push(std_file_handle)
-        });
-    }
-
-    loop {
-        tokio::select! {
-            Some(msg) = rx.recv() => {
-                // println!("message received to client.rx {:?}", msg);
-                let x = msg.to_encrypted_message(cipher)?;
-                stream.send(x).await?
-            }
-            else => {
-                println!("breaking");
-                break
-            },
-        }
-    }
-    Ok(())
-}
-
-pub async fn enqueue_file_chunks(
-    fh: &mut FileHandle,
-    tx: mpsc::UnboundedSender<EncryptedMessage>,
-) -> Result<()> {
-    let mut chunk_num = 0;
-    let mut bytes_read = 1;
-    while bytes_read != 0 {
-        let mut buf = BytesMut::with_capacity(BUFFER_SIZE);
-        bytes_read = fh.file.read_buf(&mut buf).await?;
-        // println!("Bytes_read: {:?}, The bytes: {:?}", bytes_read, &buf[..]);
-        if bytes_read != 0 {
-            let chunk = buf.freeze();
-            let file_info = fh.to_file_info();
-            let ftp = EncryptedMessage::FileTransferMessage(FileTransferPayload {
-                chunk,
-                chunk_num,
-                file_info,
-            });
-            tx.send(ftp)?;
-            chunk_num += 1;
-        }
-    }
-
-    Ok(())
-}
-
-pub async fn download_files(
-    file_infos: Vec<FileInfo>,
-    stream: &mut MessageStream,
-    cipher: &Aes256Gcm,
-) -> Result<()> {
-    // for each file_info
-    let mut info_handles: HashMap<PathBuf, mpsc::UnboundedSender<(u64, Bytes)>> = HashMap::new();
-    for fi in file_infos {
-        let (tx, rx) = mpsc::unbounded_channel::<(u64, Bytes)>();
-        let path = fi.path.clone();
-        tokio::spawn(async move { download_file(fi, rx).await });
-        info_handles.insert(path, tx);
-    }
-    loop {
-        tokio::select! {
-            result = stream.next() => match result {
-                Some(Ok(Message::EncryptedMessage(payload))) => {
-                    let ec = EncryptedMessage::from_encrypted_message(cipher, &payload)?;
-                    // println!("encrypted message received! {:?}", ec);
-                    match ec {
-                        EncryptedMessage::FileTransferMessage(payload) => {
-                            println!("matched file transfer message");
-                            if let Some(tx) = info_handles.get(&payload.file_info.path) {
-                                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);
-                }
-                None => break,
-            }
-        }
-    }
-    Ok(())
-}
-
-pub async fn download_file(
-    file_info: FileInfo,
-    rx: mpsc::UnboundedReceiver<(u64, Bytes)>,
-) -> Result<()> {
-    println!("in download file");
-    let mut rx = rx;
-    let filename = match file_info.path.file_name() {
-        Some(f) => {
-            println!("matched filename");
-            f
-        }
-        None => {
-            println!("didnt match filename");
-            OsStr::new("random.txt")
-        }
-    };
-    println!("trying to create file...filename={:?}", filename);
-    let mut file = File::create(filename).await?;
-    println!("file created ok! filename={:?}", filename);
-    while let Some((_chunk_num, chunk)) = rx.recv().await {
-        // println!("rx got message! chunk={:?}", chunk);
-        file.write_all(&chunk).await?;
-    }
-    println!("done receiving messages");
-    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!(
-        "Accept {:?}? ({:?}). (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,
     }
+    return Ok(v);
 }

src/conf.rs

@@ -1 +1,5 @@
-pub const BUFFER_SIZE: usize = 1024 * 1024;
+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 * 64; // buffer size allocated by server for each client
+pub const BUFFER_SIZE: usize = 1024 * 64; // chunk size for files sent over wire
+pub const NONCE_SIZE: usize = 96 / 8; // used for every encrypted message

src/connection.rs

@@ -0,0 +1,128 @@
+use crate::conf::BUFFER_SIZE;
+use crate::crypto::Crypt;
+use crate::file::{ChunkHeader, StdFileHandle};
+use crate::message::{FileTransferPayload, Message, MessageStream};
+use anyhow::{anyhow, Result};
+
+use futures::{SinkExt, StreamExt};
+use tokio::net::TcpStream;
+
+use bytes::{Bytes, BytesMut};
+use flate2::bufread::GzEncoder;
+use flate2::write::GzDecoder;
+use flate2::Compression;
+use std::io::{BufReader, Read, Write};
+use std::time::Instant;
+
+pub struct Connection {
+    ms: MessageStream,
+    crypt: Crypt,
+}
+
+impl Connection {
+    pub fn new(socket: TcpStream, key: Vec<u8>) -> Self {
+        let ms = Message::to_stream(socket);
+        let crypt = Crypt::new(&key);
+        Connection { ms, crypt }
+    }
+
+    pub 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<Message> {
+        match self.ms.next().await {
+            Some(Ok(msg)) => {
+                let decrypted_bytes = self.crypt.decrypt(msg.freeze())?;
+                Message::deserialize(decrypted_bytes)
+            }
+            Some(Err(e)) => Err(anyhow!(e.to_string())),
+            None => Err(anyhow!("Error awaiting msg")),
+        }
+    }
+
+    pub async fn upload_file(&mut self, handle: StdFileHandle) -> Result<()> {
+        let mut buffer = [0; BUFFER_SIZE];
+        let reader = BufReader::new(handle.file);
+        let mut gz = GzEncoder::new(reader, Compression::fast());
+        let before = Instant::now();
+        let mut count = 0;
+        let mut bytes_sent: u64 = 0;
+        loop {
+            count += 1;
+            match gz.read(&mut buffer) {
+                Ok(0) => {
+                    break;
+                }
+                Ok(n) => {
+                    bytes_sent += n as u64;
+                    let message = Message::FileTransfer(FileTransferPayload {
+                        chunk: BytesMut::from(&buffer[..n]).freeze(),
+                        chunk_header: ChunkHeader {
+                            id: handle.id,
+                            start: 0,
+                        },
+                    });
+                    self.send_msg(message).await?;
+                }
+                Err(e) => return Err(anyhow!(e.to_string())),
+            }
+        }
+        let elapsed = before.elapsed();
+        let mb_sent = bytes_sent / 1_048_576;
+        println!(
+            "{:?} mb sent, {:?} iterations. {:?} total time, {:?} avg per iteration, {:?} avg mb/sec",
+            mb_sent,
+            count,
+            elapsed,
+            elapsed / count,
+            mb_sent / elapsed.as_secs()
+        );
+        Ok(())
+    }
+
+    pub async fn upload_files(mut self, handles: Vec<StdFileHandle>) -> Result<()> {
+        for handle in handles {
+            self.upload_file(handle).await?;
+        }
+        Ok(())
+    }
+
+    pub async fn download_files(mut self, handles: Vec<StdFileHandle>) -> Result<()> {
+        for handle in handles {
+            self.download_file(handle).await?;
+        }
+        Ok(())
+    }
+
+    pub async fn download_file(&mut self, handle: StdFileHandle) -> Result<()> {
+        let mut decoder = GzDecoder::new(handle.file);
+        loop {
+            let msg = self.await_msg().await?;
+            match msg {
+                Message::FileTransfer(payload) => {
+                    if payload.chunk_header.id != handle.id {
+                        return Err(anyhow!("Wrong file"));
+                    }
+                    if payload.chunk.len() == 0 {
+                        break;
+                    }
+                    decoder.write_all(&payload.chunk[..])?
+                }
+                _ => return Err(anyhow!("Expecting file transfer message")),
+            }
+        }
+        decoder.finish()?;
+        println!("Done downloading file.");
+        Ok(())
+    }
+}

src/crypto.rs

@@ -1,69 +1,48 @@
-use crate::message::{EncryptedPayload, HandshakePayload, Message, MessageStream};
+use crate::conf::NONCE_SIZE;
-
 use aes_gcm::aead::{Aead, NewAead};
 use aes_gcm::{Aes256Gcm, Key, Nonce}; // Or `Aes128Gcm`
 use anyhow::{anyhow, Result};
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
-use futures::prelude::*;
+
 use rand::{thread_rng, Rng};
-use spake2::{Ed25519Group, Identity, Password, Spake2};
 
-pub async fn handshake(
+pub struct Crypt {
-    stream: &mut MessageStream,
+    cipher: Aes256Gcm,
-    password: Bytes,
+    arr: [u8; NONCE_SIZE],
-    id: Bytes,
+}
-) -> Result<(&mut MessageStream, Aes256Gcm)> {
+
-    let (s1, outbound_msg) =
+impl Crypt {
-        Spake2::<Ed25519Group>::start_symmetric(&Password::new(password), &Identity::new(&id));
+    pub fn new(key: &Vec<u8>) -> Crypt {
-    println!("client - sending handshake msg");
+        let key = Key::from_slice(&key[..]);
-    let handshake_msg = Message::HandshakeMessage(HandshakePayload {
+        Crypt {
-        id,
+            cipher: Aes256Gcm::new(key),
-        msg: Bytes::from(outbound_msg),
+            arr: [0u8; NONCE_SIZE],
-    });
-    // 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,
-            _ => return Err(anyhow!("Expecting handshake message response")),
-        },
-        _ => {
-            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, new_cipher(&key)));
-}
 
-pub fn new_cipher(key: &Vec<u8>) -> Aes256Gcm {
+    // Returns wire format, includes nonce as prefix
-    let key = Key::from_slice(&key[..]);
+    pub fn encrypt(&mut self, plaintext: Bytes) -> Result<Bytes> {
-    Aes256Gcm::new(key)
+        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 + body.len());
+                buffer.extend_from_slice(nonce);
+                buffer.extend_from_slice(&body);
+                Ok(buffer.freeze())
+            }
+            Err(e) => Err(anyhow!(e.to_string())),
+        }
+    }
 
-pub const NONCE_SIZE_IN_BYTES: usize = 96 / 8;
+    // Accepts wire format, includes nonce as prefix
-pub fn encrypt(cipher: &Aes256Gcm, body: &Vec<u8>) -> Result<EncryptedPayload> {
+    pub fn decrypt(&self, ciphertext: Bytes) -> Result<Bytes> {
-    let mut arr = [0u8; NONCE_SIZE_IN_BYTES];
+        let mut ciphertext_body = ciphertext;
-    thread_rng().try_fill(&mut arr[..])?;
+        let nonce_bytes = ciphertext_body.split_to(NONCE_SIZE);
-    let nonce = Nonce::from_slice(&arr);
+        let nonce = Nonce::from_slice(&nonce_bytes);
-    let plaintext = body.as_ref();
+        match self.cipher.decrypt(nonce, ciphertext_body.as_ref()) {
-    match cipher.encrypt(nonce, plaintext) {
+            Ok(payload) => Ok(Bytes::from(payload)),
-        Ok(body) => Ok(EncryptedPayload {
+            Err(e) => Err(anyhow!(e.to_string())),
-            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

@@ -1,35 +1,104 @@
 use anyhow::Result;
+use futures::future::try_join_all;
+
 use serde::{Deserialize, Serialize};
 use std::fs::Metadata;
 use std::path::PathBuf;
+
+use std::io::{Seek, SeekFrom};
+
 use tokio::fs::File;
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct FileInfo {
+pub struct ChunkHeader {
+    pub id: u8,
+    pub start: u64,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct FileOffer {
+    pub id: u8,
     pub path: PathBuf,
     pub size: u64,
 }
 
+pub struct StdFileHandle {
+    pub id: u8,
+    pub file: std::fs::File,
+    pub start: u64,
+}
+
+impl StdFileHandle {
+    pub async fn new(id: u8, file: File, start: u64) -> Result<StdFileHandle> {
+        let mut std_file = file.into_std().await;
+        std_file.seek(SeekFrom::Start(start))?;
+        Ok(StdFileHandle {
+            id: id,
+            file: std_file,
+            start: start,
+        })
+    }
+}
+
 pub struct FileHandle {
+    pub id: u8,
     pub file: File,
     pub md: Metadata,
     pub path: PathBuf,
 }
 
 impl FileHandle {
-    pub async fn new(path: PathBuf) -> Result<FileHandle> {
+    pub async fn new(id: u8, path: PathBuf) -> Result<FileHandle> {
         let file = File::open(&path).await?;
         let md = file.metadata().await?;
-        let fh = FileHandle { file, md, path };
+        let fh = FileHandle { id, file, md, path };
         return Ok(fh);
     }
 
-    pub fn to_file_info(&self) -> FileInfo {
+    pub async fn to_stds(
-        FileInfo {
+        file_handles: Vec<FileHandle>,
+        chunk_headers: Vec<ChunkHeader>,
+    ) -> Vec<StdFileHandle> {
+        let mut ret = Vec::new();
+        for handle in file_handles {
+            let chunk = chunk_headers.iter().find(|chunk| handle.id == chunk.id);
+            match chunk {
+                Some(chunk) => {
+                    match handle.to_std(chunk).await {
+                        Ok(std_file_handle) => {
+                            ret.push(std_file_handle);
+                        }
+                        _ => println!("Error seeking in file"),
+                    };
+                }
+                None => {
+                    println!("Skipping file b/c not in requested chunks");
+                }
+            }
+        }
+        ret
+    }
+
+    async fn to_std(self, chunk_header: &ChunkHeader) -> Result<StdFileHandle> {
+        StdFileHandle::new(self.id, self.file, chunk_header.start).await
+    }
+
+    pub fn to_file_offer(&self) -> FileOffer {
+        FileOffer {
+            id: self.id,
             path: self.path.clone(),
             size: self.md.len(),
         }
     }
+
+    pub async fn get_file_handles(file_paths: &Vec<PathBuf>) -> Result<Vec<FileHandle>> {
+        let tasks = file_paths
+            .into_iter()
+            .enumerate()
+            .map(|(idx, path)| FileHandle::new(idx.try_into().unwrap(), path.to_path_buf()));
+        let handles = try_join_all(tasks).await?;
+        Ok(handles)
+    }
 }
 
 const SUFFIX: [&'static str; 9] = ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"];

src/handshake.rs

@@ -0,0 +1,72 @@
+use crate::conf::{HANDSHAKE_MSG_SIZE, ID_SIZE};
+
+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 = [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 {
+        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, Vec<u8>)> {
+        let mut socket = socket;
+        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();
+        // 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, 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

@@ -1,10 +1,13 @@
 mod cli;
 mod client;
 mod conf;
+mod connection;
 mod crypto;
 mod file;
+mod handshake;
 mod message;
 mod server;
+mod ui;
 
 use clap::Parser;
 use cli::{Cli, Commands};

src/message.rs

@@ -1,109 +1,50 @@
-use crate::crypto::{decrypt, encrypt};
+use crate::file::{ChunkHeader, FileOffer};
-use crate::file::FileInfo;
 
-use aes_gcm::Aes256Gcm; // Or `Aes128Gcm`
 use anyhow::{anyhow, Result};
 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(HandshakePayload),
+    FileOffer(FileOfferPayload),
-    EncryptedMessage(EncryptedPayload),
+    FileRequest(FileRequestPayload),
-    ErrorMessage(RuckError),
+    FileTransfer(FileTransferPayload),
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct HandshakePayload {
+pub struct FileRequestPayload {
-    pub id: Bytes,
+    pub chunks: Vec<ChunkHeader>,
-    pub msg: Bytes,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct EncryptedPayload {
+pub struct FileOfferPayload {
-    pub nonce: Vec<u8>,
+    pub files: Vec<FileOffer>,
-    pub body: Vec<u8>,
-}
-
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub enum EncryptedMessage {
-    FileNegotiationMessage(FileNegotiationPayload),
-    FileTransferMessage(FileTransferPayload),
-}
-
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct FileNegotiationPayload {
-    pub files: Vec<FileInfo>,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct FileTransferPayload {
-    pub file_info: FileInfo,
+    pub chunk_header: ChunkHeader,
-    pub chunk_num: u64,
     pub chunk: Bytes,
 }
 
-impl EncryptedMessage {
+impl Message {
-    pub fn from_encrypted_message(cipher: &Aes256Gcm, payload: &EncryptedPayload) -> Result<Self> {
+    pub fn serialize(&self) -> Result<Bytes> {
-        let raw = decrypt(cipher, payload)?;
+        bincode::serialize(&self).map(|vec| Ok(Bytes::from(vec)))?
-        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> {
+    pub fn deserialize(bytes: Bytes) -> Result<Self> {
-        let raw = match bincode::serialize(&self) {
+        match bincode::deserialize(bytes.as_ref()) {
-            Ok(result) => result,
+            Ok(msg) => Ok(msg),
-            Err(e) => {
+            Err(e) => Err(anyhow!(e.to_string())),
-                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(
+        Framed::new(stream, LengthDelimitedCodec::new())
-            Framed::new(stream, LengthDelimitedCodec::new()),
-            tokio_serde::formats::SymmetricalBincode::<Message>::default(),
-        )
     }
 }
 
-pub type MessageStream = SymmetricallyFramed<
+pub type MessageStream = Framed<TcpStream, LengthDelimitedCodec>;
-    Framed<TcpStream, LengthDelimitedCodec>,
-    Message,
-    SymmetricalBincode<Message>,
->;

src/server.rs

@@ -1,16 +1,16 @@
-use crate::message::{HandshakePayload, Message, MessageStream, RuckError};
+use crate::conf::PER_CLIENT_BUFFER;
-
+use crate::handshake::Handshake;
 use anyhow::{anyhow, Result};
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
-use futures::prelude::*;
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::{mpsc, Mutex};
 
-type Tx = mpsc::UnboundedSender<Message>;
+type Tx = mpsc::UnboundedSender<Bytes>;
-type Rx = mpsc::UnboundedReceiver<Message>;
+type Rx = mpsc::UnboundedReceiver<Bytes>;
 
 pub struct Shared {
     handshake_cache: HashMap<Bytes, Tx>,
@@ -18,12 +18,12 @@ pub struct Shared {
 type State = Arc<Mutex<Shared>>;
 
 struct Client {
-    messages: MessageStream,
+    socket: TcpStream,
     rx: Rx,
     peer_tx: Option<Tx>,
 }
 struct StapledClient {
-    messages: MessageStream,
+    socket: TcpStream,
     rx: Rx,
     peer_tx: Tx,
 }
@@ -37,11 +37,11 @@ impl Shared {
 }
 
 impl Client {
-    async fn new(id: Bytes, state: State, messages: MessageStream) -> Result<Client> {
+    async fn new(id: Bytes, state: State, socket: TcpStream) -> Result<Client> {
         let (tx, rx) = mpsc::unbounded_channel();
         let mut shared = state.lock().await;
         let client = Client {
-            messages,
+            socket,
             rx,
             peer_tx: shared.handshake_cache.remove(&id),
         };
@@ -49,11 +49,7 @@ impl Client {
         Ok(client)
     }
 
-    async fn upgrade(
+    async fn upgrade(client: Client, state: State, handshake: Handshake) -> Result<StapledClient> {
-        client: Client,
-        state: State,
-        handshake_payload: HandshakePayload,
-    ) -> Result<StapledClient> {
         let mut client = client;
         let peer_tx = match client.peer_tx {
             // Receiver - already stapled at creation
@@ -61,28 +57,21 @@ impl Client {
             // Sender - needs to wait for the incoming msg to look up peer_tx
             None => {
                 tokio::select! {
+                    // Client reads handshake message sent over channel
                     Some(msg) = client.rx.recv() => {
-                        client.messages.send(msg).await?
+                        // Writes parnter handshake message over wire
-                    }
+                        client.socket.write_all(&msg[..]).await?
-                    result = client.messages.next() => match result {
-                        Some(_) => return Err(anyhow!("Client sending more messages before handshake complete")),
-                        None => return Err(anyhow!("Connection interrupted")),
                     }
                 }
-                match state
+                match state.lock().await.handshake_cache.remove(&handshake.id) {
-                    .lock()
-                    .await
-                    .handshake_cache
-                    .remove(&handshake_payload.id)
-                {
                     Some(peer_tx) => peer_tx,
                     None => return Err(anyhow!("Connection not stapled")),
                 }
             }
         };
-        peer_tx.send(Message::HandshakeMessage(handshake_payload))?;
+        peer_tx.send(handshake.outbound_msg)?;
         Ok(StapledClient {
-            messages: client.messages,
+            socket: client.socket,
             rx: client.rx,
             peer_tx,
         })
@@ -109,26 +98,14 @@ pub async fn serve() -> Result<()> {
 pub async fn handle_connection(
     state: Arc<Mutex<Shared>>,
     socket: TcpStream,
-    addr: SocketAddr,
+    _addr: SocketAddr,
 ) -> Result<()> {
-    let mut stream = Message::to_stream(socket);
+    socket.readable().await?;
-    println!("server - new conn from {:?}", addr);
+    let (handshake, socket) = Handshake::from_socket(socket).await?;
-    let handshake_payload = match stream.next().await {
+    let id = handshake.id.clone();
-        Some(Ok(Message::HandshakeMessage(payload))) => payload,
+    let client = Client::new(id.clone(), state.clone(), socket).await?;
-        Some(Ok(_)) => {
+    println!("Client created");
-            stream
+    let mut client = match Client::upgrade(client, state.clone(), handshake).await {
-                .send(Message::ErrorMessage(RuckError::NotHandshake))
-                .await?;
-            return Ok(());
-        }
-        _ => {
-            println!("No first message");
-            return Ok(());
-        }
-    };
-    let id = handshake_payload.id.clone();
-    let client = Client::new(id.clone(), state.clone(), stream).await?;
-    let mut client = match Client::upgrade(client, state.clone(), handshake_payload).await {
         Ok(client) => client,
         Err(err) => {
             // Clear handshake cache if staple is unsuccessful
@@ -136,22 +113,30 @@ pub async fn handle_connection(
             return Err(err);
         }
     };
+    println!("Client upgraded");
     // The handshake cache should be empty for {id} at this point.
+    let mut client_buffer = BytesMut::with_capacity(PER_CLIENT_BUFFER);
     loop {
         tokio::select! {
             Some(msg) = client.rx.recv() => {
-                client.messages.send(msg).await?
+                // println!("piping bytes= {:?}", msg);
+                client.socket.write_all(&msg[..]).await?
             }
-            result = client.messages.next() => match result {
+            result = client.socket.read_buf(&mut client_buffer) => match result {
-                Some(Ok(msg)) => {
+                Ok(0) => {
-                    client.peer_tx.send(msg)?
+                    break;
-                }
+                },
-                Some(Err(e)) => {
+                Ok(n) => {
+                    let b = BytesMut::from(&client_buffer[0..n]).freeze();
+                    client.peer_tx.send(b)?;
+                    client_buffer.clear();
+                },
+                Err(e) => {
                     println!("Error {:?}", e);
                 }
-                None => break,
             }
         }
     }
+    println!("done with client");
     Ok(())
 }

src/ui.rs

@@ -0,0 +1,46 @@
+use crate::file::{to_size_string, FileOffer};
+
+use futures::prelude::*;
+
+use tokio::io::{self};
+
+use tokio_util::codec::{FramedRead, LinesCodec};
+
+pub async fn prompt_user_for_file_confirmation(file_offers: Vec<FileOffer>) -> Vec<FileOffer> {
+    let mut stdin = FramedRead::new(io::stdin(), LinesCodec::new());
+    let mut files = vec![];
+    for file_offer in file_offers.into_iter() {
+        let mut reply = prompt_user_input(&mut stdin, &file_offer).await;
+        while reply.is_none() {
+            reply = prompt_user_input(&mut stdin, &file_offer).await;
+        }
+        match reply {
+            Some(true) => files.push(file_offer),
+            _ => {}
+        }
+    }
+    files
+}
+
+pub async fn prompt_user_input(
+    stdin: &mut FramedRead<io::Stdin, LinesCodec>,
+    file_offer: &FileOffer,
+) -> Option<bool> {
+    let prompt_name = file_offer.path.file_name().unwrap();
+    println!(
+        "Accept {:?}? ({:?}). (Y/n)",
+        prompt_name,
+        to_size_string(file_offer.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,
+    }
+}