diff --git a/src/client.rs b/src/client.rs index bc7387d..48b3162 100644 --- a/src/client.rs +++ b/src/client.rs @@ -30,6 +30,7 @@ pub struct Client { mode: Mode, filename: PathBuf, save_path: PathBuf, + clean_on_error: bool } /// Enum used to set the client either in Download Mode or Upload Mode @@ -52,6 +53,7 @@ impl Client { mode: config.mode, filename: config.filename.clone(), save_path: config.receive_directory.clone(), + clean_on_error: config.clean_on_error }) } @@ -240,6 +242,7 @@ impl Client { Worker::new( socket, file, + self.clean_on_error, self.blocksize, DEFAULT_TIMEOUT, self.windowsize, @@ -249,6 +252,7 @@ impl Client { Worker::new( socket, PathBuf::from(self.filename.clone()), + self.clean_on_error, self.blocksize, DEFAULT_TIMEOUT, self.windowsize, diff --git a/src/client_config.rs b/src/client_config.rs index 9440c03..952fcfe 100644 --- a/src/client_config.rs +++ b/src/client_config.rs @@ -39,6 +39,8 @@ pub struct ClientConfig { pub receive_directory: PathBuf, /// File to Upload or Download. pub filename: PathBuf, + /// Should clean (delete) files after receiving errors. (default: true) + pub clean_on_error: bool } impl Default for ClientConfig { @@ -52,6 +54,7 @@ impl Default for ClientConfig { mode: Mode::Download, receive_directory: Default::default(), filename: Default::default(), + clean_on_error: true } } } @@ -116,6 +119,9 @@ impl ClientConfig { "-d" | "--download" => { config.mode = Mode::Download; } + "--dont-clean" => { + config.clean_on_error = false; + } "-h" | "--help" => { println!("TFTP Client\n"); println!("Usage: tftpd client [OPTIONS]\n"); @@ -130,6 +136,7 @@ impl ClientConfig { println!(" -u, --upload\t\t\t\tSets the client to upload mode, Ignores all previous download flags"); println!(" -d, --download\t\t\tSet the client to download mode, Invalidates all previous upload flags"); println!(" -rd, --receive-directory \tSet the directory to receive files when in Download mode (default: current working directory)"); + println!(" --dont-clean\t\t\t\tWill prevent client from deleting files after receiving errors."); println!(" -h, --help\t\t\t\tPrint help information"); process::exit(0); } @@ -169,6 +176,7 @@ mod tests { "2", "-t", "4", + "--dont-clean" ] .iter() .map(|s| s.to_string()), @@ -183,6 +191,7 @@ mod tests { assert_eq!(config.blocksize, 1024); assert_eq!(config.mode, Mode::Upload); assert_eq!(config.timeout, Duration::from_secs(4)); + assert_eq!(config.clean_on_error, false); } #[test] diff --git a/src/config.rs b/src/config.rs index 39b4459..d6201c8 100644 --- a/src/config.rs +++ b/src/config.rs @@ -37,6 +37,8 @@ pub struct Config { pub duplicate_packets: u8, /// Overwrite existing files. (default: false) pub overwrite: bool, + /// Should clean (delete) files after receiving errors. (default: true) + pub clean_on_error: bool } impl Default for Config { @@ -51,6 +53,7 @@ impl Default for Config { read_only: Default::default(), duplicate_packets: Default::default(), overwrite: Default::default(), + clean_on_error: true } } } @@ -131,6 +134,7 @@ impl Config { println!(" -r, --read-only\t\t\tRefuse all write requests, making the server read-only (default: false)"); println!(" --duplicate-packets \t\tDuplicate all packets sent from the server (default: 0)"); println!(" --overwrite\t\t\t\tOverwrite existing files (default: false)"); + println!(" --dont-clean\t\t\t\tWill prevent daemon from deleting files after receiving errors."); println!(" -h, --help\t\t\t\tPrint help information"); process::exit(0); } @@ -153,6 +157,9 @@ impl Config { "--overwrite" => { config.overwrite = true; } + "--dont-clean" => { + config.clean_on_error = false; + } invalid => return Err(format!("Invalid flag: {invalid}").into()), } @@ -179,7 +186,7 @@ mod tests { fn parses_full_config() { let config = Config::new( [ - "/", "-i", "0.0.0.0", "-p", "1234", "-d", "/", "-rd", "/", "-sd", "/", "-s", "-r", + "/", "-i", "0.0.0.0", "-p", "1234", "-d", "/", "-rd", "/", "-sd", "/", "-s", "-r", "--dont-clean" ] .iter() .map(|s| s.to_string()), @@ -191,6 +198,7 @@ mod tests { assert_eq!(config.directory, PathBuf::from("/")); assert_eq!(config.receive_directory, PathBuf::from("/")); assert_eq!(config.send_directory, PathBuf::from("/")); + assert_eq!(config.clean_on_error, false); assert!(config.single_port); assert!(config.read_only); } diff --git a/src/server.rs b/src/server.rs index 922117d..9cd4d2e 100644 --- a/src/server.rs +++ b/src/server.rs @@ -34,6 +34,7 @@ pub struct Server { single_port: bool, read_only: bool, overwrite: bool, + clean_on_error: bool, duplicate_packets: u8, largest_block_size: usize, clients: HashMap>, @@ -50,6 +51,7 @@ impl Server { single_port: config.single_port, read_only: config.read_only, overwrite: config.overwrite, + clean_on_error: config.clean_on_error, duplicate_packets: config.duplicate_packets, largest_block_size: DEFAULT_BLOCK_SIZE, clients: HashMap::new(), @@ -179,6 +181,7 @@ impl Server { let worker = Worker::new( socket, file_path.clone(), + self.clean_on_error, worker_options.block_size, worker_options.timeout, worker_options.window_size, @@ -220,6 +223,7 @@ impl Server { let worker = Worker::new( socket, file_path.clone(), + self.clean_on_error, worker_options.block_size, worker_options.timeout, worker_options.window_size, diff --git a/src/worker.rs b/src/worker.rs index 7ec0161..84b2cd4 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -32,6 +32,7 @@ const DEFAULT_DUPLICATE_DELAY: Duration = Duration::from_millis(1); /// let worker = Worker::new( /// Box::new(socket), /// PathBuf::from_str("Cargo.toml").unwrap(), +/// true, /// 512, /// Duration::from_secs(1), /// 1, @@ -43,6 +44,7 @@ const DEFAULT_DUPLICATE_DELAY: Duration = Duration::from_millis(1); pub struct Worker { socket: Box, file_name: PathBuf, + clean_on_error: bool, blk_size: usize, timeout: Duration, windowsize: u16, @@ -54,6 +56,7 @@ impl Worker { pub fn new( socket: Box, file_name: PathBuf, + clean_on_error: bool, blk_size: usize, timeout: Duration, windowsize: u16, @@ -62,6 +65,7 @@ impl Worker { Worker { socket, file_name, + clean_on_error, blk_size, timeout, windowsize, @@ -102,6 +106,7 @@ impl Worker { /// Receives a file from the remote [`SocketAddr`] that has sent a write request using /// the supplied socket, asynchronously. pub fn receive(self) -> Result, Box> { + let clean_on_error = self.clean_on_error; let file_name = self.file_name.clone(); let remote_addr = self.socket.remote_addr().unwrap(); @@ -122,8 +127,10 @@ impl Worker { } Err(err) => { eprintln!("{err}"); - if fs::remove_file(&file_name).is_err() { - eprintln!("Error while cleaning {}", &file_name.to_str().unwrap()); + if clean_on_error { + if fs::remove_file(&file_name).is_err() { + eprintln!("Error while cleaning {}", &file_name.to_str().unwrap()); + } } } }