Merge pull request #5 from rictorlome/resume

Resume

README.md

@@ -2,11 +2,22 @@
 
 `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.
 
-### Version
+## Usage
 
-This document refers to version `0.1.0` of `ruck` as defined by the `Cargo.toml` file.
+```bash
+## tab 1
+cargo run relay # this starts the server
 
-## Server
+## tab 2
+cargo run send /path/to/file1.md /path/to/file2.md supersecretpassword
+
+## tab 3
+cargo run receive supersecretpassword
+```
+
+## Protocol
+
+### Server
 
 The server in `ruck` exposes a TCP port.
 Its only functions are to staple connections and shuttle bytes between stapled connections.
@@ -19,7 +30,7 @@ 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.
 
-## Client
+### Client
 
 There are two types of clients - `send` and `receive` clients.
 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.
@@ -29,7 +40,5 @@ Once the handshake is complete, `send` and `receive` negotiate and exchange file
 
 - `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.
+- `send` sends the specified bytes, then a completion message and hangs up.
 - `receive` hangs up once the downloads are complete.

src/client.rs

@@ -6,7 +6,6 @@ use crate::ui::prompt_user_for_file_confirmation;
 
 use anyhow::{anyhow, Result};
 
-use std::ffi::OsStr;
 use std::path::PathBuf;
 use tokio::fs::File;
 
@@ -44,9 +43,7 @@ pub async fn receive(password: &String) -> Result<()> {
     let (socket, key) = handshake.negotiate(socket, s1).await?;
     let mut connection = Connection::new(socket, key);
     // Wait for offered files, respond with desired files
-    let desired_files = request_specific_files(&mut connection).await?;
-    // Create files
-    let std_file_handles = create_files(desired_files).await?;
+    let std_file_handles = request_specific_files(&mut connection).await?;
     // Download them
     connection.download_files(std_file_handles).await?;
     return Ok(());
@@ -57,7 +54,10 @@ pub async fn offer_files(
     file_handles: &Vec<FileHandle>,
 ) -> Result<Vec<ChunkHeader>> {
     // Collect file offer
-    let files = file_handles.iter().map(|fh| fh.to_file_offer()).collect();
+    let mut files = vec![];
+    for handle in file_handles {
+        files.push(handle.to_file_offer()?);
+    }
     let msg = Message::FileOffer(FileOfferPayload { files });
     // Send file offer
     conn.send_msg(msg).await?;
@@ -70,7 +70,7 @@ pub async fn offer_files(
     }
 }
 
-pub async fn request_specific_files(conn: &mut Connection) -> Result<Vec<FileOffer>> {
+pub async fn request_specific_files(conn: &mut Connection) -> Result<Vec<StdFileHandle>> {
     // Wait for offer message
     let offer_message = conn.await_msg().await?;
     let offered_files: Vec<FileOffer> = match offer_message {
@@ -79,28 +79,50 @@ pub async fn request_specific_files(conn: &mut Connection) -> Result<Vec<FileOff
     };
     // Prompt user for confirmation of files
     let desired_files = prompt_user_for_file_confirmation(offered_files).await;
+    let std_file_handles = create_or_find_files(desired_files).await?;
     let file_request_msg = Message::FileRequest(FileRequestPayload {
-        chunks: desired_files
+        chunks: std_file_handles
             .iter()
             .map(|file| ChunkHeader {
                 id: file.id,
-                start: 0,
+                start: file.start,
             })
             .collect(),
     });
     conn.send_msg(file_request_msg).await?;
-    Ok(desired_files)
+    Ok(std_file_handles)
 }
 
-pub async fn create_files(desired_files: Vec<FileOffer>) -> Result<Vec<StdFileHandle>> {
+pub async fn create_or_find_files(desired_files: Vec<FileOffer>) -> Result<Vec<StdFileHandle>> {
     let mut v = Vec::new();
     for desired_file in desired_files {
-        let filename = desired_file
-            .path
-            .file_name()
-            .unwrap_or(OsStr::new("random.txt"));
-        let file = File::create(filename).await?;
-        let std_file_handle = StdFileHandle::new(desired_file.id, file, 0).await?;
+        let filename = desired_file.path;
+        // filename.push_str(".part");
+        let file = match File::open(filename.clone()).await {
+            Ok(file) => {
+                println!(
+                    "File {:?} already exists. Attempting to resume download.",
+                    filename
+                );
+                file
+            }
+            Err(_) => File::create(&filename).await?,
+        };
+        let metadata = file.metadata().await?;
+        println!(
+            "File: {:?}. Current len: {:?}, Full Size: {:?}",
+            filename.clone(),
+            metadata.len(),
+            desired_file.size
+        );
+        let std_file_handle = StdFileHandle::new(
+            desired_file.id,
+            filename,
+            file,
+            metadata.len(),
+            desired_file.size,
+        )
+        .await?;
         v.push(std_file_handle)
     }
     return Ok(v);

src/connection.rs

@@ -77,15 +77,17 @@ impl Connection {
                 Err(e) => return Err(anyhow!(e.to_string())),
             }
         }
+        self.send_msg(Message::FileTransferComplete).await?;
         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 (compressed), {:?} iterations. {:?} total time, {:?} avg per iteration, {:?} avg mb/sec",
+            handle.name,
             mb_sent,
             count,
             elapsed,
             elapsed / count,
-            mb_sent / elapsed.as_secs()
+            1000 * mb_sent as u128 / elapsed.as_millis()
         );
         Ok(())
     }
@@ -105,6 +107,7 @@ impl Connection {
     }
 
     pub async fn download_file(&mut self, handle: StdFileHandle) -> Result<()> {
+        let clone = handle.file.try_clone()?;
         let mut decoder = GzDecoder::new(handle.file);
         loop {
             let msg = self.await_msg().await?;
@@ -113,16 +116,32 @@ impl Connection {
                     if payload.chunk_header.id != handle.id {
                         return Err(anyhow!("Wrong file"));
                     }
-                    if payload.chunk.len() == 0 {
-                        break;
-                    }
                     decoder.write_all(&payload.chunk[..])?
                 }
+                Message::FileTransferComplete => {
+                    break;
+                }
                 _ => return Err(anyhow!("Expecting file transfer message")),
             }
         }
         decoder.finish()?;
-        println!("Done downloading file.");
+        println!("Done downloading {:?}.", handle.name);
+        Connection::check_and_finish_download(clone, handle.name, handle.size).await?;
         Ok(())
     }
