Merge pull request #74 from genonullfree/improve/client-hash-once

Reuse hash function and client hash once
Fix `../` removal

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "teleporter"
-version = "0.8.1"
+version = "0.8.2"
 authors = ["geno nullfree <[email protected]>"]
 license = "BSD-3-Clause"
 description = "A small utility to send files quickly from point A to point B"

src/client.rs

@@ -3,9 +3,7 @@ use crate::teleport::{TeleportAction, TeleportFeatures, TeleportStatus};
 use crate::teleport::{TeleportInit, TeleportInitAck};
 use crate::utils::print_updates;
 use crate::*;
-use std::hash::Hasher;
 use std::path::Path;
-use xxhash_rust::xxh3;
 
 #[derive(Debug)]
 struct Replace {
@@ -135,10 +133,10 @@ pub fn run(mut opt: Opt) -> Result<(), Error> {
             }
         };
 
-        let thread_file = filepath.clone().to_string();
+        let thread_file = File::open(&filepath)?;
         let handle = match opt.overwrite && !opt.no_delta {
             true => Some(thread::spawn(move || {
-                utils::calc_file_hash(thread_file).unwrap()
+                utils::calc_delta_hash(&thread_file).unwrap()
             })),
             false => None,
         };
@@ -235,13 +233,13 @@ pub fn run(mut opt: Opt) -> Result<(), Error> {
             }
             TeleportStatus::WrongVersion => {
                 println!(
-                    "Error: Version mismatch! Server: {:?} Us: {}",
+                    "Version mismatch! Server: {:?} Us: {}",
                     recv.version, VERSION
                 );
                 break;
             }
             TeleportStatus::RequiresEncryption => {
-                println!("Error: The server requires encryption");
+                println!("The server requires encryption");
                 break;
             }
             TeleportStatus::EncryptionError => {
@@ -252,17 +250,17 @@ pub fn run(mut opt: Opt) -> Result<(), Error> {
         };
 
         let csum_recv = recv.delta.as_ref().map(|r| r.hash);
-        let mut hash: Option<u64> = None;
+        let mut file_delta: Option<TeleportDelta> = None;
         if utils::check_feature(&recv.features, TeleportFeatures::Overwrite) {
-            hash = handle.map(|s| s.join().expect("calc_file_hash panicked"));
+            file_delta = handle.map(|s| s.join().expect("calc_file_hash panicked"));
         }
 
-        if hash != None && hash == csum_recv {
+        if file_delta.is_some() && file_delta.as_ref().unwrap().hash == csum_recv.unwrap() {
             // File matches hash
             send_data_complete(stream, &enc, file)?;
         } else {
             // Send file data
-            send(stream, file, &header, &enc, recv.delta)?;
+            send(stream, file, &header, &enc, recv.delta, file_delta)?;
         }
 
         let duration = start_time.elapsed();
@@ -298,20 +296,42 @@ fn send(
     header: &TeleportInit,
     enc: &Option<TeleportEnc>,
     delta: Option<TeleportDelta>,
+    file_delta: Option<TeleportDelta>,
 ) -> Result<(), Error> {
     let mut buf = Vec::<u8>::new();
-    let mut hash_list = Vec::<u64>::new();
     match delta {
-        Some(d) => {
-            buf.resize(d.chunk_size as usize, 0);
-            hash_list = d.chunk_hash;
-        }
+        Some(ref d) => buf.resize(d.chunk_size as usize, 0),
         None => buf.resize(4096, 0),
     }
 
+    let compare_delta = delta.is_some() && file_delta.is_some();
+    let delta_len = if delta.is_some() {
+        delta.as_ref().unwrap().chunk_hash.len()
+    } else {
+        0
+    };
+    let file_delta_len = if file_delta.is_some() {
+        file_delta.as_ref().unwrap().chunk_hash.len()
+    } else {
+        0
+    };
+
     // Send file data
     let mut sent = 0;
     loop {
+        // Check if hash matches, if so: skip chunk
+        let index = sent / buf.len();
+        if compare_delta
+            && index < delta_len
+            && index < file_delta_len
+            && delta.as_ref().unwrap().chunk_hash[index]
+                == file_delta.as_ref().unwrap().chunk_hash[index]
+        {
+            sent += buf.len();
+            continue;
+        }
+
+        file.seek(SeekFrom::Start(sent as u64))?;
         // Read a chunk of the file
         let len = match file.read(&mut buf) {
             Ok(l) => l,
@@ -323,17 +343,6 @@ fn send(
             break;
         }
 
-        // Check if hash matches, if so: skip chunk
-        let index = sent / buf.len();
-        if !hash_list.is_empty() && index < hash_list.len() {
-            let mut hasher = xxh3::Xxh3::new();
-            hasher.write(&buf);
-            if hash_list[index] == hasher.finish() {
-                sent += len;
-                continue;
-            }
-        }
-
         let data = &buf[..len];
         let mut chunk = TeleportData {
             offset: sent as u64,

src/server.rs

@@ -44,7 +44,7 @@ pub fn run(opt: Opt) -> Result<(), Error> {
         // Receive connections in recv function
         let recv_list_clone = Arc::clone(&recv_list);
         thread::spawn(move || {
-            recv(s, recv_list_clone, args).unwrap();
+            handle_connection(s, recv_list_clone, args).unwrap();
         });
     }
 
@@ -69,15 +69,18 @@ fn print_list(list: &MutexGuard<Vec<String>>) {
     io::stdout().flush().unwrap();
 }
 
-fn rm_list(filename: &str, list: &Arc<Mutex<Vec<String>>>) {
+fn rm_filename_from_list(filename: &str, list: &Arc<Mutex<Vec<String>>>) {
     let mut recv_data = list.lock().unwrap();
     recv_data.retain(|x| x != filename);
     print_list(&recv_data);
     drop(recv_data);
 }
 
-/// Recv receives filenames and file data for a file
-fn recv(mut stream: TcpStream, recv_list: Arc<Mutex<Vec<String>>>, opt: Opt) -> Result<(), Error> {
+fn handle_connection(
+    mut stream: TcpStream,
+    recv_list: Arc<Mutex<Vec<String>>>,
+    opt: Opt,
+) -> Result<(), Error> {
     let start_time = Instant::now();
     let ip = stream.peer_addr().unwrap();
     let mut file: File;
@@ -123,9 +126,11 @@ fn recv(mut stream: TcpStream, recv_list: Arc<Mutex<Vec<String>>>, opt: Opt) ->
         return send_ack(resp, &mut stream, &enc);
     }
 
-    if !opt.allow_dangerous_filepath && filename.starts_with('/') {
-        // Remove any preceeding '/'
-        filename.remove(0);
+    if !opt.allow_dangerous_filepath {
+        if filename.starts_with('/') {
+            // Remove any preceeding '/'
+            filename.remove(0);
+        }
 
         // Prohibit directory traversal
         filename = filename.replace("../", "");
@@ -214,7 +219,7 @@ fn recv(mut stream: TcpStream, recv_list: Arc<Mutex<Vec<String>>>, opt: Opt) ->
                 "Connection closed (reason: {:?}). Aborted {} transfer.",
                 e, &filename
             );
-            rm_list(&filename, &recv_list);
+            rm_filename_from_list(&filename, &recv_list);
             return Ok(());
         }
     }
@@ -279,7 +284,7 @@ fn recv(mut stream: TcpStream, recv_list: Arc<Mutex<Vec<String>>>, opt: Opt) ->
         }
     }
 
-    rm_list(&filename, &recv_list);
+    rm_filename_from_list(&filename, &recv_list);
 
     Ok(())
 }

src/utils.rs

@@ -155,36 +155,6 @@ fn gen_chunk_size(file_size: u64) -> usize {
     chunk as usize
 }
 
-// Called from client
-pub fn calc_file_hash(filename: String) -> Result<u64, Error> {
-    let mut hasher = xxh3::Xxh3::new();
-    let mut buf = Vec::<u8>::new();
-
-    let mut file = File::open(filename)?;
-    let meta = file.metadata()?;
-
-    buf.resize(gen_chunk_size(meta.len()), 0);
-
-    file.seek(SeekFrom::Start(0))?;
-
-    loop {
-        // Read a chunk of the file
-        let len = match file.read(&mut buf) {
-            Ok(l) => l,
-            Err(s) => return Err(s),
-        };
-        if len == 0 {
-            break;
-        }
-
-        hasher.write(&buf);
-    }
-
-    file.seek(SeekFrom::Start(0))?;
-
-    Ok(hasher.finish())
-}
-
 pub fn add_feature(opt: &mut Option<u32>, add: TeleportFeatures) -> Result<(), Error> {
     if let Some(o) = opt {
         *o |= add as u32;