+
+    pub async fn check_and_finish_download(
+        file: std::fs::File,
+        filename: String,
+        size: u64,
+    ) -> Result<()> {
+        let metadata = file.metadata()?;
+        if metadata.len() == size {
+            println!("OK: downloaded {:?} matches advertised size.", filename);
+            return Ok(());
+        }
+        return Err(anyhow!(
+            "Downloaded file does not match expected size. Try again"
+        ));
+    }
 }

src/file.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use futures::future::try_join_all;
 
 use serde::{Deserialize, Serialize};
@@ -18,24 +18,34 @@ pub struct ChunkHeader {
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct FileOffer {
     pub id: u8,
-    pub path: PathBuf,
+    pub path: String,
     pub size: u64,
 }
 
 pub struct StdFileHandle {
     pub id: u8,
+    pub name: String,
     pub file: std::fs::File,
     pub start: u64,
+    pub size: 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))?;
+    pub async fn new(
+        id: u8,
+        name: String,
+        file: File,
+        start: u64,
+        size: u64,
+    ) -> Result<StdFileHandle> {
+        let mut file = file.into_std().await;
+        file.seek(SeekFrom::Start(start))?;
         Ok(StdFileHandle {
-            id: id,
-            file: std_file,
-            start: start,
+            id,
+            name,
+            file,
+            start,
+            size,
         })
     }
 }
@@ -80,15 +90,22 @@ impl FileHandle {
     }
 
     async fn to_std(self, chunk_header: &ChunkHeader) -> Result<StdFileHandle> {
-        StdFileHandle::new(self.id, self.file, chunk_header.start).await
+        StdFileHandle::new(
+            self.id,
+            pathbuf_to_string(&self.path)?,
+            self.file,
+            chunk_header.start,
+            self.md.len(),
+        )
+        .await
     }
 
-    pub fn to_file_offer(&self) -> FileOffer {
-        FileOffer {
+    pub fn to_file_offer(&self) -> Result<FileOffer> {
+        Ok(FileOffer {
             id: self.id,
-            path: self.path.clone(),
+            path: pathbuf_to_string(&self.path)?,
             size: self.md.len(),
-        }
+        })
     }
 
     pub async fn get_file_handles(file_paths: &Vec<PathBuf>) -> Result<Vec<FileHandle>> {
@@ -115,3 +132,15 @@ pub fn to_size_string(size: u64) -> String {
 
     result
 }
+
+pub fn pathbuf_to_string(path: &PathBuf) -> Result<String> {
+    let filename = match path.file_name() {
+        Some(s) => s,
+        None => return Err(anyhow!("Could not get filename from file offer.")),
+    };
+    let filename = filename.to_os_string();
+    match filename.into_string() {
+        Ok(s) => Ok(s),
+        Err(_) => Err(anyhow!("Error converting {:?} to String", path)),
+    }
+}

src/handshake.rs

@@ -49,7 +49,7 @@ impl Handshake {
     ) -> Result<(TcpStream, Vec<u8>)> {
         let mut socket = socket;
         let bytes = &self.to_bytes();
-        println!("client - sending handshake msg= {:?}", &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?;

src/message.rs

@@ -11,6 +11,7 @@ pub enum Message {
     FileOffer(FileOfferPayload),
     FileRequest(FileRequestPayload),
     FileTransfer(FileTransferPayload),
+    FileTransferComplete,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

src/ui.rs

@@ -26,7 +26,7 @@ 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();
+    let prompt_name = &file_offer.path;
     println!(
         "Accept {:?}? ({:?}). (Y/n)",
         prompt_name,